dbt-selly/dbt-env/lib/python3.8/site-packages/dbt/task/rpc/project_commands.py

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()