Merge pull request #2042 from fishtown-analytics/feature/select-on-seeds

support --select on dbt seed (#1711)

core/dbt/contracts/rpc.py

@@ -62,6 +62,8 @@ class RPCTestParameters(RPCCompileParameters):
 @dataclass
 class RPCSeedParameters(RPCParameters):
     threads: Optional[int] = None
+    select: Union[None, str, List[str]] = None
+    exclude: Union[None, str, List[str]] = None
     show: bool = False
 
 

core/dbt/main.py

@@ -863,7 +863,7 @@ def parse_args(args, cls=DBTArgumentParser):
                           rpc_sub, seed_sub)
     # --models, --exclude
     _add_selection_arguments(run_sub, compile_sub, generate_sub, test_sub)
-    _add_selection_arguments(snapshot_sub, models_name='select')
+    _add_selection_arguments(snapshot_sub, seed_sub, models_name='select')
     # --full-refresh
     _add_table_mutability_arguments(run_sub, compile_sub)
 

core/dbt/task/rpc/project_commands.py

@@ -73,6 +73,9 @@ 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.models = self._listify(params.select)
+        self.args.exclude = self._listify(params.exclude)
         if params.threads is not None:
             self.args.threads = params.threads
         self.args.show = params.show

core/dbt/task/seed.py

@@ -13,8 +13,8 @@ def raise_on_first_error(self):
 
     def build_query(self):
         return {
-            "include": ["*"],
-            "exclude": [],
+            "include": self.args.models,
+            "exclude": self.args.exclude,
             "resource_types": [NodeType.Seed],
         }
 

test/integration/005_simple_seed_test/test_simple_seed.py

@@ -136,6 +136,22 @@ def test_postgres_simple_seed_with_disabled(self):
         self.assertTableDoesExist('seed_enabled')
         self.assertTableDoesNotExist('seed_disabled')
 
+    @use_profile('postgres')
+    def test_postgres_simple_seed_selection(self):
+        results = self.run_dbt(['seed', '--select', 'seed_enabled'])
+        self.assertEqual(len(results),  1)
+        self.assertTableDoesExist('seed_enabled')
+        self.assertTableDoesNotExist('seed_disabled')
+        self.assertTableDoesNotExist('seed_tricky')
+
+    @use_profile('postgres')
+    def test_postgres_simple_seed_exclude(self):
+        results = self.run_dbt(['seed', '--exclude', 'seed_enabled'])
+        self.assertEqual(len(results),  1)
+        self.assertTableDoesNotExist('seed_enabled')
+        self.assertTableDoesNotExist('seed_disabled')
+        self.assertTableDoesExist('seed_tricky')
+
 
 class TestSeedParsing(DBTIntegrationTest):
     def setUp(self):

test/rpc/test_base.py

@@ -622,6 +622,40 @@ def test_rpc_seed_threads(
         assert_has_threads(results, 7)
 
 
+def test_rpc_seed_include_exclude(
+    project_root, profiles_root, postgres_profile, unique_schema
+):
+    project = ProjectDefinition(
+        project_data={'seeds': {'quote_columns': False}},
+        seeds={
+            'data_1.csv': 'a,b\n1,hello\n2,goodbye',
+            'data_2.csv': 'a,b\n1,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:
+        results = querier.async_wait_for_result(querier.seed(select=['data_1']))
+        assert len(results['results']) == 1
+        results = querier.async_wait_for_result(querier.seed(select='data_1'))
+        assert len(results['results']) == 1
+        results = querier.async_wait_for_result(querier.cli_args('seed --select=data_1'))
+        assert len(results['results']) == 1
+
+        results = querier.async_wait_for_result(querier.seed(exclude=['data_2']))
+        assert len(results['results']) == 1
+        results = querier.async_wait_for_result(querier.seed(exclude='data_2'))
+        assert len(results['results']) == 1
+        results = querier.async_wait_for_result(querier.cli_args('seed --exclude=data_2'))
+        assert len(results['results']) == 1
+
+
 sleeper_sql = '''
 {{ log('test output', info=True) }}
 {{ run_query('select * from pg_sleep(20)') }}

test/rpc/util.py

@@ -240,11 +240,17 @@ def run_operation(
 
     def seed(
         self,
+        select: Optional[Union[str, List[str]]] = None,
+        exclude: Optional[Union[str, List[str]]] = None,
         show: bool = None,
         threads: Optional[int] = None,
         request_id: int = 1,
     ):
         params = {}
+        if select is not None:
+            params['select'] = select
+        if exclude is not None:
+            params['exclude'] = exclude
         if show is not None:
             params['show'] = show
         if threads is not None: