dbt clone

.changes/unreleased/Features-20230616-104849.yaml

@@ -0,0 +1,6 @@
+kind: Features
+body: dbt clone
+time: 2023-06-16T10:48:49.079961-05:00
+custom:
+  Author: jtcohen6 aranke
+  Issue: "7258"

core/dbt/cli/flags.py

@@ -373,6 +373,7 @@ def command_args(command: CliCommand) -> ArgsList:
     CMD_DICT: Dict[CliCommand, ClickCommand] = {
         CliCommand.BUILD: cli.build,
         CliCommand.CLEAN: cli.clean,
+        CliCommand.CLONE: cli.clone,
         CliCommand.COMPILE: cli.compile,
         CliCommand.DOCS_GENERATE: cli.docs_generate,
         CliCommand.DOCS_SERVE: cli.docs_serve,

core/dbt/cli/main.py

@@ -23,6 +23,7 @@
 from dbt.events.base_types import EventMsg
 from dbt.task.build import BuildTask
 from dbt.task.clean import CleanTask
+from dbt.task.clone import CloneTask
 from dbt.task.compile import CompileTask
 from dbt.task.debug import DebugTask
 from dbt.task.deps import DepsTask
@@ -608,6 +609,43 @@ def retry(ctx, **kwargs):
     return results, success
 
 
+# dbt clone
+@cli.command("clone")
+@click.pass_context
+@p.defer_state
+@p.exclude
+@p.full_refresh
+@p.profile
+@p.profiles_dir
+@p.project_dir
+@p.resource_type
+@p.select
+@p.selector
+@p.state  # required
+@p.target
+@p.target_path
+@p.threads
+@p.vars
+@p.version_check
+@requires.preflight
+@requires.profile
+@requires.project
+@requires.runtime_config
+@requires.manifest
+@requires.postflight
+def clone(ctx, **kwargs):
+    """Create clones of selected nodes based on their location in the manifest provided to --state."""
+    task = CloneTask(
+        ctx.obj["flags"],
+        ctx.obj["runtime_config"],
+        ctx.obj["manifest"],
+    )
+
+    results = task.run()
+    success = task.interpret_results(results)
+    return results, success
+
+
 # dbt run operation
 @cli.command("run-operation")
 @click.pass_context

core/dbt/cli/types.py

@@ -8,6 +8,7 @@ class Command(Enum):
     BUILD = "build"
     CLEAN = "clean"
     COMPILE = "compile"
+    CLONE = "clone"
     DOCS_GENERATE = "generate"
     DOCS_SERVE = "serve"
     DEBUG = "debug"

core/dbt/context/providers.py

@@ -1432,6 +1432,20 @@ def this(self) -> Optional[RelationProxy]:
             return None
         return self.db_wrapper.Relation.create_from(self.config, self.model)
 
+    @contextproperty
+    def state_relation(self) -> Optional[RelationProxy]:
+        """
+        For commands which add information about this node's corresponding
+        production version (via a --state artifact), access the Relation
+        object for that stateful other
+        """
+        if getattr(self.model, "state_relation", None):
+            return self.db_wrapper.Relation.create_from_node(
+                self.config, self.model.state_relation  # type: ignore
+            )
+        else:
+            return None
+
 
 # This is called by '_context_for', used in 'render_with_context'
 def generate_parser_model_context(

core/dbt/include/global_project/macros/materializations/models/clone/can_clone_table.sql

@@ -0,0 +1,7 @@
+{% macro can_clone_table() %}
+    {{ return(adapter.dispatch('can_clone_table', 'dbt')()) }}
+{% endmacro %}
+
+{% macro default__can_clone_table() %}
+    {{ return(False) }}
+{% endmacro %}

core/dbt/include/global_project/macros/materializations/models/clone/clone.sql

@@ -0,0 +1,68 @@
+{%- materialization clone, default -%}
+
+  {%- set relations = {'relations': []} -%}
+
+  {%- if not state_relation -%}
+      -- nothing to do
+      {{ log("No relation found in state manifest for " ~ model.unique_id, info=True) }}
+      {{ return(relations) }}
+  {%- endif -%}
+
+  {%- set existing_relation = load_cached_relation(this) -%}
+
+  {%- if existing_relation and not flags.FULL_REFRESH -%}
+      -- noop!
+      {{ log("Relation " ~ existing_relation ~ " already exists", info=True) }}
+      {{ return(relations) }}
+  {%- endif -%}
+
+  {%- set other_existing_relation = load_cached_relation(state_relation) -%}
+
+  -- If this is a database that can do zero-copy cloning of tables, and the other relation is a table, then this will be a table
+  -- Otherwise, this will be a view
+
+  {% set can_clone_table = can_clone_table() %}
+
+  {%- if other_existing_relation and other_existing_relation.type == 'table' and can_clone_table -%}
+
+      {%- set target_relation = this.incorporate(type='table') -%}
+      {% if existing_relation is not none and not existing_relation.is_table %}
+        {{ log("Dropping relation " ~ existing_relation ~ " because it is of type " ~ existing_relation.type) }}
+        {{ drop_relation_if_exists(existing_relation) }}
+      {% endif %}
+
+      -- as a general rule, data platforms that can clone tables can also do atomic 'create or replace'
+      {% call statement('main') %}
+          {{ create_or_replace_clone(target_relation, state_relation) }}
+      {% endcall %}
+
+      {% set should_revoke = should_revoke(existing_relation, full_refresh_mode=True) %}
+      {% do apply_grants(target_relation, grant_config, should_revoke=should_revoke) %}
+      {% do persist_docs(target_relation, model) %}
+
+      {{ return({'relations': [target_relation]}) }}
+
+  {%- else -%}
+
+      {%- set target_relation = this.incorporate(type='view') -%}
+
+      -- TODO: this should probably be illegal
+      -- I'm just doing it out of convenience to reuse the 'view' materialization logic
+      {%- do context.update({
+          'sql': get_clone_target(state_relation),
+          'compiled_code': get_clone_target(state_relation)
+      }) -%}
+
+      -- reuse the view materialization
+      -- TODO: support actual dispatch for materialization macros
+      {% set search_name = "materialization_view_" ~ adapter.type() %}
+      {% if not search_name in context %}
+          {% set search_name = "materialization_view_default" %}
+      {% endif %}
+      {% set materialization_macro = context[search_name] %}
+      {% set relations = materialization_macro() %}
+      {{ return(relations) }}
+
+  {%- endif -%}
+
+{%- endmaterialization -%}

core/dbt/include/global_project/macros/materializations/models/clone/create_or_replace_clone.sql

@@ -0,0 +1,7 @@
+{% macro create_or_replace_clone(this_relation, state_relation) %}
+    {{ return(adapter.dispatch('create_or_replace_clone', 'dbt')(this_relation, state_relation)) }}
+{% endmacro %}
+
+{% macro default__create_or_replace_clone(this_relation, state_relation) %}
+    create or replace table {{ this_relation }} clone {{ state_relation }}
+{% endmacro %}

core/dbt/include/global_project/macros/materializations/models/clone/get_clone_target.sql

@@ -0,0 +1,10 @@
+{% macro get_clone_target(to_relation) %}
+    {{ return(adapter.dispatch('get_clone_target', 'dbt')(to_relation)) }}
+{% endmacro %}
+
+{% macro default__get_clone_target(to_relation) %}
+    {% set target_sql %}
+        select * from {{ to_relation }}
+    {% endset %}
+    {{ return(target_sql) }}
+{% endmacro %}

core/dbt/task/clone.py

@@ -0,0 +1,183 @@
+import threading
+from typing import AbstractSet, Optional, Any, List, Iterable, Set
+
+from dbt.dataclass_schema import dbtClassMixin
+
+from dbt.contracts.graph.manifest import WritableManifest
+from dbt.contracts.results import RunStatus, RunResult
+from dbt.exceptions import DbtInternalError, DbtRuntimeError, CompilationError
+from dbt.graph import ResourceTypeSelector
+from dbt.node_types import NodeType
+from dbt.parser.manifest import write_manifest
+from dbt.task.base import BaseRunner
+from dbt.task.runnable import GraphRunnableTask
+from dbt.task.run import _validate_materialization_relations_dict
+from dbt.adapters.base import BaseRelation
+from dbt.clients.jinja import MacroGenerator
+from dbt.context.providers import generate_runtime_model_context
+
+
+class CloneRunner(BaseRunner):
+    def before_execute(self):
+        pass
+
+    def after_execute(self, result):
+        pass
+
+    def _build_run_model_result(self, model, context):
+        result = context["load_result"]("main")
+        if result:
+            status = RunStatus.Success
+            message = str(result.response)
+        else:
+            status = RunStatus.Success
+            message = "No-op"
+        adapter_response = {}
+        if result and isinstance(result.response, dbtClassMixin):
+            adapter_response = result.response.to_dict(omit_none=True)
+        return RunResult(
+            node=model,
+            status=status,
+            timing=[],
+            thread_id=threading.current_thread().name,
+            execution_time=0,
+            message=message,
+            adapter_response=adapter_response,
+            failures=None,
+        )
+
+    def compile(self, manifest):
+        # no-op
+        return self.node
+
+    def _materialization_relations(self, result: Any, model) -> List[BaseRelation]:
+        if isinstance(result, str):
+            msg = (
+                'The materialization ("{}") did not explicitly return a '
+                "list of relations to add to the cache.".format(str(model.get_materialization()))
+            )
+            raise CompilationError(msg, node=model)
+
+        if isinstance(result, dict):
+            return _validate_materialization_relations_dict(result, model)
+
+        msg = (
+            "Invalid return value from materialization, expected a dict "
+            'with key "relations", got: {}'.format(str(result))
+        )
+        raise CompilationError(msg, node=model)
+
+    def execute(self, model, manifest):
+        context = generate_runtime_model_context(model, self.config, manifest)
+        materialization_macro = manifest.find_materialization_macro_by_name(
+            self.config.project_name, "clone", self.adapter.type()
+        )
+
+        if "config" not in context:
+            raise DbtInternalError(
+                "Invalid materialization context generated, missing config: {}".format(context)
+            )
+
+        context_config = context["config"]
+
+        hook_ctx = self.adapter.pre_model_hook(context_config)
+        try:
+            result = MacroGenerator(materialization_macro, context)()
+        finally:
+            self.adapter.post_model_hook(context_config, hook_ctx)
+
+        for relation in self._materialization_relations(result, model):
+            self.adapter.cache_added(relation.incorporate(dbt_created=True))
+
+        return self._build_run_model_result(model, context)
+
+
+class CloneTask(GraphRunnableTask):
+    def raise_on_first_error(self):
+        return False
+
+    def get_model_schemas(self, adapter, selected_uids: Iterable[str]) -> Set[BaseRelation]:
+        if self.manifest is None:
+            raise DbtInternalError("manifest was None in get_model_schemas")
+        result: Set[BaseRelation] = set()
+
+        for node in self.manifest.nodes.values():
+            if node.unique_id not in selected_uids:
+                continue
+            if node.is_relational and not node.is_ephemeral:
+                relation = adapter.Relation.create_from(self.config, node)
+                result.add(relation.without_identifier())
+
+                # cache the 'other' schemas too!
+                if node.state_relation:  # type: ignore
+                    other_relation = adapter.Relation.create_from_node(
+                        self.config, node.state_relation  # type: ignore
+                    )
+                    result.add(other_relation.without_identifier())
+
+        return result
+
+    def before_run(self, adapter, selected_uids: AbstractSet[str]):
+        with adapter.connection_named("master"):
+            # unlike in other tasks, we want to add information from the --state manifest *before* caching!
+            self.defer_to_manifest(adapter, selected_uids)
+            # only create *our* schemas, but cache *other* schemas in addition
+            schemas_to_create = super().get_model_schemas(adapter, selected_uids)
+            self.create_schemas(adapter, schemas_to_create)
+            schemas_to_cache = self.get_model_schemas(adapter, selected_uids)
+            self.populate_adapter_cache(adapter, schemas_to_cache)
+
+    @property
+    def resource_types(self):
+        if not self.args.resource_types:
+            return NodeType.refable()
+
+        values = set(self.args.resource_types)
+
+        if "all" in values:
+            values.remove("all")
+            values.update(NodeType.refable())
+
+        values = [NodeType(val) for val in values if val in NodeType.refable()]
+
+        return list(values)
+
+    def get_node_selector(self) -> ResourceTypeSelector:
+        resource_types = self.resource_types
+
+        if self.manifest is None or self.graph is None:
+            raise DbtInternalError("manifest and graph must be set to get perform node selection")
+        return ResourceTypeSelector(
+            graph=self.graph,
+            manifest=self.manifest,
+            previous_state=self.previous_state,
+            resource_types=resource_types,
+        )
+
+    def get_runner_type(self, _):
+        return CloneRunner
+
+    def _get_deferred_manifest(self) -> Optional[WritableManifest]:
+        state = self.previous_state
+        if state is None:
+            raise DbtRuntimeError(
+                "--state is required for cloning relations from another environment"
+            )
+
+        if state.manifest is None:
+            raise DbtRuntimeError(f'Could not find manifest in --state path: "{self.args.state}"')
+        return state.manifest
+
+    # Note that this is different behavior from --defer with other commands, which *merge*
+    # selected nodes from this manifest + unselected nodes from the other manifest
+    def defer_to_manifest(self, adapter, selected_uids: AbstractSet[str]):
+        deferred_manifest = self._get_deferred_manifest()
+        if deferred_manifest is None:
+            return
+        if self.manifest is None:
+            raise DbtInternalError(
+                "Expected to defer to manifest, but there is no runtime manifest to defer from!"
+            )
+        self.manifest.add_from_artifact(other=deferred_manifest)
+        # TODO: is it wrong to write the manifest here? I think it's right...
+        write_manifest(self.manifest, self.config.target_path)

core/dbt/task/retry.py

@@ -9,6 +9,7 @@
 from dbt.graph import GraphQueue
 from dbt.task.base import ConfiguredTask
 from dbt.task.build import BuildTask
+from dbt.task.clone import CloneTask
 from dbt.task.compile import CompileTask
 from dbt.task.generate import GenerateTask
 from dbt.task.run import RunTask
@@ -22,6 +23,7 @@
 TASK_DICT = {
     "build": BuildTask,
     "compile": CompileTask,
+    "clone": CloneTask,
     "generate": GenerateTask,
     "seed": SeedTask,
     "snapshot": SnapshotTask,
@@ -33,6 +35,7 @@
 CMD_DICT = {
     "build": CliCommand.BUILD,
     "compile": CliCommand.COMPILE,
+    "clone": CliCommand.CLONE,
     "generate": CliCommand.DOCS_GENERATE,
     "seed": CliCommand.SEED,
     "snapshot": CliCommand.SNAPSHOT,