Repository: dbt-labs/dbt-core

Issue:
Expose state and defer arguments to RPC run method
### Describe the feature
The new `--state` and `--defer` arguments are available via the CLI and the environment, but aren't available to RPC `run` method. It should be!

We should also add tests for using it via the RPC server - it happens to work for `cli_args` now, but that's just because of how that method handles arguments.

### Describe alternatives you've considered
`cli_args` works with `--state` and `--defer`, so you can use that.

### Who will this benefit?
RPC users.



Discussion:

Resolving patch:
diff --git a/core/dbt/contracts/rpc.py b/core/dbt/contracts/rpc.py
--- a/core/dbt/contracts/rpc.py
+++ b/core/dbt/contracts/rpc.py
@@ -45,6 +45,17 @@ class RPCCompileParameters(RPCParameters):
     models: Union[None, str, List[str]] = None
     exclude: Union[None, str, List[str]] = None
     selector: Optional[str] = None
+    state: Optional[str] = None
+
+
+@dataclass
+class RPCRunParameters(RPCParameters):
+    threads: Optional[int] = None
+    models: Union[None, str, List[str]] = None
+    exclude: Union[None, str, List[str]] = None
+    selector: Optional[str] = None
+    state: Optional[str] = None
+    defer: Optional[bool] = None
 
 
 @dataclass
@@ -53,12 +64,14 @@ class RPCSnapshotParameters(RPCParameters):
     select: Union[None, str, List[str]] = None
     exclude: Union[None, str, List[str]] = None
     selector: Optional[str] = None
+    state: Optional[str] = None
 
 
 @dataclass
 class RPCTestParameters(RPCCompileParameters):
     data: bool = False
     schema: bool = False
+    state: Optional[str] = None
 
 
 @dataclass
@@ -68,11 +81,13 @@ class RPCSeedParameters(RPCParameters):
     exclude: Union[None, str, List[str]] = None
     selector: Optional[str] = None
     show: bool = False
+    state: Optional[str] = None
 
 
 @dataclass
 class RPCDocsGenerateParameters(RPCParameters):
     compile: bool = True
+    state: Optional[str] = None
 
 
 @dataclass
diff --git a/core/dbt/task/rpc/project_commands.py b/core/dbt/task/rpc/project_commands.py
--- a/core/dbt/task/rpc/project_commands.py
+++ b/core/dbt/task/rpc/project_commands.py
@@ -1,12 +1,15 @@
 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,
@@ -54,6 +57,15 @@ 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
 ):
@@ -66,16 +78,28 @@ def set_args(self, params: RPCCompileParameters) -> None:
         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[RPCCompileParameters], RunTask):
+
+class RemoteRunProjectTask(RPCCommandTask[RPCRunParameters], RunTask):
     METHOD_NAME = 'run'
 
-    def set_args(self, params: RPCCompileParameters) -> None:
+    def set_args(self, params: RPCRunParameters) -> None:
         self.args.models = self._listify(params.models)
         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):
@@ -90,6 +114,9 @@ def set_args(self, params: RPCSeedParameters) -> 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'
@@ -103,6 +130,9 @@ def set_args(self, params: RPCTestParameters) -> None:
         if params.threads is not None:
             self.args.threads = params.threads
 
+        self.args.state = state_path(params.state)
+        self.set_previous_state()
+
 
 class RemoteDocsGenerateProjectTask(
     RPCCommandTask[RPCDocsGenerateParameters],
@@ -116,6 +146,8 @@ def set_args(self, params: RPCDocsGenerateParameters) -> 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:
@@ -185,6 +217,9 @@ def set_args(self, params: RPCSnapshotParameters) -> None:
         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],
diff --git a/core/dbt/task/runnable.py b/core/dbt/task/runnable.py
--- a/core/dbt/task/runnable.py
+++ b/core/dbt/task/runnable.py
@@ -91,6 +91,9 @@ def __init__(self, args, config):
         self._skipped_children = {}
         self._raise_next_tick = None
         self.previous_state: Optional[PreviousState] = None
