def __create_job_config(
self, ems_query_job_config: EmsQueryJobConfig) -> QueryJobConfig:
job_config = QueryJobConfig()
job_config.priority = ems_query_job_config.priority.value
job_config.use_legacy_sql = False
job_config.use_query_cache = ems_query_job_config.use_query_cache
job_config.labels = ems_query_job_config.labels
if ems_query_job_config.destination_table is not None:
job_config.time_partitioning = TimePartitioning("DAY")
table_reference = TableReference(
DatasetReference(
ems_query_job_config.destination_project_id
or self.__project_id,
ems_query_job_config.destination_dataset),
ems_query_job_config.destination_table)
job_config.destination = table_reference
job_config.write_disposition = ems_query_job_config.write_disposition.value
job_config.create_disposition = ems_query_job_config.create_disposition.value
partitioning = ems_query_job_config.time_partitioning
if partitioning is not None:
job_config.time_partitioning = TimePartitioning(
partitioning.type.value, partitioning.field,
partitioning.expiration_ms,
partitioning.require_partition_filter)
if ems_query_job_config.table_definitions is not None:
job_config.table_definitions = ems_query_job_config.table_definitions
return job_config
def load(project_id, bq_client, src_dataset_id, dst_dataset_id):
"""
Transform safely loaded tables and store results in target dataset.
:param project_id: Identifies the BQ project
:param bq_client: a BigQuery client object
:param src_dataset_id: reference to source dataset object
:param dst_dataset_id: reference to destination dataset object
:return: List of BQ job_ids
"""
dst_dataset = Dataset(f'{bq_client.project}.{dst_dataset_id}')
dst_dataset.description = f'Vocabulary cleaned and loaded from {src_dataset_id}'
dst_dataset.labels = {'type': 'vocabulary'}
dst_dataset.location = "US"
bq_client.create_dataset(dst_dataset, exists_ok=True)
src_tables = list(bq_client.list_tables(dataset=src_dataset_id))
job_config = QueryJobConfig()
query_jobs = []
for src_table in src_tables:
schema = bq.get_table_schema(src_table.table_id)
destination = f'{project_id}.{dst_dataset_id}.{src_table.table_id}'
table = bq_client.create_table(Table(destination, schema=schema),
exists_ok=True)
job_config.destination = table
query = SELECT_TPL.render(project_id=project_id,
dataset_id=src_dataset_id,
table=src_table.table_id,
fields=schema)
query_job = bq_client.query(query, job_config=job_config)
LOGGER.info(f'table:{destination} job_id:{query_job.job_id}')
query_jobs.append(query_job)
query_job.result()
return query_jobs
def schema_upgrade_cdm52_to_cdm531(project_id,
dataset_id,
snapshot_dataset_id,
hpo_id=None):
"""
:param project_id:
:param dataset_id: Dataset to convert
:param snapshot_dataset_id: Dataset with converted tables. Overwritten if tables already exist
:param hpo_id: Identifies the hpo_id of the site
:return:
"""
# Create dataset if not exists
client = bq.get_client(project_id)
client.create_dataset(snapshot_dataset_id, exists_ok=True)
sq.create_empty_cdm_tables(snapshot_dataset_id, hpo_id)
copy_table_job_ids = []
tables = [table.table_id for table in list(client.list_tables(dataset_id))]
if hpo_id:
hpo_tables = [
resources.get_table_id(table, hpo_id)
for table in resources.CDM_TABLES + PII_TABLES
]
# Filter tables that do not exist
tables = [table for table in hpo_tables if table in tables]
for table_id in tables:
q = get_upgrade_table_query(client, dataset_id, table_id, hpo_id)
job_config = QueryJobConfig()
job_config.destination = f'{client.project}.{snapshot_dataset_id}.{table_id}'
job_config.use_legacy_sql = False
job = client.query(q, job_config)
copy_table_job_ids.append(job.job_id)
job.result()
return copy_table_job_ids
def bq_create_table_as_select(google_client, dataset_id, table_name, query):
table_ref = google_client.dataset(dataset_id).table(table_name)
job_config = QueryJobConfig()
job_config.destination = table_ref
job_config.write_disposition = "WRITE_TRUNCATE"
query_job = google_client.query(query=query, job_config=job_config)
retry_count = 100
while retry_count > 0 and query_job.state != 'DONE':
retry_count -= 1
time.sleep(3)
query_job.reload() # API call
logging.info("job state : %s " % (query_job.state))
logging.info("job state : %s at %s" % (query_job.state, query_job.ended))
def load(project_id,
bq_client,
src_dataset_id,
dst_dataset_id,
overwrite_ok=False):
"""
Transform safely loaded tables and store results in target dataset.
:param project_id:
:param bq_client:
:param src_dataset_id:
:param dst_dataset_id:
:param overwrite_ok: if True and the dest dataset already exists the dataset is recreated
:return:
"""
if overwrite_ok:
bq_client.delete_dataset(dst_dataset_id,
delete_contents=True,
not_found_ok=True)
bq_client.create_dataset(dst_dataset_id)
src_tables = list(bq_client.list_tables(dataset=src_dataset_id))
job_config = QueryJobConfig()
query_jobs = []
for src_table in src_tables:
schema = bq.get_table_schema(src_table.table_id)
destination = f'{project_id}.{dst_dataset_id}.{src_table.table_id}'
table = bq_client.create_table(Table(destination, schema=schema),
exists_ok=True)
job_config.destination = table
query = SELECT_TPL.render(project_id=project_id,
dataset_id=src_dataset_id,
table=src_table.table_id,
fields=schema)
query_job = bq_client.query(query, job_config=job_config)
LOGGER.info(f'table:{destination} job_id:{query_job.job_id}')
query_jobs.append(query_job)
return query_jobs
def enrich_task():
client = Client()
# Need to use a temporary table because bq query sets field modes to NULLABLE and descriptions to null
# when writeDisposition is WRITE_TRUNCATE
# Create a temporary table
temp_table_name = '{task}_{milliseconds}'.format(
task=task, milliseconds=int(round(time.time() * 1000)))
temp_table_ref = client.dataset(dataset_name_temp).table(
temp_table_name)
table = Table(temp_table_ref)
description_path = os.path.join(
dags_folder,
'resources/stages/enrich/descriptions/{task}.txt'.format(
task=task))
table.description = read_file(description_path)
if time_partitioning_field is not None:
table.time_partitioning = TimePartitioning(
field=time_partitioning_field)
logging.info('Creating table: ' + json.dumps(table.to_api_repr()))
schema_path = os.path.join(
dags_folder,
'resources/stages/enrich/schemas/{task}.json'.format(
task=task))
schema = read_bigquery_schema_from_file(schema_path)
table.schema = schema
table = client.create_table(table)
assert table.table_id == temp_table_name
# Query from raw to temporary table
query_job_config = QueryJobConfig()
# Finishes faster, query limit for concurrent interactive queries is 50
query_job_config.priority = QueryPriority.INTERACTIVE
query_job_config.destination = temp_table_ref
sql_path = os.path.join(
dags_folder,
'resources/stages/enrich/sqls/{task}.sql'.format(task=task))
sql = read_file(sql_path, environment)
query_job = client.query(sql,
location='US',
job_config=query_job_config)
submit_bigquery_job(query_job, query_job_config)
assert query_job.state == 'DONE'
# Copy temporary table to destination
copy_job_config = CopyJobConfig()
copy_job_config.write_disposition = 'WRITE_TRUNCATE'
dest_table_name = '{task}'.format(task=task)
dest_table_ref = client.dataset(
dataset_name,
project=destination_dataset_project_id).table(dest_table_name)
copy_job = client.copy_table(temp_table_ref,
dest_table_ref,
location='US',
job_config=copy_job_config)
submit_bigquery_job(copy_job, copy_job_config)
assert copy_job.state == 'DONE'
# Delete temp table
client.delete_table(temp_table_ref)
def _config_query(self, use_legacy_sql):
job_config = QueryJobConfig()
job_config.destination = self.temp_table
job_config.use_legacy_sql = use_legacy_sql
job_config.allow_large_results = True
return job_config
def parse_url(url): # noqa: C901
query = dict(url.query) # need mutable query.
# use_legacy_sql (legacy)
if "use_legacy_sql" in query:
raise ValueError("legacy sql is not supported by this dialect")
# allow_large_results (legacy)
if "allow_large_results" in query:
raise ValueError(
"allow_large_results is only allowed for legacy sql, which is not supported by this dialect"
)
# flatten_results (legacy)
if "flatten_results" in query:
raise ValueError(
"flatten_results is only allowed for legacy sql, which is not supported by this dialect"
)
# maximum_billing_tier (deprecated)
if "maximum_billing_tier" in query:
raise ValueError("maximum_billing_tier is a deprecated argument")
project_id = url.host
location = None
dataset_id = url.database or None
arraysize = None
credentials_path = None
# location
if "location" in query:
location = query.pop("location")
# credentials_path
if "credentials_path" in query:
credentials_path = query.pop("credentials_path")
# arraysize
if "arraysize" in query:
str_arraysize = query.pop("arraysize")
try:
arraysize = int(str_arraysize)
except ValueError:
raise ValueError("invalid int in url query arraysize: " +
str_arraysize)
# if only these "non-config" values were present, the dict will now be empty
if not query:
# if a dataset_id exists, we need to return a job_config that isn't None
# so it can be updated with a dataset reference from the client
if dataset_id:
return (
project_id,
location,
dataset_id,
arraysize,
credentials_path,
QueryJobConfig(),
)
else:
return project_id, location, dataset_id, arraysize, credentials_path, None
job_config = QueryJobConfig()
# clustering_fields list(str)
if "clustering_fields" in query:
clustering_fields = GROUP_DELIMITER.split(query["clustering_fields"])
job_config.clustering_fields = list(clustering_fields)
# create_disposition
if "create_disposition" in query:
create_disposition = query["create_disposition"]
try:
job_config.create_disposition = getattr(CreateDisposition,
create_disposition)
except AttributeError:
raise ValueError("invalid create_disposition in url query: " +
create_disposition)
# default_dataset
if "default_dataset" in query or "dataset_id" in query or "project_id" in query:
raise ValueError(
"don't pass default_dataset, dataset_id, project_id in url query, instead use the url host and database"
)
# destination
if "destination" in query:
dest_project = None
dest_dataset = None
dest_table = None
try:
dest_project, dest_dataset, dest_table = query[
"destination"].split(".")
except ValueError:
raise ValueError(
"url query destination parameter should be fully qualified with project, dataset, and table"
)
job_config.destination = TableReference(
DatasetReference(dest_project, dest_dataset), dest_table)
# destination_encryption_configuration
if "destination_encryption_configuration" in query:
job_config.destination_encryption_configuration = EncryptionConfiguration(
query["destination_encryption_configuration"])
# dry_run
if "dry_run" in query:
try:
job_config.dry_run = parse_boolean(query["dry_run"])
except ValueError:
raise ValueError("invalid boolean in url query for dry_run: " +
query["dry_run"])
# labels
if "labels" in query:
label_groups = GROUP_DELIMITER.split(query["labels"])
labels = {}
for label_group in label_groups:
try:
key, value = KEY_VALUE_DELIMITER.split(label_group)
except ValueError:
raise ValueError("malformed url query in labels: " +
label_group)
labels[key] = value
job_config.labels = labels
# maximum_bytes_billed
if "maximum_bytes_billed" in query:
try:
job_config.maximum_bytes_billed = int(
query["maximum_bytes_billed"])
except ValueError:
raise ValueError(
"invalid int in url query maximum_bytes_billed: " +
query["maximum_bytes_billed"])
# priority
if "priority" in query:
try:
job_config.priority = getattr(QueryPriority, query["priority"])
except AttributeError:
raise ValueError("invalid priority in url query: " +
query["priority"])
# query_parameters
if "query_parameters" in query:
raise NotImplementedError("url query query_parameters not implemented")
# schema_update_options
if "schema_update_options" in query:
schema_update_options = GROUP_DELIMITER.split(
query["schema_update_options"])
try:
job_config.schema_update_options = [
getattr(SchemaUpdateOption, schema_update_option)
for schema_update_option in schema_update_options
]
except AttributeError:
raise ValueError("invalid schema_update_options in url query: " +
query["schema_update_options"])
# table_definitions
if "table_definitions" in query:
raise NotImplementedError(
"url query table_definitions not implemented")
# time_partitioning
if "time_partitioning" in query:
raise NotImplementedError(
"url query time_partitioning not implemented")
# udf_resources
if "udf_resources" in query:
raise NotImplementedError("url query udf_resources not implemented")
# use_query_cache
if "use_query_cache" in query:
try:
job_config.use_query_cache = parse_boolean(
query["use_query_cache"])
except ValueError:
raise ValueError(
"invalid boolean in url query for use_query_cache: " +
query["use_query_cache"])
# write_disposition
if "write_disposition" in query:
try:
job_config.write_disposition = getattr(WriteDisposition,
query["write_disposition"])
except AttributeError:
raise ValueError("invalid write_disposition in url query: " +
query["write_disposition"])
return project_id, location, dataset_id, arraysize, credentials_path, job_config
def enrich_task(ds, **kwargs):
template_context = kwargs.copy()
template_context['ds'] = ds
template_context['params'] = environment
client = Client()
# Need to use a temporary table because bq query sets field modes to NULLABLE and descriptions to null
# when writeDisposition is WRITE_TRUNCATE
# Create a temporary table
temp_table_name = '{task}_{milliseconds}'.format(
task=task, milliseconds=int(round(time.time() * 1000)))
temp_table_ref = client.dataset(dataset_name_temp).table(
temp_table_name)
table = Table(temp_table_ref)
description_path = os.path.join(
dags_folder,
'resources/stages/enrich/descriptions/{task}.txt'.format(
task=task))
table.description = read_file(description_path)
table.time_partitioning = TimePartitioning(
field=time_partitioning_field)
logging.info('Creating table: ' + json.dumps(table.to_api_repr()))
schema_path = os.path.join(
dags_folder,
'resources/stages/enrich/schemas/{task}.json'.format(
task=task))
schema = read_bigquery_schema_from_file(schema_path)
table.schema = schema
table = client.create_table(table)
assert table.table_id == temp_table_name
# Query from raw to temporary table
query_job_config = QueryJobConfig()
# Finishes faster, query limit for concurrent interactive queries is 50
query_job_config.priority = QueryPriority.INTERACTIVE
query_job_config.destination = temp_table_ref
sql_path = os.path.join(
dags_folder,
'resources/stages/enrich/sqls/{task}.sql'.format(task=task))
sql_template = read_file(sql_path)
sql = kwargs['task'].render_template('', sql_template,
template_context)
print('Enrichment sql:')
print(sql)
query_job = client.query(sql,
location='US',
job_config=query_job_config)
submit_bigquery_job(query_job, query_job_config)
assert query_job.state == 'DONE'
if load_all_partitions:
# Copy temporary table to destination
copy_job_config = CopyJobConfig()
copy_job_config.write_disposition = 'WRITE_TRUNCATE'
dest_table_name = '{task}'.format(task=task)
dest_table_ref = client.dataset(
dataset_name,
project=destination_dataset_project_id).table(
dest_table_name)
copy_job = client.copy_table(temp_table_ref,
dest_table_ref,
location='US',
job_config=copy_job_config)
submit_bigquery_job(copy_job, copy_job_config)
assert copy_job.state == 'DONE'
else:
# Merge
# https://cloud.google.com/bigquery/docs/reference/standard-sql/dml-syntax#merge_statement
merge_job_config = QueryJobConfig()
# Finishes faster, query limit for concurrent interactive queries is 50
merge_job_config.priority = QueryPriority.INTERACTIVE
merge_sql_path = os.path.join(
dags_folder,
'resources/stages/enrich/sqls/merge_{task}.sql'.format(
task=task))
merge_sql_template = read_file(merge_sql_path)
template_context['params']['source_table'] = temp_table_name
merge_sql = kwargs['task'].render_template(
'', merge_sql_template, template_context)
print('Merge sql:')
print(merge_sql)
merge_job = client.query(merge_sql,
location='US',
job_config=merge_job_config)
submit_bigquery_job(merge_job, merge_job_config)
assert merge_job.state == 'DONE'
# Delete temp table
client.delete_table(temp_table_ref)
