Repository: dbt-labs/dbt-core
Pull request: add run-operation + snapshot to the RPC server (#1875)
Fixes #1875 

`snapshot` + `run-operation` support. 


`run-operation`:
In addition to the defaults of 'task_tags' and 'timeout' there are 2 parameters:
    - macro: str (required)
    - args: Dict[str, Any] (optional)

As in the CLI version, run operations return a `success` boolean that indicates whether the macro raised an exception or not. They will only return an `error` object if the request itself is somehow wrong. The result returned by the macro is discarded.


`snapshot`:
This has the same arguments as `compile` and `run`. It executes the `dbt snapshot` task.


Both methods work via `cli_args`.

Problem statement:
Support run-operation and snapshot methods in rpc server
### Describe the feature
The rpc server should support the `run_operation` and `snapshot` methods. Additionally, the `cli_args` method should support running operations and snapshots using the CLI syntax.

### Methods

*run-operation*:
- args

*snapshot*:
 - select
 - exclude

### Who will this benefit?
Parity for all methods (except for diagnostic ones (ls, debug)) between the RPC server and the CLI for dbt.

Hints from issue discussion:

Commit message: add run_operation + snapshot to the RPC server

Patch to review:
diff --git a/core/dbt/contracts/rpc.py b/core/dbt/contracts/rpc.py
index 226d8982aad..8c5aed779a9 100644
--- a/core/dbt/contracts/rpc.py
+++ b/core/dbt/contracts/rpc.py
@@ -116,9 +116,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 +167,11 @@ class ResultTable(JsonSchemaMixin):
     rows: List[Any]
 
 
+@dataclass
+class RemoteRunOperationResult(RemoteResult):
+    success: bool
+
+
 @dataclass
 class RemoteRunResult(RemoteCompileResult):
     table: ResultTable
@@ -431,6 +442,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(
diff --git a/core/dbt/main.py b/core/dbt/main.py
index 7f79afe4c3f..9d03d2830d3 100644
--- a/core/dbt/main.py
+++ b/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
 
 
diff --git a/core/dbt/rpc/builtins.py b/core/dbt/rpc/builtins.py
index 25edc9de7dc..a3601c54778 100644
--- a/core/dbt/rpc/builtins.py
+++ b/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)
diff --git a/core/dbt/task/rpc/project_commands.py b/core/dbt/task/rpc/project_commands.py
index 93793c7f2ed..e2c2f4a5532 100644
--- a/core/dbt/task/rpc/project_commands.py
+++ b/core/dbt/task/rpc/project_commands.py
@@ -5,10 +5,12 @@
 from dbt.contracts.rpc import (
     RPCCompileParameters,
     RPCDocsGenerateParameters,
+    RPCRunOperationParameters,
     RPCSeedParameters,
     RPCTestParameters,
     RemoteCatalogResults,
     RemoteExecutionResult,
+    RemoteRunOperationResult,
 )
 from dbt.rpc.method import (
     Parameters,
@@ -16,7 +18,9 @@
 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 +101,40 @@ 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[RPCCompileParameters], SnapshotTask):
+    METHOD_NAME = 'snapshot'
+
+    def set_args(self, params: RPCCompileParameters) -> None:
+        self.args.models = self._listify(params.models)
+        self.args.exclude = self._listify(params.exclude)
diff --git a/core/dbt/utils.py b/core/dbt/utils.py
index 4299cfd2c70..d5d924947b9 100644
--- a/core/dbt/utils.py
+++ b/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)
 
 
diff --git a/test/integration/048_rpc_test/test_execute_fetch_and_serialize.py b/test/integration/048_rpc_test/test_execute_fetch_and_serialize.py
index 73d140a84d0..7850818909d 100644
--- a/test/integration/048_rpc_test/test_execute_fetch_and_serialize.py
+++ b/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')
diff --git a/test/rpc/fixtures.py b/test/rpc/fixtures.py
index c858aad4bcc..c5e1820a50d 100644
--- a/test/rpc/fixtures.py
+++ b/test/rpc/fixtures.py
@@ -216,6 +216,19 @@ def run(
             method='run', params=params, request_id=request_id
         )
 
+    def run_operation(
+        self,
+        macro: str,
+        args: Optional[Dict[str, Any]],
+        request_id: int = 1,
+    ):
+        params = {'macro': macro}
+        if args is not None:
+            params['args'] = args
+        return self.request(
+            method='run-operation', params=params, request_id=request_id
+        )
+
     def seed(self, show: bool = None, request_id: int = 1):
         params = {}
         if show is not None:
@@ -224,6 +237,21 @@ def seed(self, show: bool = None, request_id: int = 1):
             method='seed', params=params, request_id=request_id
         )
 
