dvt-core 1.11.0b4__py3-none-any.whl

dvt/task/seed.py

@@ -0,0 +1,104 @@
+import random
+from typing import Optional, Type
+
+from dvt.artifacts.schemas.results import NodeStatus, RunStatus
+from dvt.contracts.graph.manifest import Manifest
+from dvt.events.types import LogSeedResult, LogStartLine, SeedHeader
+from dvt.graph import ResourceTypeSelector
+from dvt.node_types import NodeType
+from dvt.task import group_lookup
+from dvt.task.base import BaseRunner
+from dvt.task.printer import print_run_end_messages
+from dvt.task.run import ModelRunner, RunTask
+
+from dbt_common.events.base_types import EventLevel
+from dbt_common.events.functions import fire_event
+from dbt_common.events.types import Formatting
+from dbt_common.exceptions import DbtInternalError
+
+
+class SeedRunner(ModelRunner):
+    def describe_node(self) -> str:
+        return "seed file {}".format(self.get_node_representation())
+
+    def before_execute(self) -> None:
+        fire_event(
+            LogStartLine(
+                description=self.describe_node(),
+                index=self.node_index,
+                total=self.num_nodes,
+                node_info=self.node.node_info,
+            )
+        )
+
+    def _build_run_model_result(self, model, context):
+        result = super()._build_run_model_result(model, context)
+        agate_result = context["load_result"]("agate_table")
+        result.agate_table = agate_result.table
+        return result
+
+    def compile(self, manifest: Manifest):
+        return self.node
+
+    def print_result_line(self, result):
+        model = result.node
+        group = group_lookup.get(model.unique_id)
+        level = EventLevel.ERROR if result.status == NodeStatus.Error else EventLevel.INFO
+        fire_event(
+            LogSeedResult(
+                status=result.status,
+                result_message=result.message,
+                index=self.node_index,
+                total=self.num_nodes,
+                execution_time=result.execution_time,
+                schema=self.node.schema,
+                relation=model.alias,
+                node_info=model.node_info,
+                group=group,
+            ),
+            level=level,
+        )
+
+
+class SeedTask(RunTask):
+    def raise_on_first_error(self) -> bool:
+        return False
+
+    def get_node_selector(self):
+        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=[NodeType.Seed],
+        )
+
+    def get_runner_type(self, _) -> Optional[Type[BaseRunner]]:
+        return SeedRunner
+
+    def task_end_messages(self, results) -> None:
+        if self.args.show:
+            self.show_tables(results)
+
+        print_run_end_messages(results)
+
+    def show_table(self, result):
+        table = result.agate_table
+        rand_table = table.order_by(lambda x: random.random())
+
+        schema = result.node.schema
+        alias = result.node.alias
+
+        header = "Random sample of table: {}.{}".format(schema, alias)
+        fire_event(Formatting(""))
+        fire_event(SeedHeader(header=header))
+        fire_event(Formatting("-" * len(header)))
+
+        rand_table.print_table(max_rows=10, max_columns=None)
+        fire_event(Formatting(""))
+
+    def show_tables(self, results):
+        for result in results:
+            if result.status != RunStatus.Error:
+                self.show_table(result)

dvt/task/show.py

