From e29eccd741250d157aceb48f41f0dfd38a6bac69 Mon Sep 17 00:00:00 2001 From: Jacob Beck Date: Wed, 3 Apr 2019 19:13:22 -0600 Subject: [PATCH 1/2] pass the model name along in get_relations --- plugins/bigquery/dbt/adapters/bigquery/impl.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/plugins/bigquery/dbt/adapters/bigquery/impl.py b/plugins/bigquery/dbt/adapters/bigquery/impl.py index ad5fbf724bf..c36337f45c9 100644 --- a/plugins/bigquery/dbt/adapters/bigquery/impl.py +++ b/plugins/bigquery/dbt/adapters/bigquery/impl.py @@ -156,7 +156,8 @@ def get_relation(self, database, schema, identifier, model_name=None): ) try: - table = self.connections.get_bq_table(database, schema, identifier) + table = self.connections.get_bq_table(database, schema, identifier, + conn_name=model_name) except google.api_core.exceptions.NotFound: table = None return self._bq_table_to_relation(table) From c19644882bdec827021eeb3a5f18e6b5e152b845 Mon Sep 17 00:00:00 2001 From: Jacob Beck Date: Fri, 5 Apr 2019 10:22:54 -0600 Subject: [PATCH 2/2] add unit tests --- test/unit/test_bigquery_adapter.py | 49 +++++++++++++++++++++++++++++- 1 file changed, 48 insertions(+), 1 deletion(-) diff --git a/test/unit/test_bigquery_adapter.py b/test/unit/test_bigquery_adapter.py index b667cf38a6f..971c6b18968 100644 --- a/test/unit/test_bigquery_adapter.py +++ b/test/unit/test_bigquery_adapter.py @@ -18,7 +18,7 @@ def _bq_conn(): return conn -class TestBigQueryAdapter(unittest.TestCase): +class BaseTestBigQueryAdapter(unittest.TestCase): def setUp(self): flags.STRICT_MODE = True @@ -73,6 +73,7 @@ def get_adapter(self, target): return adapter +class TestBigQueryAdapterAcquire(BaseTestBigQueryAdapter): @patch('dbt.adapters.bigquery.BigQueryConnectionManager.open', return_value=_bq_conn()) def test_acquire_connection_oauth_validations(self, mock_open_connection): adapter = self.get_adapter('oauth') @@ -134,6 +135,52 @@ def test_location_value(self, mock_bq, mock_auth_default): location='Luna Station') +class TestConnectionNamePassthrough(BaseTestBigQueryAdapter): + + def setUp(self): + super(TestConnectionNamePassthrough, self).setUp() + self._conn_patch = patch.object(BigQueryAdapter, 'ConnectionManager') + self.conn_manager_cls = self._conn_patch.start() + + self._relation_patch = patch.object(BigQueryAdapter, 'Relation') + self.relation_cls = self._relation_patch.start() + + self.mock_connection_manager = self.conn_manager_cls.return_value + self.conn_manager_cls.TYPE = 'bigquery' + self.relation_cls.DEFAULTS = BigQueryRelation.DEFAULTS + + self.adapter = self.get_adapter('oauth') + + def tearDown(self): + super(TestConnectionNamePassthrough, self).tearDown() + self._conn_patch.stop() + self._relation_patch.stop() + + def test_get_relation(self): + self.adapter.get_relation('db', 'schema', 'my_model', 'my_conn') + self.mock_connection_manager.get_bq_table.assert_called_once_with('db', 'schema', 'my_model', conn_name='my_conn') + + def test_create_schema(self): + self.adapter.create_schema('db', 'schema', 'my_conn') + self.mock_connection_manager.create_dataset.assert_called_once_with('db', 'schema', 'my_conn') + + @patch.object(BigQueryAdapter, 'check_schema_exists') + def test_drop_schema(self, mock_check_schema): + mock_check_schema.return_value = True + self.adapter.drop_schema('db', 'schema', 'my_conn') + self.mock_connection_manager.drop_dataset.assert_called_once_with('db', 'schema', 'my_conn') + + def test_get_columns_in_relation(self): + self.mock_connection_manager.get_bq_table.side_effect = ValueError + self.adapter.get_columns_in_relation( + MagicMock(database='db', schema='schema', table_name='ident'), + 'my_conn' + ) + self.mock_connection_manager.get_bq_table.assert_called_once_with( + database='db', schema='schema', identifier='ident', conn_name='my_conn' + ) + + class TestBigQueryRelation(unittest.TestCase): def setUp(self): flags.STRICT_MODE = True