+        self.set_previous_state()
+
+    def set_previous_state(self):
         if self.args.state is not None:
             self.previous_state = PreviousState(self.args.state)
 


Test patch:
diff --git a/test/rpc/conftest.py b/test/rpc/conftest.py
--- a/test/rpc/conftest.py
+++ b/test/rpc/conftest.py
@@ -75,6 +75,16 @@ def postgres_profile_data(unique_schema):
                     'dbname': 'dbt',
                     'schema': unique_schema,
                 },
+                'other_schema': {
+                    'type': 'postgres',
+                    'threads': 4,
+                    'host': 'database',
+                    'port': 5432,
+                    'user': 'root',
+                    'pass': 'password',
+                    'dbname': 'dbt',
+                    'schema': unique_schema+'_alt',
+                }
             },
             'target': 'default'
         }
diff --git a/test/rpc/test_compile.py b/test/rpc/test_compile.py
--- a/test/rpc/test_compile.py
+++ b/test/rpc/test_compile.py
@@ -1,7 +1,9 @@
+import os
 import pytest
 from .util import (
     assert_has_threads,
     get_querier,
+    get_write_manifest,
     ProjectDefinition,
 )
 
@@ -28,3 +30,45 @@ def test_rpc_compile_threads(
             querier.cli_args('compile --threads=7')
         )
         assert_has_threads(results, 7)
+
+
+@pytest.mark.supported('postgres')
+def test_rpc_compile_state(
+    project_root, profiles_root, dbt_profile, unique_schema
+):
+    project = ProjectDefinition(models={'my_model.sql': 'select 1 as id'})
+    querier_ctx = get_querier(
+        project_def=project,
+        project_dir=project_root,
+        profiles_dir=profiles_root,
+        schema=unique_schema,
+        test_kwargs={},
+    )
+    with querier_ctx as querier:
+        state_dir = os.path.join(project_root, 'state')
+        os.makedirs(state_dir)
+
+        results = querier.async_wait_for_result(
+            querier.compile()
+        )
+        assert len(results['results']) == 1
+
+        get_write_manifest(querier, os.path.join(state_dir, 'manifest.json'))
+
+        project.models['my_model.sql'] = 'select 2 as id'
+        project.write_models(project_root, remove=True)
+
+        querier.sighup()
+        assert querier.wait_for_status('ready') is True
+
+        results = querier.async_wait_for_result(
+            querier.compile(state='./state', models=['state:modified'])
+        )
+        assert len(results['results']) == 1
+
+        get_write_manifest(querier, os.path.join(state_dir, 'manifest.json'))
+
+        results = querier.async_wait_for_result(
+            querier.compile(state='./state', models=['state:modified']),
+        )
+        assert len(results['results']) == 0
diff --git a/test/rpc/test_run.py b/test/rpc/test_run.py
--- a/test/rpc/test_run.py
+++ b/test/rpc/test_run.py
@@ -1,7 +1,9 @@
+import os
 import pytest
 from .util import (
     assert_has_threads,
     get_querier,
+    get_write_manifest,
     ProjectDefinition,
 )
 
