add run-operation + snapshot to the RPC server (#1875)

core/dbt/contracts/rpc.py

@@ -45,6 +45,12 @@ class RPCCompileParameters(RPCParameters):
     exclude: Union[None, str, List[str]] = None
 
 
+@dataclass
+class RPCSnapshotParameters(RPCParameters):
+    select: Union[None, str, List[str]] = None
+    exclude: Union[None, str, List[str]] = None
+
+
 @dataclass
 class RPCTestParameters(RPCCompileParameters):
     data: bool = False
@@ -116,9 +122,15 @@ class GCParameters(RPCParameters):
         will be applied to the task manager before GC starts. By default the
         existing gc settings remain.
     """
-    task_ids: Optional[List[TaskID]]
-    before: Optional[datetime]
-    settings: Optional[GCSettings]
+    task_ids: Optional[List[TaskID]] = None
+    before: Optional[datetime] = None
+    settings: Optional[GCSettings] = None
+
+
+@dataclass
+class RPCRunOperationParameters(RPCParameters):
+    macro: str
+    args: Dict[str, Any] = field(default_factory=dict)
 
 
 # Outputs
@@ -161,6 +173,11 @@ class ResultTable(JsonSchemaMixin):
     rows: List[Any]
 
 
+@dataclass
+class RemoteRunOperationResult(RemoteResult):
+    success: bool
+
+
 @dataclass
 class RemoteRunResult(RemoteCompileResult):
     table: ResultTable
@@ -431,6 +448,31 @@ def from_result(
         )
 
 
+@dataclass
+class PollRunOperationCompleteResult(RemoteRunOperationResult, PollResult):
+    state: TaskHandlerState = field(
+        metadata=restrict_to(TaskHandlerState.Success,
+                             TaskHandlerState.Failed),
+    )
+
+    @classmethod
+    def from_result(
+        cls: Type['PollRunOperationCompleteResult'],
+        base: RemoteRunOperationResult,
+        tags: TaskTags,
+        timing: TaskTiming,
+    ) -> 'PollRunOperationCompleteResult':
+        return cls(
+            success=base.success,
+            logs=base.logs,
+            tags=tags,
+            state=timing.state,
+            start=timing.start,
+            end=timing.end,
+            elapsed=timing.elapsed,
+        )
+
+
 @dataclass
 class PollCatalogCompleteResult(RemoteCatalogResults, PollResult):
     state: TaskHandlerState = field(

core/dbt/main.py

@@ -386,7 +386,7 @@ def _build_snapshot_subparser(subparsers, base_subparser):
         '''
     )
     sub.set_defaults(cls=snapshot_task.SnapshotTask, which='snapshot',
-                     rpc_method=None)
+                     rpc_method='snapshot')
     return sub
 
 
@@ -707,7 +707,7 @@ def _build_run_operation_subparser(subparsers, base_subparser):
         '''
     )
     sub.set_defaults(cls=run_operation_task.RunOperationTask,
-                     which='run-operation', rpc_method=None)
+                     which='run-operation', rpc_method='run-operation')
     return sub
 
 

core/dbt/rpc/builtins.py

@@ -21,6 +21,7 @@
     RemoteCompileResult,
     RemoteCatalogResults,
     RemoteEmptyResult,
+    RemoteRunOperationResult,
     PollParameters,
     PollResult,
     PollInProgressResult,
@@ -30,6 +31,7 @@
     PollCompileCompleteResult,
     PollCatalogCompleteResult,
     PollRemoteEmptyCompleteResult,
+    PollRunOperationCompleteResult,
     TaskHandlerState,
     TaskTiming,
 )
@@ -141,6 +143,7 @@ def poll_complete(
         PollCompileCompleteResult,
         PollCatalogCompleteResult,
         PollRemoteEmptyCompleteResult,
+        PollRunOperationCompleteResult,
     ]]
 
     if isinstance(result, RemoteExecutionResult):
@@ -154,6 +157,8 @@ def poll_complete(
         cls = PollCatalogCompleteResult
     elif isinstance(result, RemoteEmptyResult):
         cls = PollRemoteEmptyCompleteResult
+    elif isinstance(result, RemoteRunOperationResult):
+        cls = PollRunOperationCompleteResult
     else:
         raise dbt.exceptions.InternalException(
             'got invalid result in poll_complete: {}'.format(result)

core/dbt/task/rpc/project_commands.py

@@ -5,18 +5,23 @@
 from dbt.contracts.rpc import (
     RPCCompileParameters,
     RPCDocsGenerateParameters,
+    RPCRunOperationParameters,
     RPCSeedParameters,
     RPCTestParameters,
     RemoteCatalogResults,
     RemoteExecutionResult,
+    RemoteRunOperationResult,
+    RPCSnapshotParameters,
 )
 from dbt.rpc.method import (
     Parameters,
 )
 from dbt.task.compile import CompileTask
 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 .base import RPCTask
@@ -97,3 +102,41 @@ def get_catalog_results(
             _compile_results=compile_results,
             logs=[],
         )
+
+
+class RemoteRunOperationTask(
+    RPCTask[RPCRunOperationParameters],
+    HasCLI[RPCRunOperationParameters, RemoteRunOperationResult],
+    RunOperationTask,
+):
+    METHOD_NAME = 'run-operation'
+
+    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:
+        success, _ = RunOperationTask.run(self)
+        result = RemoteRunOperationResult(logs=[], success=success)
+        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.models = self._listify(params.select)
+        self.args.exclude = self._listify(params.exclude)

core/dbt/utils.py

@@ -455,6 +455,10 @@ def default(self, obj):
             return float(obj)
         if isinstance(obj, (datetime.datetime, datetime.date, datetime.time)):
             return obj.isoformat()
+        if hasattr(obj, 'to_dict'):
+            # if we have a to_dict we should try to serialize the result of
+            # that!
+            obj = obj.to_dict()
         return super().default(obj)
 
 

plugins/snowflake/setup.py

@@ -31,6 +31,9 @@
     install_requires=[
         'dbt-core=={}'.format(package_version),
         'snowflake-connector-python>=1.6.12,<2.1',
+        'azure-storage-blob~=2.1',
+        'azure-storage-common~=2.1',
+
     ],
     zip_safe=False,
 )

test/integration/048_rpc_test/test_execute_fetch_and_serialize.py

@@ -2,6 +2,7 @@
 import pickle
 import os
 
+
 class TestRpcExecuteReturnsResults(DBTIntegrationTest):
 
     @property
@@ -26,7 +27,7 @@ def test_pickle(self, agate_table):
 
         pickle.dumps(table)
 
-    def test_file(self, filename):
+    def do_test_file(self, filename):
         file_path = os.path.join("sql", filename)
         with open(file_path) as fh:
             query = fh.read()
@@ -39,12 +40,12 @@ def test_file(self, filename):
 
     @use_profile('bigquery')
     def test__bigquery_fetch_and_serialize(self):
-        self.test_file('bigquery.sql')
+        self.do_test_file('bigquery.sql')
 
     @use_profile('snowflake')
     def test__snowflake_fetch_and_serialize(self):
-        self.test_file('snowflake.sql')
+        self.do_test_file('snowflake.sql')
 
     @use_profile('redshift')
     def test__redshift_fetch_and_serialize(self):
-        self.test_file('redshift.sql')
+        self.do_test_file('redshift.sql')

test/rpc/conftest.py

@@ -0,0 +1,54 @@
+import os
+import pytest
+import random
+import time
+from typing import Dict, Any
+
+import yaml
+
+
+@pytest.fixture
+def unique_schema() -> str:
+    return "test{}{:04}".format(int(time.time()), random.randint(0, 9999))
+
+
+@pytest.fixture
+def profiles_root(tmpdir):
+    return tmpdir.mkdir('profile')
+
+
+@pytest.fixture
+def project_root(tmpdir):
+    return tmpdir.mkdir('project')
+
+
+@pytest.fixture
+def postgres_profile_data(unique_schema):
+    return {
+        'config': {
+            'send_anonymous_usage_stats': False
+        },
+        'test': {
+            'outputs': {
+                'default': {
+                    'type': 'postgres',
+                    'threads': 4,
+                    'host': 'database',
+                    'port': 5432,
+                    'user': 'root',
+                    'pass': 'password',
+                    'dbname': 'dbt',
+                    'schema': unique_schema,
+                },
+            },
+            'target': 'default'
+        }
+    }
+
+
+@pytest.fixture
+def postgres_profile(profiles_root, postgres_profile_data) -> Dict[str, Any]:
+    path = os.path.join(profiles_root, 'profiles.yml')
+    with open(path, 'w') as fp:
+        fp.write(yaml.safe_dump(postgres_profile_data))
+    return postgres_profile_data