Skip to content

Commit

Permalink
Fix lint.
Browse files Browse the repository at this point in the history
  • Loading branch information
jmcarp committed Mar 4, 2021
1 parent 62be9f9 commit 7b0c74c
Showing 1 changed file with 13 additions and 9 deletions.
22 changes: 13 additions & 9 deletions plugins/bigquery/dbt/adapters/bigquery/connections.py
Original file line number Diff line number Diff line change
Expand Up @@ -307,21 +307,25 @@ def raw_execute(self, sql, fetch=False, *, use_legacy_sql=False):

logger.debug('On {}: {}', conn.name, sql)

job_params = {'use_legacy_sql': use_legacy_sql, 'labels': {}}
labels = {}

if active_user:
job_params['labels']['dbt_invocation_id'] = active_user.invocation_id
labels['dbt_invocation_id'] = active_user.invocation_id

if self.profile.query_comment.job_label:
try:
labels = json.loads(self.query_header.comment.query_comment)
job_params['labels'].update({
_sanitize_bigquery_label(key): _sanitize_bigquery_label(str(value))
for key, value in labels.items()
comment_labels = json.loads(
self.query_header.comment.query_comment
)
labels.update({
_sanitize_label(key): _sanitize_label(str(value))
for key, value in comment_labels.items()
})
except (TypeError, ValueError):
pass

job_params = {'use_legacy_sql': use_legacy_sql, 'labels': labels}

priority = conn.credentials.priority
if priority == Priority.Batch:
job_params['priority'] = google.cloud.bigquery.QueryPriority.BATCH
Expand Down Expand Up @@ -585,12 +589,12 @@ def _is_retryable(error):
return False


_SANITIZE_BIGQUERY_LABEL_PATTERN = re.compile(r"[^a-z0-9_-]")
_SANITIZE_LABEL_PATTERN = re.compile(r"[^a-z0-9_-]")


def _sanitize_bigquery_label(value: str, max_length: int = 63) -> str:
def _sanitize_label(value: str, max_length: int = 63) -> str:
"""Return a legal value for a BigQuery label."""
value = value.lower()
value = _SANITIZE_BIGQUERY_LABEL_PATTERN.sub("_", value)
value = _SANITIZE_LABEL_PATTERN.sub("_", value)
value = value[: max_length - 1]
return value

0 comments on commit 7b0c74c

Please sign in to comment.