@@ -0,0 +1,150 @@
+import io
+import threading
+import time
+
+from dvt.artifacts.schemas.run import RunResult, RunStatus
+from dvt.context.providers import generate_runtime_model_context
+from dvt.contracts.graph.nodes import SeedNode
+from dvt.events.types import ShowNode
+from dvt.flags import get_flags
+from dvt.task.base import ConfiguredTask
+from dvt.task.compile import CompileRunner, CompileTask
+from dvt.task.seed import SeedRunner
+
+from dbt.adapters.factory import get_adapter
+from dbt_common.events.base_types import EventLevel
+from dbt_common.events.functions import fire_event
+from dbt_common.events.types import Note
+from dbt_common.exceptions import DbtRuntimeError
+
+
+class ShowRunner(CompileRunner):
+    def __init__(self, config, adapter, node, node_index, num_nodes) -> None:
+        super().__init__(config, adapter, node, node_index, num_nodes)
+        self.run_ephemeral_models = True
+
+    def execute(self, compiled_node, manifest):
+        start_time = time.time()
+
+        # Allow passing in -1 (or any negative number) to get all rows
+        limit = None if self.config.args.limit < 0 else self.config.args.limit
+
+        model_context = generate_runtime_model_context(compiled_node, self.config, manifest)
+        compiled_node.compiled_code = self.adapter.execute_macro(
+            macro_name="get_show_sql",
+            macro_resolver=manifest,
+            context_override=model_context,
+            kwargs={
+                "compiled_code": model_context["compiled_code"],
+                "sql_header": model_context["config"].get("sql_header"),
+                "limit": limit,
+            },
+        )
+        adapter_response, execute_result = self.adapter.execute(
+            compiled_node.compiled_code, fetch=True
+        )
+
+        end_time = time.time()
+
+        return RunResult(
+            node=compiled_node,
+            status=RunStatus.Success,
+            timing=[],
+            thread_id=threading.current_thread().name,
+            execution_time=end_time - start_time,
+            message=None,
+            adapter_response=adapter_response.to_dict(),
+            agate_table=execute_result,
+            failures=None,
+            batch_results=None,
+        )
+
+
+class ShowTask(CompileTask):
+    def _runtime_initialize(self):
+        if not (self.args.select or getattr(self.args, "inline", None)):
+            raise DbtRuntimeError("Either --select or --inline must be passed to show")
+        super()._runtime_initialize()
+
+    def get_runner_type(self, node):
+        if isinstance(node, SeedNode):
+            return SeedRunner
+        else:
+            return ShowRunner
+
+    def task_end_messages(self, results) -> None:
+        is_inline = bool(getattr(self.args, "inline", None))
+
+        if is_inline:
+            matched_results = [result for result in results if result.node.name == "inline_query"]
+        else:
+            matched_results = []
+            for result in results:
+                if result.node.name in self.selection_arg[0]:
+                    matched_results.append(result)
+                else:
+                    fire_event(
+                        Note(msg=f"Excluded node '{result.node.name}' from results"),
+                        EventLevel.DEBUG,
+                    )
+
+        for result in matched_results:
+            table = result.agate_table
+
+            # Hack to get Agate table output as string
+            output = io.StringIO()
+            if self.args.output == "json":
+                table.to_json(path=output)
+            else:
+                table.print_table(output=output, max_rows=None)
+
+            node_name = result.node.name
+
+            if hasattr(result.node, "version") and result.node.version:
+                node_name += f".v{result.node.version}"
+
+            fire_event(
+                ShowNode(
+                    node_name=node_name,
+                    preview=output.getvalue(),
+                    is_inline=is_inline,
+                    output_format=self.args.output,
+                    unique_id=result.node.unique_id,
+                    quiet=get_flags().QUIET,
+                )
+            )
+
+    def _handle_result(self, result) -> None:
+        super()._handle_result(result)
+
+        if (
+            result.node.is_ephemeral_model
+            and type(self) is ShowTask
+            and (self.args.select or getattr(self.args, "inline", None))
+        ):
+            self.node_results.append(result)
+
+
+class ShowTaskDirect(ConfiguredTask):
+    def run(self):
+        adapter = get_adapter(self.config)
+        with adapter.connection_named("show", should_release_connection=False):
+            limit = None if self.args.limit < 0 else self.args.limit
+            response, table = adapter.execute(self.args.inline_direct, fetch=True, limit=limit)
+
+            output = io.StringIO()
+            if self.args.output == "json":
+                table.to_json(path=output)
+            else:
+                table.print_table(output=output, max_rows=None)
+
+            fire_event(
+                ShowNode(
+                    node_name="direct-query",
+                    preview=output.getvalue(),
+                    is_inline=True,
+                    output_format=self.args.output,
+                    unique_id="direct-query",
+                    quiet=get_flags().QUIET,
+                )
+            )

dvt/task/snapshot.py

@@ -0,0 +1,57 @@
+from typing import Optional, Type
+
+from dvt.artifacts.schemas.results import NodeStatus
+from dvt.events.types import LogSnapshotResult
+from dvt.graph import ResourceTypeSelector
+from dvt.node_types import NodeType
+from dvt.task import group_lookup
+from dvt.task.base import BaseRunner
+from dvt.task.run import ModelRunner, RunTask
+
+from dbt_common.events.base_types import EventLevel
+from dbt_common.events.functions import fire_event
+from dbt_common.exceptions import DbtInternalError
+from dbt_common.utils import cast_dict_to_dict_of_strings
+
+
+class SnapshotRunner(ModelRunner):
+    def describe_node(self) -> str:
+        return "snapshot {}".format(self.get_node_representation())
+
+    def print_result_line(self, result):
+        model = result.node
+        group = group_lookup.get(model.unique_id)
+        cfg = model.config.to_dict(omit_none=True)
+        level = EventLevel.ERROR if result.status == NodeStatus.Error else EventLevel.INFO
+        fire_event(
+            LogSnapshotResult(
+                status=result.status,
+                description=self.get_node_representation(),
+                cfg=cast_dict_to_dict_of_strings(cfg),
+                index=self.node_index,
+                total=self.num_nodes,
+                execution_time=result.execution_time,
+                node_info=model.node_info,
+                result_message=result.message,
+                group=group,
+            ),
+            level=level,
+        )
+
+
+class SnapshotTask(RunTask):
+    def raise_on_first_error(self) -> bool:
+        return False
+
+    def get_node_selector(self):
+        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=[NodeType.Snapshot],
+        )
+
+    def get_runner_type(self, _) -> Optional[Type[BaseRunner]]:
+        return SnapshotRunner

