Compare commits

...

1 Commits

Author SHA1 Message Date
Michelle Ark
a912b25fd7 first pass: --sample 2023-08-07 14:21:37 -04:00
4 changed files with 25 additions and 4 deletions

View File

@@ -1,6 +1,7 @@
from collections.abc import Hashable
from dataclasses import dataclass, field
from typing import Optional, TypeVar, Any, Type, Dict, Iterator, Tuple, Set
import uuid
from dbt.contracts.graph.nodes import SourceDefinition, ManifestNode, ResultNode, ParsedNode
from dbt.contracts.relation import (
@@ -35,6 +36,7 @@ class BaseRelation(FakeAPIObject, Hashable):
include_policy: Policy = field(default_factory=lambda: Policy())
quote_policy: Policy = field(default_factory=lambda: Policy())
dbt_created: bool = False
sample: Optional[int] = None
def _is_exactish_match(self, field: ComponentName, value: str) -> bool:
if self.dbt_created and self.quote_policy.get_part(field) is False:
@@ -180,7 +182,11 @@ class BaseRelation(FakeAPIObject, Hashable):
def render(self) -> str:
# if there is nothing set, this will return the empty string.
return ".".join(part for _, part in self._render_iterator() if part is not None)
rendered_parts = ".".join(part for _, part in self._render_iterator() if part is not None)
if self.sample and rendered_parts:
alias = f"_dbt_sample_{uuid.uuid4().hex.upper()[:6]}"
return f"(select * from {rendered_parts} limit {self.sample}) {alias}"
return rendered_parts
def quoted(self, identifier):
return "{quote_char}{identifier}{quote_char}".format(

View File

@@ -323,6 +323,7 @@ def docs_serve(ctx, **kwargs):
@p.profile
@p.profiles_dir
@p.project_dir
@p.sample
@p.select
@p.selector
@p.inline
@@ -551,6 +552,7 @@ def parse(ctx, **kwargs):
@p.profile
@p.profiles_dir
@p.project_dir
@p.sample
@p.select
@p.selector
@p.state

View File

@@ -369,6 +369,15 @@ resource_type = click.option(
default=(),
)
sample = click.option(
"--sample",
envvar="DBT_SAMPLE",
help="Limit by sample rows when resolving dbt ref and sources.",
type=click.INT,
default=None,
)
model_decls = ("-m", "--models", "--model")
select_decls = ("-s", "--select")
select_attrs = {

View File

@@ -531,9 +531,13 @@ class RuntimeRefResolver(BaseRefResolver):
def create_relation(self, target_model: ManifestNode) -> RelationProxy:
if target_model.is_ephemeral_model:
self.model.set_cte(target_model.unique_id, None)
return self.Relation.create_ephemeral_from_node(self.config, target_model)
return self.Relation.create_ephemeral_from_node(
self.config, target_model, sample=self.config.args.sample
)
else:
return self.Relation.create_from(self.config, target_model)
return self.Relation.create_from(
self.config, target_model, sample=self.config.args.sample
)
def validate(
self,
@@ -590,7 +594,7 @@ class RuntimeSourceResolver(BaseSourceResolver):
target_kind="source",
disabled=(isinstance(target_source, Disabled)),
)
return self.Relation.create_from_source(target_source)
return self.Relation.create_from_source(target_source, sample=self.config.args.sample)
# metric` implementations