Skip to content

Commit

Permalink
fix validate connection failing (dbt-labs#168)
Browse files Browse the repository at this point in the history
* fix validate connection failing

* fix code check

* more format
  • Loading branch information
ChenyuLInx authored Apr 27, 2022
1 parent 5b8281c commit ebeeb88
Show file tree
Hide file tree
Showing 2 changed files with 21 additions and 5 deletions.
22 changes: 17 additions & 5 deletions dbt/adapters/bigquery/connections.py
Original file line number Diff line number Diff line change
Expand Up @@ -380,8 +380,11 @@ def raw_execute(self, sql, fetch=False, *, use_legacy_sql=False):
client = conn.handle

fire_event(SQLQuery(conn_name=conn.name, sql=sql))

if self.profile.query_comment and self.profile.query_comment.job_label:
if (
hasattr(self.profile, "query_comment")
and self.profile.query_comment
and self.profile.query_comment.job_label
):
query_comment = self.query_header.comment.query_comment
labels = self._labels_from_query_comment(query_comment)
else:
Expand Down Expand Up @@ -475,7 +478,10 @@ def execute(
message = f"{code} ({num_rows_formated} rows, {processed_bytes} processed)"

response = BigQueryAdapterResponse( # type: ignore[call-arg]
_message=message, rows_affected=num_rows, code=code, bytes_processed=bytes_processed
_message=message,
rows_affected=num_rows,
code=code,
bytes_processed=bytes_processed,
)

return response, table
Expand Down Expand Up @@ -530,7 +536,8 @@ def copy_and_results():

self._retry_and_handle(
msg='copy table "{}" to "{}"'.format(
", ".join(source_ref.path for source_ref in source_ref_array), destination_ref.path
", ".join(source_ref.path for source_ref in source_ref_array),
destination_ref.path,
),
conn=conn,
fn=copy_and_results,
Expand Down Expand Up @@ -572,7 +579,12 @@ def fn():
self._retry_and_handle(msg="create dataset", conn=conn, fn=fn)

def _query_and_results(
self, client, sql, job_params, job_creation_timeout=None, job_execution_timeout=None
self,
client,
sql,
job_params,
job_creation_timeout=None,
job_execution_timeout=None,
):
"""Query the client and wait for results."""
# Cannot reuse job_config if destination is set and ddl is used
Expand Down
4 changes: 4 additions & 0 deletions tests/functional/adapter/test_basic.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
from dbt.tests.adapter.basic.test_snapshot_check_cols import BaseSnapshotCheckCols
from dbt.tests.adapter.basic.test_snapshot_timestamp import BaseSnapshotTimestamp
from dbt.tests.adapter.basic.test_adapter_methods import BaseAdapterMethod
from dbt.tests.adapter.basic.test_validate_connection import BaseValidateConnection


class TestSimpleMaterializationsBigQuery(BaseSimpleMaterializations):
Expand Down Expand Up @@ -54,3 +55,6 @@ class TestSnapshotTimestampBigQuery(BaseSnapshotTimestamp):

class TestBaseAdapterMethodBigQuery(BaseAdapterMethod):
pass

class TestBigQueryValidateConnection(BaseValidateConnection):
pass

0 comments on commit ebeeb88

Please sign in to comment.