dvt/task/sql.py

@@ -0,0 +1,111 @@
+import traceback
+from abc import abstractmethod
+from datetime import datetime, timezone
+from typing import Generic, TypeVar
+
+import dvt.exceptions
+from dvt.contracts.graph.manifest import Manifest
+from dvt.contracts.sql import (
+    RemoteCompileResult,
+    RemoteCompileResultMixin,
+    RemoteRunResult,
+    ResultTable,
+)
+from dvt.events.types import SQLRunnerException
+from dvt.task.compile import CompileRunner
+
+import dbt_common.exceptions.base
+from dbt_common.events.functions import fire_event
+
+SQLResult = TypeVar("SQLResult", bound=RemoteCompileResultMixin)
+
+
+class GenericSqlRunner(CompileRunner, Generic[SQLResult]):
+    def __init__(self, config, adapter, node, node_index, num_nodes) -> None:
+        CompileRunner.__init__(self, config, adapter, node, node_index, num_nodes)
+
+    def handle_exception(self, e, ctx):
+        fire_event(
+            SQLRunnerException(
+                exc=str(e), exc_info=traceback.format_exc(), node_info=self.node.node_info
+            )
+        )
+        # REVIEW: This code is invalid and will always throw.
+        if isinstance(e, dbt.exceptions.Exception):
+            if isinstance(e, dbt_common.exceptions.DbtRuntimeError):
+                e.add_node(ctx.node)
+            return e
+
+    def before_execute(self) -> None:
+        pass
+
+    def after_execute(self, result) -> None:
+        pass
+
+    def compile(self, manifest: Manifest):
+        return self.compiler.compile_node(self.node, manifest, {}, write=False)
+
+    @abstractmethod
+    def execute(self, compiled_node, manifest) -> SQLResult:
+        pass
+
+    @abstractmethod
+    def from_run_result(self, result, start_time, timing_info) -> SQLResult:
+        pass
+
+    def error_result(self, node, error, start_time, timing_info):
+        raise error
+
+    def ephemeral_result(self, node, start_time, timing_info):
+        raise dbt_common.exceptions.base.NotImplementedError(
+            "cannot execute ephemeral nodes remotely!"
+        )
+
+
+class SqlCompileRunner(GenericSqlRunner[RemoteCompileResult]):
+    def execute(self, compiled_node, manifest) -> RemoteCompileResult:
+        return RemoteCompileResult(
+            raw_code=compiled_node.raw_code,
+            compiled_code=compiled_node.compiled_code,
+            node=compiled_node,
+            timing=[],  # this will get added later
+            generated_at=datetime.now(timezone.utc).replace(tzinfo=None),
+        )
+
+    def from_run_result(self, result, start_time, timing_info) -> RemoteCompileResult:
+        return RemoteCompileResult(
+            raw_code=result.raw_code,
+            compiled_code=result.compiled_code,
+            node=result.node,
+            timing=timing_info,
+            generated_at=datetime.now(timezone.utc).replace(tzinfo=None),
+        )
+
+
+class SqlExecuteRunner(GenericSqlRunner[RemoteRunResult]):
+    def execute(self, compiled_node, manifest) -> RemoteRunResult:
+        _, execute_result = self.adapter.execute(compiled_node.compiled_code, fetch=True)
+
+        table = ResultTable(
+            column_names=list(execute_result.column_names),
+            rows=[list(row) for row in execute_result],
+        )
+
+        return RemoteRunResult(
+            raw_code=compiled_node.raw_code,
+            compiled_code=compiled_node.compiled_code,
+            node=compiled_node,
+            table=table,
+            timing=[],
+            generated_at=datetime.now(timezone.utc).replace(tzinfo=None),
+        )
+
+    def from_run_result(self, result, start_time, timing_info) -> RemoteRunResult:
+        return RemoteRunResult(
+            raw_code=result.raw_code,
+            compiled_code=result.compiled_code,
+            node=result.node,
+            table=result.table,
+            timing=timing_info,
+            generated_at=datetime.now(timezone.utc).replace(tzinfo=None),
+        )