Merge pull request #1151 from fishtown-analytics/feature/bq-profile-l…

…ocation Add 'location' to google bigquery profile (#969)
dbt-labs · Nov 26, 2018 · d16ca86 · d16ca86
2 parents f72e0a8 + b92d669
commit d16ca86
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 7 deletions.
diff --git a/dbt/adapters/bigquery/impl.py b/dbt/adapters/bigquery/impl.py
@@ -142,8 +142,9 @@ def get_bigquery_credentials(cls, profile_credentials):
     def get_bigquery_client(cls, profile_credentials):
         project_name = profile_credentials.project
         creds = cls.get_bigquery_credentials(profile_credentials)
-
-        return google.cloud.bigquery.Client(project_name, creds)
+        location = getattr(profile_credentials, 'location', None)
+        return google.cloud.bigquery.Client(project_name, creds,
+                                            location=location)
 
     @classmethod
     def open_connection(cls, connection):

diff --git a/dbt/contracts/connection.py b/dbt/contracts/connection.py
@@ -139,6 +139,9 @@
         'timeout_seconds': {
             'type': 'integer',
         },
+        'location': {
+            'type': 'string',
+        },
     },
     'required': ['method', 'project', 'schema'],
 }

diff --git a/test/unit/test_bigquery_adapter.py b/test/unit/test_bigquery_adapter.py
@@ -35,6 +35,14 @@ def setUp(self):
                     'keyfile': '/tmp/dummy-service-account.json',
                     'threads': 1,
                 },
+                'loc': {
+                    'type': 'bigquery',
+                    'method': 'oauth',
+                    'project': 'dbt-unit-000000',
+                    'schema': 'dummy_schema',
+                    'threads': 1,
+                    'location': 'Luna Station',
+                },
             },
             'target': 'oauth',
         }
@@ -43,15 +51,17 @@ def setUp(self):
             'name': 'X',
             'version': '0.1',
             'project-root': '/tmp/dbt/does-not-exist',
+            'profile': 'default',
         }
 
-    def get_adapter(self, profile):
+    def get_adapter(self, target):
         project = self.project_cfg.copy()
-        project['profile'] = profile
+        profile = self.raw_profile.copy()
+        profile['target'] = target
 
         config = config_from_parts_or_dicts(
             project=project,
-            profile=self.raw_profile,
+            profile=profile,
         )
         return BigQueryAdapter(config)
 
@@ -68,7 +78,6 @@ def test_acquire_connection_oauth_validations(self, mock_open_connection):
 
         except BaseException as e:
             raise
-            self.fail('validation failed with unknown exception: {}'.format(str(e)))
 
         mock_open_connection.assert_called_once()
 
@@ -84,10 +93,21 @@ def test_acquire_connection_service_account_validations(self, mock_open_connecti
 
         except BaseException as e:
             raise
-            self.fail('validation failed with unknown exception: {}'.format(str(e)))
 
         mock_open_connection.assert_called_once()
 
+    @patch('dbt.adapters.bigquery.impl.google.auth.default')
+    @patch('dbt.adapters.bigquery.impl.google.cloud.bigquery')
+    def test_location_value(self, mock_bq, mock_auth_default):
+        creds = MagicMock()
+        mock_auth_default.return_value = (creds, MagicMock())
+        adapter = self.get_adapter('loc')
+
+        adapter.acquire_connection('dummy')
+        mock_client = mock_bq.Client
+        mock_client.assert_called_once_with('dbt-unit-000000', creds,
+                                            location='Luna Station')
+
 
 class TestBigQueryRelation(unittest.TestCase):
     def setUp(self):