Add source snapshot-freshness to dbt rpc (#2040)

core/dbt/contracts/rpc.py

@@ -136,6 +136,12 @@ class RPCRunOperationParameters(RPCParameters):
     args: Dict[str, Any] = field(default_factory=dict)
 
 
+@dataclass
+class RPCSourceFreshnessParameters(RPCParameters):
+    threads: Optional[int] = None
+    select: Union[None, str, List[str]] = None
+
+
 # Outputs
 
 @dataclass

core/dbt/main.py

@@ -594,8 +594,11 @@ def _build_source_snapshot_freshness_subparser(subparsers, base_subparser):
         Specify number of threads to use. Overrides settings in profiles.yml
         '''
     )
-    sub.set_defaults(cls=freshness_task.FreshnessTask,
-                     which='snapshot-freshness', rpc_method=None)
+    sub.set_defaults(
+        cls=freshness_task.FreshnessTask,
+        which='snapshot-freshness',
+        rpc_method='snapshot-freshness',
+    )
     return sub
 
 

core/dbt/task/rpc/project_commands.py

@@ -12,11 +12,13 @@
     RemoteExecutionResult,
     RemoteRunOperationResult,
     RPCSnapshotParameters,
+    RPCSourceFreshnessParameters,
 )
 from dbt.rpc.method import (
     Parameters,
 )
 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
@@ -150,3 +152,16 @@ def set_args(self, params: RPCSnapshotParameters) -> None:
         self.args.exclude = self._listify(params.exclude)
         if params.threads is not None:
             self.args.threads = params.threads
+
+
+class RemoteSourceFreshnessTask(
+    RPCCommandTask[RPCSourceFreshnessParameters],
+    FreshnessTask
+):
+    METHOD_NAME = 'snapshot-freshness'
+
+    def set_args(self, params: RPCSourceFreshnessParameters) -> None:
+        self.args.selected = self._listify(params.select)
+        if params.threads is not None:
+            self.args.threads = params.threads
+        self.args.output = None

test/rpc/test_base.py

@@ -1,5 +1,5 @@
 # flake8: disable=redefined-outer-name
-
+from datetime import datetime, timedelta
 import time
 import yaml
 from .util import (
@@ -712,3 +712,81 @@ def test_get_status(
         result = querier.is_result(querier.poll(token, logs=False))
         assert 'logs' in result
         assert len(result['logs']) == 0
+
+
+source_freshness_schema_yml = '''
+version: 2
+sources:
+  - name: test_source
+    loaded_at_field: b
+    schema: {schema}
+    freshness:
+      warn_after: {{count: 10, period: hour}}
+      error_after: {{count: 1, period: day}}
+    tables:
+      - name: test_table
+        identifier: source
+      - name: failure_table
+        identifier: other_source
+'''
+
+
+def test_source_freshness(
+    project_root, profiles_root, postgres_profile, unique_schema
+):
+    start_time = datetime.utcnow()
+    warn_me = start_time - timedelta(hours=18)
+    error_me = start_time - timedelta(days=2)
+    # this should trigger a 'warn'
+    project = ProjectDefinition(
+        project_data={'seeds': {'quote_columns': False}},
+        seeds={
+            'source.csv': 'a,b\n1,{}\n'.format(error_me.strftime('%Y-%m-%d %H:%M:%S')),
+            'other_source.csv': 'a,b\n1,{}\n'.format(error_me.strftime('%Y-%m-%d %H:%M:%S'))
+        },
+        models={
+            'sources.yml': source_freshness_schema_yml.format(schema=unique_schema),
+        },
+    )
+    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:
+        seeds = querier.async_wait_for_result(querier.seed())
+        assert len(seeds['results']) == 2
+        # should error
+        error_results = querier.async_wait_for_result(querier.snapshot_freshness(), state='failed')
+        assert len(error_results['results']) == 2
+        for result in error_results['results']:
+            assert result['status'] == 'error'
+        error_results = querier.async_wait_for_result(querier.cli_args('source snapshot-freshness'), state='failed')
+        assert len(error_results['results']) == 2
+        for result in error_results['results']:
+            assert result['status'] == 'error'
+
+        project.seeds['source.csv'] += '2,{}\n'.format(warn_me.strftime('%Y-%m-%d %H:%M:%S'))
+        project.write_seeds(project_root, remove=True)
+        querier.async_wait_for_result(querier.seed())
+        # should warn
+        warn_results = querier.async_wait_for_result(querier.snapshot_freshness(select='test_source.test_table'))
+        assert len(warn_results['results']) == 1
+        assert warn_results['results'][0]['status'] == 'warn'
+        warn_results = querier.async_wait_for_result(querier.cli_args('source snapshot-freshness -s test_source.test_table'))
+        assert len(warn_results['results']) == 1
+        assert warn_results['results'][0]['status'] == 'warn'
+
+        project.seeds['source.csv'] += '3,{}\n'.format(start_time.strftime('%Y-%m-%d %H:%M:%S'))
+        project.write_seeds(project_root, remove=True)
+        querier.async_wait_for_result(querier.seed())
+        # should pass!
+        pass_results = querier.async_wait_for_result(querier.snapshot_freshness(select=['test_source.test_table']))
+        assert len(pass_results['results']) == 1
+        assert pass_results['results'][0]['status'] == 'pass'
+        pass_results = querier.async_wait_for_result(querier.cli_args('source snapshot-freshness --select test_source.test_table'))
+        assert len(pass_results['results']) == 1
+        assert pass_results['results'][0]['status'] == 'pass'

test/rpc/util.py

@@ -271,6 +271,21 @@ def snapshot(
             method='snapshot', params=params, request_id=request_id
         )
 
+    def snapshot_freshness(
+        self,
+        select: Optional[Union[str, List[str]]] = None,
+        threads: Optional[int] = None,
+        request_id: int = 1,
+    ):
+        params = {}
+        if select is not None:
+            params['select'] = select
+        if threads is not None:
+            params['threads'] = threads
+        return self.request(
+            method='snapshot-freshness', params=params, request_id=request_id
+        )
+
     def test(
         self,
         models: Optional[Union[str, List[str]]] = None,