@@ -79,3 +81,56 @@ def test_rpc_run_vars_compiled(
         results = querier.async_wait_for_result(querier.cli_args('run'))
         assert len(results['results']) == 1
         assert results['results'][0]['node']['config']['materialized'] == 'view'
+
+
+@pytest.mark.supported('postgres')
+def test_rpc_run_state_defer(
+    project_root, profiles_root, dbt_profile, unique_schema
+):
+    project = ProjectDefinition(models={'my_model.sql': 'select 1 as id'})
+    querier_ctx = get_querier(
+        project_def=project,
+        project_dir=project_root,
+        profiles_dir=profiles_root,
+        schema=unique_schema,
+        test_kwargs={},
+    )
+    with querier_ctx as querier:
+        state_dir = os.path.join(project_root, 'state')
+        os.makedirs(state_dir)
+
+        results = querier.async_wait_for_result(
+            querier.run()
+        )
+        assert len(results['results']) == 1
+
+        get_write_manifest(querier, os.path.join(state_dir, 'manifest.json'))
+
+        project.models['my_model.sql'] = 'select 2 as id'
+        project.write_models(project_root, remove=True)
+        querier.sighup()
+        assert querier.wait_for_status('ready') is True
+
+        results = querier.async_wait_for_result(
+            querier.run(state='./state', models=['state:modified'])
+        )
+        assert len(results['results']) == 1
+
+        get_write_manifest(querier, os.path.join(state_dir, 'manifest.json'))
+
+        results = querier.async_wait_for_result(
+            querier.run(state='./state', models=['state:modified']),
+        )
+        assert len(results['results']) == 0
+
+        project.models['my_model.sql'] = '{% if execute %}{% do exceptions.raise_compiler_error("should not see this") %}{% endif %}select 2 as id'
+        project.models['my_second_model.sql'] = 'select * from {{ ref("my_model") }}'
+        project.write_models(project_root, remove=True)
+        querier.sighup()
+        assert querier.wait_for_status('ready') is True
+
+        # if 'defer' is ignored, this will fail
+        results = querier.async_wait_for_result(
+            querier.run(state='./state', models=['my_second_model'], defer=True)
+        )
+        assert len(results['results']) == 1
diff --git a/test/rpc/test_seed.py b/test/rpc/test_seed.py
--- a/test/rpc/test_seed.py
+++ b/test/rpc/test_seed.py
@@ -1,7 +1,9 @@
 import pytest
+import os
 from .util import (
     assert_has_threads,
     get_querier,
+    get_write_manifest,
     ProjectDefinition,
 )
 
@@ -65,3 +67,45 @@ def test_rpc_seed_include_exclude(
         assert len(results['results']) == 1
         results = querier.async_wait_for_result(querier.cli_args('seed --exclude=data_2'))
         assert len(results['results']) == 1
+
+
+@pytest.mark.supported('postgres')
+def test_rpc_seed_state(
+    project_root, profiles_root, dbt_profile, unique_schema
+):
+    project = ProjectDefinition(seeds={'my_seed.csv': 'a,b\n1,hello'})
+    querier_ctx = get_querier(
+        project_def=project,
+        project_dir=project_root,
+        profiles_dir=profiles_root,
+        schema=unique_schema,
+        test_kwargs={},
+    )
+    with querier_ctx as querier:
+        state_dir = os.path.join(project_root, 'state')
+        os.makedirs(state_dir)
+
+        results = querier.async_wait_for_result(
+            querier.seed()
+        )
+        assert len(results['results']) == 1
+
+        get_write_manifest(querier, os.path.join(state_dir, 'manifest.json'))
+
+        project.seeds['my_seed.csv'] = 'a,b\n1,hello\n2,goodbye'
+        project.write_seeds(project_root, remove=True)
+
+        querier.sighup()
+        assert querier.wait_for_status('ready') is True
+
+        results = querier.async_wait_for_result(
+            querier.seed(state='./state', select=['state:modified'])
+        )
+        assert len(results['results']) == 1
+
+        get_write_manifest(querier, os.path.join(state_dir, 'manifest.json'))
+
+        results = querier.async_wait_for_result(
+            querier.seed(state='./state', select=['state:modified']),
+        )
+        assert len(results['results']) == 0
diff --git a/test/rpc/test_snapshots.py b/test/rpc/test_snapshots.py
--- a/test/rpc/test_snapshots.py
+++ b/test/rpc/test_snapshots.py
@@ -1,7 +1,9 @@
+import os
 import pytest
 from .util import (
     assert_has_threads,
     get_querier,
+    get_write_manifest,
     ProjectDefinition,
 )
 
@@ -22,6 +24,23 @@
 {% endsnapshot %}
 '''
 
+snapshot_data_2 = '''
+{% snapshot snapshot_actual %}
+
+    {{
+        config(
+            target_database=database,
+            target_schema=schema,
+            unique_key='id',
+            strategy='timestamp',
+            updated_at='updated_at',
+        )
+    }}
+    select 2 as id, '2019-10-31 23:59:40' as updated_at
+
+{% endsnapshot %}
+'''
+
 
 @pytest.mark.supported('postgres')
 def test_snapshots(
@@ -108,3 +127,46 @@ def test_rpc_snapshot_threads(
         )
         assert_has_threads(results, 7)
 
+
+@pytest.mark.supported('postgres')
+def test_rpc_snapshot_state(
+    project_root, profiles_root, dbt_profile, unique_schema
+):
+    project = ProjectDefinition(
+        snapshots={'my_snapshots.sql': snapshot_data},
+    )
+    querier_ctx = get_querier(
+        project_def=project,
+        project_dir=project_root,
+        profiles_dir=profiles_root,
+        schema=unique_schema,
+        test_kwargs={},
+    )
+    with querier_ctx as querier:
+        state_dir = os.path.join(project_root, 'state')
+        os.makedirs(state_dir)
+
+        results = querier.async_wait_for_result(
+            querier.snapshot()
+        )
+        assert len(results['results']) == 1
+
+        get_write_manifest(querier, os.path.join(state_dir, 'manifest.json'))
+
+        project.snapshots['my_snapshots.sql'] = snapshot_data_2
+        project.write_snapshots(project_root, remove=True)
+
+        querier.sighup()
+        assert querier.wait_for_status('ready') is True
+
+        results = querier.async_wait_for_result(
+            querier.snapshot(state='./state', select=['state:modified'])
+        )
+        assert len(results['results']) == 1
+
+        get_write_manifest(querier, os.path.join(state_dir, 'manifest.json'))
+
+        results = querier.async_wait_for_result(
+            querier.snapshot(state='./state', select=['state:modified']),
+        )
+        assert len(results['results']) == 0
diff --git a/test/rpc/test_test.py b/test/rpc/test_test.py
--- a/test/rpc/test_test.py
+++ b/test/rpc/test_test.py
@@ -1,8 +1,10 @@
+import os
 import pytest
 import yaml
 from .util import (
     assert_has_threads,
     get_querier,
+    get_write_manifest,
     ProjectDefinition,
 )
 
@@ -46,3 +48,60 @@ def test_rpc_test_threads(
             querier.cli_args('test --threads=7')
         )
         assert_has_threads(results, 7)
+
+
+@pytest.mark.supported('postgres')
+def test_rpc_test_state(
+    project_root, profiles_root, dbt_profile, unique_schema
+):
+    schema_yaml = {
+        'version': 2,
+        'models': [{
+            'name': 'my_model',
+            'columns': [
+                {
+                    'name': 'id',
+                    'tests': ['not_null', 'unique'],
+                },
+            ],
+        }],
+    }
+    project = ProjectDefinition(
+        models={
+            'my_model.sql': 'select 1 as id',
+            'schema.yml': yaml.safe_dump(schema_yaml)
+        }
+    )
+    querier_ctx = get_querier(
+        project_def=project,
+        project_dir=project_root,
+        profiles_dir=profiles_root,
+        schema=unique_schema,
+        test_kwargs={},
+    )
+    with querier_ctx as querier:
+        state_dir = os.path.join(project_root, 'state')
+        os.makedirs(state_dir)
+
+        results = querier.async_wait_for_result(querier.run())
+        results = querier.async_wait_for_result(querier.test())
+        assert len(results['results']) == 2
+
+        get_write_manifest(querier, os.path.join(state_dir, 'manifest.json'))
+
+        project.models['my_model.sql'] = 'select 2 as id'
+        project.write_models(project_root, remove=True)
+        querier.sighup()
+        assert querier.wait_for_status('ready') is True
+
+        results = querier.async_wait_for_result(
+            querier.test(state='./state', models=['state:modified'])
+        )
+        assert len(results['results']) == 2
+
+        get_write_manifest(querier, os.path.join(state_dir, 'manifest.json'))
+
+        results = querier.async_wait_for_result(
+            querier.test(state='./state', models=['state:modified']),
+        )
+        assert len(results['results']) == 0
diff --git a/test/rpc/util.py b/test/rpc/util.py
--- a/test/rpc/util.py
+++ b/test/rpc/util.py
@@ -215,6 +215,7 @@ def compile(
         exclude: Optional[Union[str, List[str]]] = None,
         threads: Optional[int] = None,
         request_id: int = 1,
+        state: Optional[bool] = None,
     ):
         params = {}
         if models is not None:
@@ -223,6 +224,8 @@ def compile(
             params['exclude'] = exclude
         if threads is not None:
             params['threads'] = threads
+        if state is not None:
+            params['state'] = state
         return self.request(
             method='compile', params=params, request_id=request_id
         )
@@ -233,6 +236,8 @@ def run(
         exclude: Optional[Union[str, List[str]]] = None,
         threads: Optional[int] = None,
         request_id: int = 1,
+        defer: Optional[bool] = None,
+        state: Optional[str] = None,
     ):
         params = {}
         if models is not None:
@@ -241,6 +246,10 @@ def run(
             params['exclude'] = exclude
         if threads is not None:
             params['threads'] = threads
+        if defer is not None:
+            params['defer'] = defer
+        if state is not None:
+            params['state'] = state
         return self.request(
             method='run', params=params, request_id=request_id
         )
@@ -265,6 +274,7 @@ def seed(
         show: bool = None,
         threads: Optional[int] = None,
         request_id: int = 1,
+        state: Optional[bool] = None,
     ):
         params = {}
         if select is not None:
@@ -275,6 +285,8 @@ def seed(
             params['show'] = show
         if threads is not None:
             params['threads'] = threads
+        if state is not None:
+            params['state'] = state
         return self.request(
             method='seed', params=params, request_id=request_id
         )
@@ -285,6 +297,7 @@ def snapshot(
         exclude: Optional[Union[str, List[str]]] = None,
         threads: Optional[int] = None,
         request_id: int = 1,
+        state: Optional[bool] = None,
     ):
         params = {}
         if select is not None:
@@ -293,6 +306,8 @@ def snapshot(
             params['exclude'] = exclude
         if threads is not None:
             params['threads'] = threads
+        if state is not None:
+            params['state'] = state
         return self.request(
             method='snapshot', params=params, request_id=request_id
         )
@@ -320,6 +335,7 @@ def test(
         data: bool = None,
         schema: bool = None,
         request_id: int = 1,
+        state: Optional[bool] = None,
     ):
         params = {}
         if models is not None:
@@ -332,6 +348,8 @@ def test(
             params['schema'] = schema
         if threads is not None:
             params['threads'] = threads
+        if state is not None:
+            params['state'] = state
         return self.request(
             method='test', params=params, request_id=request_id
         )
@@ -638,3 +656,11 @@ def assert_has_threads(results, num_threads):
         f'Got invalid number of concurrency logs ({len(c_logs)})'
     assert 'message' in c_logs[0]
     assert f'Concurrency: {num_threads} threads' in c_logs[0]['message']
+
+
+def get_write_manifest(querier, dest):
+    result = querier.async_wait_for_result(querier.get_manifest())
+    manifest = result['manifest']
+
+    with open(dest, 'w') as fp:
+        json.dump(manifest, fp)