341 lines
10 KiB
Python
341 lines
10 KiB
Python
import json
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
from typing import List, Optional, Union
|
|
|
|
from dbt import flags
|
|
from dbt.contracts.graph.manifest import WritableManifest
|
|
from dbt.contracts.rpc import (
|
|
GetManifestParameters,
|
|
GetManifestResult,
|
|
RPCCompileParameters,
|
|
RPCDocsGenerateParameters,
|
|
RPCRunParameters,
|
|
RPCRunOperationParameters,
|
|
RPCSeedParameters,
|
|
RPCTestParameters,
|
|
RemoteCatalogResults,
|
|
RemoteExecutionResult,
|
|
RemoteListResults,
|
|
RemoteRunOperationResult,
|
|
RPCSnapshotParameters,
|
|
RPCSourceFreshnessParameters,
|
|
RPCListParameters,
|
|
RPCBuildParameters,
|
|
)
|
|
from dbt.exceptions import RuntimeException
|
|
from dbt.rpc.method import (
|
|
Parameters, RemoteManifestMethod
|
|
)
|
|
|
|
from dbt.task.base import BaseTask
|
|
from dbt.task.compile import CompileTask
|
|
from dbt.task.freshness import FreshnessTask
|
|
from dbt.task.generate import GenerateTask
|
|
from dbt.task.run import RunTask
|
|
from dbt.task.run_operation import RunOperationTask
|
|
from dbt.task.seed import SeedTask
|
|
from dbt.task.snapshot import SnapshotTask
|
|
from dbt.task.test import TestTask
|
|
from dbt.task.list import ListTask
|
|
from dbt.task.build import BuildTask
|
|
|
|
from .base import RPCTask
|
|
from .cli import HasCLI
|
|
|
|
|
|
class RPCCommandTask(
|
|
RPCTask[Parameters],
|
|
HasCLI[Parameters, RemoteExecutionResult],
|
|
BaseTask,
|
|
):
|
|
@staticmethod
|
|
def _listify(
|
|
value: Optional[Union[str, List[str]]]
|
|
) -> Optional[List[str]]:
|
|
if value is None:
|
|
return None
|
|
elif isinstance(value, str):
|
|
return [value]
|
|
else:
|
|
return value
|
|
|
|
def handle_request(self) -> RemoteExecutionResult:
|
|
return self.run()
|
|
|
|
|
|
def state_path(state: Optional[str]) -> Optional[Path]:
|
|
if state is not None:
|
|
return Path(state)
|
|
elif flags.ARTIFACT_STATE_PATH is not None:
|
|
return Path(flags.ARTIFACT_STATE_PATH)
|
|
else:
|
|
return None
|
|
|
|
|
|
class RemoteCompileProjectTask(
|
|
RPCCommandTask[RPCCompileParameters], CompileTask
|
|
):
|
|
METHOD_NAME = 'compile'
|
|
|
|
def set_args(self, params: RPCCompileParameters) -> None:
|
|
if params.models:
|
|
self.args.select = self._listify(params.models)
|
|
else:
|
|
self.args.select = self._listify(params.select)
|
|
self.args.exclude = self._listify(params.exclude)
|
|
self.args.selector_name = params.selector
|
|
if params.threads is not None:
|
|
self.args.threads = params.threads
|
|
|
|
self.args.state = state_path(params.state)
|
|
|
|
self.set_previous_state()
|
|
|
|
|
|
class RemoteRunProjectTask(RPCCommandTask[RPCRunParameters], RunTask):
|
|
METHOD_NAME = 'run'
|
|
|
|
def set_args(self, params: RPCRunParameters) -> None:
|
|
if params.models:
|
|
self.args.select = self._listify(params.models)
|
|
else:
|
|
self.args.select = self._listify(params.select)
|
|
self.args.exclude = self._listify(params.exclude)
|
|
self.args.selector_name = params.selector
|
|
|
|
if params.threads is not None:
|
|
self.args.threads = params.threads
|
|
if params.defer is None:
|
|
self.args.defer = flags.DEFER_MODE
|
|
else:
|
|
self.args.defer = params.defer
|
|
|
|
self.args.state = state_path(params.state)
|
|
self.set_previous_state()
|
|
|
|
|
|
class RemoteSeedProjectTask(RPCCommandTask[RPCSeedParameters], SeedTask):
|
|
METHOD_NAME = 'seed'
|
|
|
|
def set_args(self, params: RPCSeedParameters) -> None:
|
|
# select has an argparse `dest` value of `models`.
|
|
self.args.select = self._listify(params.select)
|
|
self.args.exclude = self._listify(params.exclude)
|
|
self.args.selector_name = params.selector
|
|
if params.threads is not None:
|
|
self.args.threads = params.threads
|
|
self.args.show = params.show
|
|
|
|
self.args.state = state_path(params.state)
|
|
self.set_previous_state()
|
|
|
|
|
|
class RemoteTestProjectTask(RPCCommandTask[RPCTestParameters], TestTask):
|
|
METHOD_NAME = 'test'
|
|
|
|
def set_args(self, params: RPCTestParameters) -> None:
|
|
if params.models:
|
|
self.args.select = self._listify(params.models)
|
|
else:
|
|
self.args.select = self._listify(params.select)
|
|
self.args.exclude = self._listify(params.exclude)
|
|
self.args.selector_name = params.selector
|
|
self.args.data = params.data
|
|
self.args.schema = params.schema
|
|
if params.threads is not None:
|
|
self.args.threads = params.threads
|
|
if params.defer is None:
|
|
self.args.defer = flags.DEFER_MODE
|
|
else:
|
|
self.args.defer = params.defer
|
|
|
|
self.args.state = state_path(params.state)
|
|
self.set_previous_state()
|
|
|
|
|
|
class RemoteDocsGenerateProjectTask(
|
|
RPCCommandTask[RPCDocsGenerateParameters],
|
|
GenerateTask,
|
|
):
|
|
METHOD_NAME = 'docs.generate'
|
|
|
|
def set_args(self, params: RPCDocsGenerateParameters) -> None:
|
|
self.args.select = None
|
|
self.args.exclude = None
|
|
self.args.selector_name = None
|
|
self.args.compile = params.compile
|
|
|
|
self.args.state = state_path(params.state)
|
|
|
|
def get_catalog_results(
|
|
self, nodes, sources, generated_at, compile_results, errors
|
|
) -> RemoteCatalogResults:
|
|
return RemoteCatalogResults(
|
|
nodes=nodes,
|
|
sources=sources,
|
|
generated_at=datetime.utcnow(),
|
|
_compile_results=compile_results,
|
|
errors=errors,
|
|
logs=[],
|
|
)
|
|
|
|
|
|
class RemoteRunOperationTask(
|
|
RunOperationTask,
|
|
RemoteManifestMethod[RPCRunOperationParameters, RemoteRunOperationResult],
|
|
HasCLI[RPCRunOperationParameters, RemoteRunOperationResult],
|
|
):
|
|
METHOD_NAME = 'run-operation'
|
|
|
|
def __init__(self, args, config, manifest):
|
|
super().__init__(args, config)
|
|
RemoteManifestMethod.__init__(
|
|
self, args, config, manifest # type: ignore
|
|
)
|
|
|
|
def load_manifest(self):
|
|
# we started out with a manifest!
|
|
pass
|
|
|
|
def set_args(self, params: RPCRunOperationParameters) -> None:
|
|
self.args.macro = params.macro
|
|
self.args.args = params.args
|
|
|
|
def _get_kwargs(self):
|
|
if isinstance(self.args.args, dict):
|
|
return self.args.args
|
|
else:
|
|
return RunOperationTask._get_kwargs(self)
|
|
|
|
def _runtime_initialize(self):
|
|
return RunOperationTask._runtime_initialize(self)
|
|
|
|
def handle_request(self) -> RemoteRunOperationResult:
|
|
base = RunOperationTask.run(self)
|
|
result = RemoteRunOperationResult.from_local_result(base=base, logs=[])
|
|
return result
|
|
|
|
def interpret_results(self, results):
|
|
return results.success
|
|
|
|
|
|
class RemoteSnapshotTask(RPCCommandTask[RPCSnapshotParameters], SnapshotTask):
|
|
METHOD_NAME = 'snapshot'
|
|
|
|
def set_args(self, params: RPCSnapshotParameters) -> None:
|
|
# select has an argparse `dest` value of `models`.
|
|
self.args.select = self._listify(params.select)
|
|
self.args.exclude = self._listify(params.exclude)
|
|
self.args.selector_name = params.selector
|
|
if params.threads is not None:
|
|
self.args.threads = params.threads
|
|
|
|
self.args.state = state_path(params.state)
|
|
self.set_previous_state()
|
|
|
|
|
|
class RemoteSourceFreshnessTask(
|
|
RPCCommandTask[RPCSourceFreshnessParameters],
|
|
FreshnessTask
|
|
):
|
|
METHOD_NAME = 'source-freshness'
|
|
|
|
def set_args(self, params: RPCSourceFreshnessParameters) -> None:
|
|
self.args.select = self._listify(params.select)
|
|
self.args.exclude = self._listify(params.exclude)
|
|
self.args.selector_name = params.selector
|
|
if params.threads is not None:
|
|
self.args.threads = params.threads
|
|
self.args.output = None
|
|
|
|
|
|
class RemoteSourceSnapshotFreshnessTask(
|
|
RemoteSourceFreshnessTask
|
|
):
|
|
""" Deprecated task method name, aliases to `source-freshness` """
|
|
METHOD_NAME = 'snapshot-freshness'
|
|
|
|
|
|
# this is a weird and special method.
|
|
class GetManifest(
|
|
RemoteManifestMethod[GetManifestParameters, GetManifestResult]
|
|
):
|
|
METHOD_NAME = 'get-manifest'
|
|
|
|
def set_args(self, params: GetManifestParameters) -> None:
|
|
self.args.select = None
|
|
self.args.exclude = None
|
|
self.args.selector_name = None
|
|
|
|
def handle_request(self) -> GetManifestResult:
|
|
task = RemoteCompileProjectTask(self.args, self.config, self.manifest)
|
|
task.handle_request()
|
|
|
|
manifest: Optional[WritableManifest] = None
|
|
if task.manifest is not None:
|
|
manifest = task.manifest.writable_manifest()
|
|
|
|
return GetManifestResult(
|
|
logs=[],
|
|
manifest=manifest,
|
|
)
|
|
|
|
def interpret_results(self, results):
|
|
return results.manifest is not None
|
|
|
|
|
|
class RemoteListTask(
|
|
RPCCommandTask[RPCListParameters], ListTask
|
|
):
|
|
METHOD_NAME = 'list'
|
|
|
|
def set_args(self, params: RPCListParameters) -> None:
|
|
self.args.output = params.output
|
|
self.args.output_keys = params.output_keys
|
|
self.args.resource_types = self._listify(params.resource_types)
|
|
self.args.models = self._listify(params.models)
|
|
self.args.exclude = self._listify(params.exclude)
|
|
self.args.selector_name = params.selector
|
|
self.args.select = self._listify(params.select)
|
|
self.args.single_threaded = True
|
|
|
|
if self.args.models:
|
|
if self.args.select:
|
|
raise RuntimeException(
|
|
'"models" and "select" are mutually exclusive arguments'
|
|
)
|
|
if self.args.resource_types:
|
|
raise RuntimeException(
|
|
'"models" and "resource_type" are mutually exclusive '
|
|
'arguments'
|
|
)
|
|
|
|
@staticmethod
|
|
def output_results(results):
|
|
return RemoteListResults(
|
|
output=[json.loads(x) for x in results],
|
|
logs=None
|
|
)
|
|
|
|
|
|
class RemoteBuildProjectTask(RPCCommandTask[RPCBuildParameters], BuildTask):
|
|
|
|
METHOD_NAME = 'build'
|
|
|
|
def set_args(self, params: RPCBuildParameters) -> None:
|
|
self.args.resource_types = self._listify(params.resource_types)
|
|
self.args.select = self._listify(params.select)
|
|
self.args.exclude = self._listify(params.exclude)
|
|
self.args.selector_name = params.selector
|
|
|
|
if params.threads is not None:
|
|
self.args.threads = params.threads
|
|
if params.defer is None:
|
|
self.args.defer = flags.DEFER_MODE
|
|
else:
|
|
self.args.defer = params.defer
|
|
|
|
self.args.state = state_path(params.state)
|
|
self.set_previous_state()
|