Skip to content

Commit

Permalink
add unit test and move default logic to mashumaro hook
Browse files Browse the repository at this point in the history
  • Loading branch information
Kyle Wigley committed Apr 28, 2021
1 parent 7bf2fd6 commit 4abcd18
Show file tree
Hide file tree
Showing 2 changed files with 25 additions and 5 deletions.
9 changes: 5 additions & 4 deletions plugins/bigquery/dbt/adapters/bigquery/connections.py
Original file line number Diff line number Diff line change
Expand Up @@ -111,16 +111,17 @@ def type(self):
def _connection_keys(self):
return ('method', 'database', 'schema', 'location', 'priority',
'timeout_seconds', 'maximum_bytes_billed')

def __post_init__(self):
@classmethod
def __pre_deserialize__(cls, d: Dict[Any, Any]) -> Dict[Any, Any]:
# We need to inject the correct value of the database (aka project) at
# this stage, ref
# https:/fishtown-analytics/dbt/pull/2908#discussion_r532927436.

# `database` is an alias of `project` in BigQuery
if self.database is None:
if 'database' not in d:
_, database = get_bigquery_defaults()
self.database = database
d['database'] = database
return d


class BigQueryConnectionManager(BaseConnectionManager):
Expand Down
21 changes: 20 additions & 1 deletion test/unit/test_bigquery_adapter.py
Original file line number Diff line number Diff line change
Expand Up @@ -97,7 +97,7 @@ def setUp(self):
'priority': 'batch',
'maximum_bytes_billed': 0,
},
'oauth--no-project': {
'oauth-no-project': {
'type': 'bigquery',
'method': 'oauth',
'schema': 'dummy_schema',
Expand Down Expand Up @@ -144,6 +144,25 @@ def get_adapter(self, target):


class TestBigQueryAdapterAcquire(BaseTestBigQueryAdapter):
@patch('dbt.adapters.bigquery.connections.get_bigquery_defaults', return_value=('credentials', 'project_id'))
@patch('dbt.adapters.bigquery.BigQueryConnectionManager.open', return_value=_bq_conn())
def test_acquire_connection_oauth_no_project_validations(self, mock_open_connection, mock_get_bigquery_defaults):
adapter = self.get_adapter('oauth-no-project')
mock_get_bigquery_defaults.assert_called_once()
try:
connection = adapter.acquire_connection('dummy')
self.assertEqual(connection.type, 'bigquery')

except dbt.exceptions.ValidationException as e:
self.fail('got ValidationException: {}'.format(str(e)))

except BaseException as e:
raise

mock_open_connection.assert_not_called()
connection.handle
mock_open_connection.assert_called_once()

@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')
Expand Down

0 comments on commit 4abcd18

Please sign in to comment.