+    def snapshot(
+        self,
+        models: Optional[Union[str, List[str]]] = None,
+        exclude: Optional[Union[str, List[str]]] = None,
+        request_id: int = 1,
+    ):
+        params = {}
+        if models is not None:
+            params['models'] = models
+        if exclude is not None:
+            params['exclude'] = exclude
+        return self.request(
+            method='snapshot', params=params, request_id=request_id
+        )
+
     def test(
         self,
         models: Optional[Union[str, List[str]]] = None,
@@ -241,7 +269,7 @@ def test(
             params['data'] = data
         if schema is not None:
             params['schema'] = schema
-        return self.requuest(
+        return self.request(
             method='test', params=params, request_id=request_id
         )
 
@@ -419,6 +447,7 @@ def __init__(
         packages=None,
         models=None,
         macros=None,
+        snapshots=None,
     ):
         self.project = {
             'name': name,
@@ -430,6 +459,7 @@ def __init__(
         self.packages = packages
         self.models = models
         self.macros = macros
+        self.snapshots = snapshots
 
     def _write_recursive(self, path, inputs):
         for name, value in inputs.items():
@@ -460,19 +490,22 @@ def write_config(self, project_dir, remove=False):
             cfg.remove()
         cfg.write(yaml.safe_dump(self.project))
 
-    def write_models(self, project_dir, remove=False):
+    def _write_values(self, project_dir, remove, name, value):
         if remove:
-            project_dir.join('models').remove()
+            project_dir.join(name).remove()
+
+        if value is not None:
+            self._write_recursive(project_dir.mkdir(name), value)
 
-        if self.models is not None:
-            self._write_recursive(project_dir.mkdir('models'), self.models)
+
+    def write_models(self, project_dir, remove=False):
+        self._write_values(project_dir, remove, 'models', self.models)
 
     def write_macros(self, project_dir, remove=False):
-        if remove:
-            project_dir.join('macros').remove()
+        self._write_values(project_dir, remove, 'macros', self.macros)
 
-        if self.macros is not None:
-            self._write_recursive(project_dir.mkdir('macros'), self.macros)
+    def write_snapshots(self, project_dir, remove=False):
+        self._write_values(project_dir, remove, 'snapshots', self.snapshots)
 
     def write_to(self, project_dir, remove=False):
         if remove:
@@ -482,6 +515,7 @@ def write_to(self, project_dir, remove=False):
         self.write_config(project_dir)
         self.write_models(project_dir)
         self.write_macros(project_dir)
+        self.write_snapshots(project_dir)
 
 
 class TestArgs:
diff --git a/test/rpc/test_base.py b/test/rpc/test_base.py
index d7c6255123c..ed608014eba 100644
--- a/test/rpc/test_base.py
+++ b/test/rpc/test_base.py
@@ -296,3 +296,152 @@ def test_ps_poll_output_match(project_dir, profiles_dir, postgres_profile, uniqu
 
         for key in ('start', 'end', 'elapsed', 'state'):
             assert ps_result[key] == poll_result[key]
+
+
+macros_data = '''
+{% macro foo() %}
+    {{ return(1) }}
+{% endmacro %}
+{% macro bar(value) %}
+    {{ return(value + 1) }}
+{% endmacro %}
+{% macro quux(value) %}
+    {{ return(asdf) }}
+{% endmacro %}
+'''
+
+
+def test_run_operation(project_dir, profiles_dir, postgres_profile, unique_schema):
+    project = ProjectDefinition(
+        models={'my_model.sql': 'select 1 as id'},
+        macros={
+            'my_macros.sql': macros_data,
+        }
+    )
+    server_ctx = rpc_server(
+        project_dir=project_dir, schema=unique_schema, profiles_dir=profiles_dir
+    )
+    schema_ctx = built_schema(
+        project_dir=project_dir, schema=unique_schema, profiles_dir=profiles_dir, test_kwargs={}, project_def=project,
+    )
+    with schema_ctx, server_ctx as server:
+        querier = Querier(server)
+        token = querier.is_async_result(querier.run_operation(macro='foo', args={}))
+        poll_result = querier.is_result(querier.async_wait(token))
+
+        assert 'success' in poll_result
+        assert poll_result['success'] is True
+
+        token = querier.is_async_result(querier.run_operation(macro='bar', args={'value': 10}))
+        poll_result = querier.is_result(querier.async_wait(token))
+
+        assert 'success' in poll_result
+        assert poll_result['success'] is True
+
+        token = querier.is_async_result(querier.run_operation(macro='baz', args={}))
+        poll_result = querier.is_result(querier.async_wait(token, state='failed'))
+        assert 'state' in poll_result
+        assert poll_result['state'] == 'failed'
+
+        token = querier.is_async_result(querier.run_operation(macro='quux', args={}))
+        poll_result = querier.is_result(querier.async_wait(token))
+        assert 'success' in poll_result
+        assert poll_result['success'] is True
+
+
+def test_run_operation_cli(project_dir, profiles_dir, postgres_profile, unique_schema):
+    project = ProjectDefinition(
+        models={'my_model.sql': 'select 1 as id'},
+        macros={
+            'my_macros.sql': macros_data,
+        }
+    )
+    server_ctx = rpc_server(
+        project_dir=project_dir, schema=unique_schema, profiles_dir=profiles_dir
+    )
+    schema_ctx = built_schema(
+        project_dir=project_dir, schema=unique_schema, profiles_dir=profiles_dir, test_kwargs={}, project_def=project,
+    )
+    with schema_ctx, server_ctx as server:
+        querier = Querier(server)
+        token = querier.is_async_result(querier.cli_args(cli='run-operation foo'))
+        poll_result = querier.is_result(querier.async_wait(token))
+
+        assert 'success' in poll_result
+        assert poll_result['success'] is True
+
+        token = querier.is_async_result(querier.cli_args(cli='''run-operation bar --args="{'value': 10}"'''))
+        poll_result = querier.is_result(querier.async_wait(token))
+
+        assert 'success' in poll_result
+        assert poll_result['success'] is True
+
+        token = querier.is_async_result(querier.cli_args(cli='run-operation baz'))
+        poll_result = querier.is_result(querier.async_wait(token, state='failed'))
+        assert 'state' in poll_result
+        assert poll_result['state'] == 'failed'
+
+        token = querier.is_async_result(querier.cli_args(cli='run-operation quux'))
+        poll_result = querier.is_result(querier.async_wait(token))
+        assert 'success' in poll_result
+        assert poll_result['success'] is True
+
+
+snapshot_data = '''
+{% snapshot snapshot_actual %}
+
+    {{
+        config(
+            target_database=database,
+            target_schema=schema,
+            unique_key='id',
+            strategy='timestamp',
+            updated_at='updated_at',
+        )
+    }}
+    select 1 as id, '2019-10-31 23:59:40' as updated_at
+
+{% endsnapshot %}
+'''
+
+
+def test_snapshots(project_dir, profiles_dir, postgres_profile, unique_schema):
+    project = ProjectDefinition(
+        snapshots={'my_snapshots.sql': snapshot_data},
+    )
+    server_ctx = rpc_server(
+        project_dir=project_dir, schema=unique_schema, profiles_dir=profiles_dir
+    )
+    schema_ctx = built_schema(
+        project_dir=project_dir, schema=unique_schema, profiles_dir=profiles_dir, test_kwargs={}, project_def=project,
+    )
+    with schema_ctx, server_ctx as server:
+        querier = Querier(server)
+        token = querier.is_async_result(querier.snapshot())
+        results = querier.is_result(querier.async_wait(token))
+        assert len(results['results']) == 1
+
+        token = querier.is_async_result(querier.snapshot(exclude=['snapshot_actual']))
+        results = querier.is_result(querier.async_wait(token))
+        assert len(results['results']) == 0
+
+
+def test_snapshots_cli(project_dir, profiles_dir, postgres_profile, unique_schema):
+    project = ProjectDefinition(
+        snapshots={'my_snapshots.sql': snapshot_data},
+    )
+    server_ctx = rpc_server(
+        project_dir=project_dir, schema=unique_schema, profiles_dir=profiles_dir
+    )
+    schema_ctx = built_schema(
+        project_dir=project_dir, schema=unique_schema, profiles_dir=profiles_dir, test_kwargs={}, project_def=project,
+    )
+    with schema_ctx, server_ctx as server:
+        querier = Querier(server)
+        token = querier.is_async_result(querier.cli_args(cli='snapshot'))
+        results = querier.is_result(querier.async_wait(token))
+        assert len(results['results']) == 1
+
+        token = querier.is_async_result(querier.cli_args(cli='snapshot --exclude=snapshot_actual'))
+        results = querier.is_result(querier.async_wait(token))
+        assert len(results['results']) == 0


Code review comments:
diff_hunk: @@ -97,3 +101,40 @@ 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[RPCCompileParameters], SnapshotTask):
+    METHOD_NAME = 'snapshot'
+
+    def set_args(self, params: RPCCompileParameters) -> None:
+        self.args.models = self._listify(params.models)
path: core/dbt/task/rpc/project_commands.py
text: @user1:
I think the input to `snapshot` on the CLI is `-s --select`, not `--models` - can we update this accordingly?