Source code for airflow.providers.google.cloud.hooks.bigquery
## Licensed to the Apache Software Foundation (ASF) under one# or more contributor license agreements. See the NOTICE file# distributed with this work for additional information# regarding copyright ownership. The ASF licenses this file# to you under the Apache License, Version 2.0 (the# "License"); you may not use this file except in compliance# with the License. You may obtain a copy of the License at## http://www.apache.org/licenses/LICENSE-2.0## Unless required by applicable law or agreed to in writing,# software distributed under the License is distributed on an# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY# KIND, either express or implied. See the License for the# specific language governing permissions and limitations# under the License."""This module contains a BigQuery Hook, as well as a very basic PEP 249implementation for BigQuery."""from__future__importannotationsimporthashlibimportjsonimportloggingimportreimporttimeimportuuidimportwarningsfromcopyimportdeepcopyfromdatetimeimportdatetime,timedeltafromtypingimportAny,Iterable,Mapping,NoReturn,Sequence,Union,castfromaiohttpimportClientSessionasClientSessionfromgcloud.aio.bigqueryimportJob,TableasTable_asyncfromgoogle.api_core.page_iteratorimportHTTPIteratorfromgoogle.api_core.retryimportRetryfromgoogle.cloud.bigqueryimport(DEFAULT_RETRY,Client,CopyJob,ExternalConfig,ExtractJob,LoadJob,QueryJob,SchemaField,)fromgoogle.cloud.bigquery.datasetimportAccessEntry,Dataset,DatasetListItem,DatasetReferencefromgoogle.cloud.bigquery.tableimportEncryptionConfiguration,Row,RowIterator,Table,TableReferencefromgoogle.cloud.exceptionsimportNotFoundfromgoogleapiclient.discoveryimportResource,buildfrompandasimportDataFramefrompandas_gbqimportread_gbqfrompandas_gbq.gbqimportGbqConnector# noqafromrequestsimportSessionfromsqlalchemyimportcreate_enginefromairflow.exceptionsimportAirflowExceptionfromairflow.providers.common.sql.hooks.sqlimportDbApiHookfromairflow.providers.google.cloud.utils.bigqueryimportbq_castfromairflow.providers.google.common.constsimportCLIENT_INFOfromairflow.providers.google.common.hooks.base_googleimportGoogleBaseAsyncHook,GoogleBaseHook,get_fieldfromairflow.utils.helpersimportconvert_camel_to_snakefromairflow.utils.log.logging_mixinimportLoggingMixin
[docs]classBigQueryHook(GoogleBaseHook,DbApiHook):""" Interact with BigQuery. This hook uses the Google Cloud connection. :param gcp_conn_id: The Airflow connection used for GCP credentials. :param use_legacy_sql: This specifies whether to use legacy SQL dialect. :param location: The location of the BigQuery resource. :param api_resource_configs: This contains params configuration applied for Google BigQuery jobs. :param impersonation_chain: This is the optional service account to impersonate using short term credentials. :param labels: The BigQuery resource label. """
def__init__(self,gcp_conn_id:str=GoogleBaseHook.default_conn_name,use_legacy_sql:bool=True,location:str|None=None,api_resource_configs:dict|None=None,impersonation_chain:str|Sequence[str]|None=None,labels:dict|None=None,**kwargs,)->None:ifkwargs.get("delegate_to")isnotNone:raiseRuntimeError("The `delegate_to` parameter has been deprecated before and finally removed in this version"" of Google Provider. You MUST convert it to `impersonate_chain`")super().__init__(gcp_conn_id=gcp_conn_id,impersonation_chain=impersonation_chain,)self.use_legacy_sql=use_legacy_sqlself.location=locationself.running_job_id:str|None=Noneself.api_resource_configs:dict=api_resource_configsifapi_resource_configselse{}self.labels=labelsself.credentials_path="bigquery_hook_credentials.json"
[docs]defget_conn(self)->BigQueryConnection:"""Returns a BigQuery PEP 249 connection object."""service=self.get_service()returnBigQueryConnection(service=service,project_id=self.project_id,use_legacy_sql=self.use_legacy_sql,location=self.location,num_retries=self.num_retries,hook=self,
)
[docs]defget_service(self)->Resource:"""Returns a BigQuery service object."""warnings.warn("This method will be deprecated. Please use `BigQueryHook.get_client` method",DeprecationWarning)http_authorized=self._authorize()returnbuild("bigquery","v2",http=http_authorized,cache_discovery=False)
[docs]defget_client(self,project_id:str|None=None,location:str|None=None)->Client:""" Returns authenticated BigQuery Client. :param project_id: Project ID for the project which the client acts on behalf of. :param location: Default location for jobs / datasets / tables. :return: """returnClient(client_info=CLIENT_INFO,project=project_id,location=location,credentials=self.get_credentials(),
)
[docs]defget_uri(self)->str:"""Override DbApiHook get_uri method for get_sqlalchemy_engine()"""returnf"bigquery://{self.project_id}"
[docs]defget_sqlalchemy_engine(self,engine_kwargs=None):""" Get an sqlalchemy_engine object. :param engine_kwargs: Kwargs used in :func:`~sqlalchemy.create_engine`. :return: the created engine. """ifengine_kwargsisNone:engine_kwargs={}extras=self.get_connection(self.gcp_conn_id).extra_dejsoncredentials_path=get_field(extras,"key_path")ifcredentials_path:returncreate_engine(self.get_uri(),credentials_path=credentials_path,**engine_kwargs)keyfile_dict=get_field(extras,"keyfile_dict")ifkeyfile_dict:keyfile_content=keyfile_dictifisinstance(keyfile_dict,dict)elsejson.loads(keyfile_dict)returncreate_engine(self.get_uri(),credentials_info=keyfile_content,**engine_kwargs)try:# 1. If the environment variable GOOGLE_APPLICATION_CREDENTIALS is set# ADC uses the service account key or configuration file that the variable points to.# 2. If the environment variable GOOGLE_APPLICATION_CREDENTIALS isn't set# ADC uses the service account that is attached to the resource that is running your code.returncreate_engine(self.get_uri(),**engine_kwargs)exceptExceptionase:self.log.error(e)raiseAirflowException("For now, we only support instantiating SQLAlchemy engine by"
" using ADC or extra fields `key_path` and `keyfile_dict`.")
[docs]defget_records(self,sql,parameters=None):ifself.locationisNone:raiseAirflowException("Need to specify 'location' to use BigQueryHook.get_records()")returnsuper().get_records(sql,parameters=parameters)
@staticmethoddef_resolve_table_reference(table_resource:dict[str,Any],project_id:str|None=None,dataset_id:str|None=None,table_id:str|None=None,)->dict[str,Any]:try:# Check if tableReference is present and is validTableReference.from_api_repr(table_resource["tableReference"])exceptKeyError:# Something is wrong so we try to build the referencetable_resource["tableReference"]=table_resource.get("tableReference",{})values=[("projectId",project_id),("tableId",table_id),("datasetId",dataset_id)]forkey,valueinvalues:# Check if value is already present if no use the provided oneresolved_value=table_resource["tableReference"].get(key,value)ifnotresolved_value:# If there's no value in tableReference and provided one is None raise errorraiseAirflowException(f"Table resource is missing proper `tableReference` and `{key}` is None")table_resource["tableReference"][key]=resolved_valuereturntable_resource
[docs]definsert_rows(self,table:Any,rows:Any,target_fields:Any=None,commit_every:Any=1000,replace:Any=False,**kwargs,)->None:""" Insertion is currently unsupported. Theoretically, you could use BigQuery's streaming API to insert rows into a table, but this hasn't been implemented. """raiseNotImplementedError()
[docs]defget_pandas_df(self,sql:str,parameters:Iterable|Mapping|None=None,dialect:str|None=None,**kwargs,)->DataFrame:""" Returns a Pandas DataFrame for the results produced by a BigQuery query. The DbApiHook method must be overridden because Pandas doesn't support PEP 249 connections, except for SQLite. See: https://github.com/pandas-dev/pandas/blob/055d008615272a1ceca9720dc365a2abd316f353/pandas/io/sql.py#L415 https://github.com/pandas-dev/pandas/issues/6900 :param sql: The BigQuery SQL to execute. :param parameters: The parameters to render the SQL query with (not used, leave to override superclass method) :param dialect: Dialect of BigQuery SQL – legacy SQL or standard SQL defaults to use `self.use_legacy_sql` if not specified :param kwargs: (optional) passed into pandas_gbq.read_gbq method """ifdialectisNone:dialect="legacy"ifself.use_legacy_sqlelse"standard"credentials,project_id=self.get_credentials_and_project_id()returnread_gbq(sql,project_id=project_id,dialect=dialect,verbose=False,credentials=credentials,**kwargs
)@GoogleBaseHook.fallback_to_default_project_id
[docs]deftable_exists(self,dataset_id:str,table_id:str,project_id:str)->bool:""" Checks for the existence of a table in Google BigQuery. :param project_id: The Google cloud project in which to look for the table. The connection supplied to the hook must provide access to the specified project. :param dataset_id: The name of the dataset in which to look for the table. :param table_id: The name of the table to check the existence of. """table_reference=TableReference(DatasetReference(project_id,dataset_id),table_id)try:self.get_client(project_id=project_id).get_table(table_reference)returnTrueexceptNotFound:returnFalse
@GoogleBaseHook.fallback_to_default_project_id
[docs]deftable_partition_exists(self,dataset_id:str,table_id:str,partition_id:str,project_id:str)->bool:""" Checks for the existence of a partition in a table in Google BigQuery. :param project_id: The Google cloud project in which to look for the table. The connection supplied to the hook must provide access to the specified project. :param dataset_id: The name of the dataset in which to look for the table. :param table_id: The name of the table to check the existence of. :param partition_id: The name of the partition to check the existence of. """table_reference=TableReference(DatasetReference(project_id,dataset_id),table_id)try:returnpartition_idinself.get_client(project_id=project_id).list_partitions(table_reference)exceptNotFound:returnFalse
@GoogleBaseHook.fallback_to_default_project_id
[docs]defcreate_empty_table(self,project_id:str|None=None,dataset_id:str|None=None,table_id:str|None=None,table_resource:dict[str,Any]|None=None,schema_fields:list|None=None,time_partitioning:dict|None=None,cluster_fields:list[str]|None=None,labels:dict|None=None,view:dict|None=None,materialized_view:dict|None=None,encryption_configuration:dict|None=None,retry:Retry|None=DEFAULT_RETRY,location:str|None=None,exists_ok:bool=True,)->Table:""" Creates a new, empty table in the dataset. To create a view, which is defined by a SQL query, parse a dictionary to 'view' kwarg :param project_id: The project to create the table into. :param dataset_id: The dataset to create the table into. :param table_id: The Name of the table to be created. :param table_resource: Table resource as described in documentation: https://cloud.google.com/bigquery/docs/reference/rest/v2/tables#Table If provided all other parameters are ignored. :param schema_fields: If set, the schema field list as defined here: https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs#configuration.load.schema :param labels: a dictionary containing labels for the table, passed to BigQuery :param retry: Optional. How to retry the RPC. **Example**: :: schema_fields=[{"name": "emp_name", "type": "STRING", "mode": "REQUIRED"}, {"name": "salary", "type": "INTEGER", "mode": "NULLABLE"}] :param time_partitioning: configure optional time partitioning fields i.e. partition by field, type and expiration as per API specifications. .. seealso:: https://cloud.google.com/bigquery/docs/reference/rest/v2/tables#timePartitioning :param cluster_fields: [Optional] The fields used for clustering. BigQuery supports clustering for both partitioned and non-partitioned tables. https://cloud.google.com/bigquery/docs/reference/rest/v2/tables#clustering.fields :param view: [Optional] A dictionary containing definition for the view. If set, it will create a view instead of a table: https://cloud.google.com/bigquery/docs/reference/rest/v2/tables#ViewDefinition **Example**: :: view = { "query": "SELECT * FROM `test-project-id.test_dataset_id.test_table_prefix*` LIMIT 1000", "useLegacySql": False } :param materialized_view: [Optional] The materialized view definition. :param encryption_configuration: [Optional] Custom encryption configuration (e.g., Cloud KMS keys). **Example**: :: encryption_configuration = { "kmsKeyName": "projects/testp/locations/us/keyRings/test-kr/cryptoKeys/test-key" } :param num_retries: Maximum number of retries in case of connection problems. :param location: (Optional) The geographic location where the table should reside. :param exists_ok: If ``True``, ignore "already exists" errors when creating the table. :return: Created table """_table_resource:dict[str,Any]={}ifself.location:_table_resource["location"]=self.locationifschema_fields:_table_resource["schema"]={"fields":schema_fields}iftime_partitioning:_table_resource["timePartitioning"]=time_partitioningifcluster_fields:_table_resource["clustering"]={"fields":cluster_fields}iflabels:_table_resource["labels"]=labelsifview:_table_resource["view"]=viewifmaterialized_view:_table_resource["materializedView"]=materialized_viewifencryption_configuration:_table_resource["encryptionConfiguration"]=encryption_configurationtable_resource=table_resourceor_table_resourcetable_resource=self._resolve_table_reference(table_resource=table_resource,project_id=project_id,dataset_id=dataset_id,table_id=table_id,)table=Table.from_api_repr(table_resource)returnself.get_client(project_id=project_id,location=location).create_table(table=table,exists_ok=exists_ok,retry=retry
)@GoogleBaseHook.fallback_to_default_project_id
[docs]defcreate_empty_dataset(self,dataset_id:str|None=None,project_id:str|None=None,location:str|None=None,dataset_reference:dict[str,Any]|None=None,exists_ok:bool=True,)->dict[str,Any]:""" Create a new empty dataset: https://cloud.google.com/bigquery/docs/reference/rest/v2/datasets/insert :param project_id: The name of the project where we want to create an empty a dataset. Don't need to provide, if projectId in dataset_reference. :param dataset_id: The id of dataset. Don't need to provide, if datasetId in dataset_reference. :param location: (Optional) The geographic location where the dataset should reside. There is no default value but the dataset will be created in US if nothing is provided. :param dataset_reference: Dataset reference that could be provided with request body. More info: https://cloud.google.com/bigquery/docs/reference/rest/v2/datasets#resource :param exists_ok: If ``True``, ignore "already exists" errors when creating the dataset. """dataset_reference=dataset_referenceor{}if"datasetReference"notindataset_reference:dataset_reference["datasetReference"]={}forparam,valueinzip(["datasetId","projectId"],[dataset_id,project_id]):specified_param=dataset_reference["datasetReference"].get(param)ifspecified_param:ifvalue:self.log.info("`%s` was provided in both `dataset_reference` and as `%s`. ""Using value from `dataset_reference`",param,convert_camel_to_snake(param),)continue# use specified valueifnotvalue:raiseValueError(f"Please specify `{param}` either in `dataset_reference` "f"or by providing `{convert_camel_to_snake(param)}`",)# dataset_reference has no param but we can fallback to default valueself.log.info("%s was not specified in `dataset_reference`. Will use default value %s.",param,value)dataset_reference["datasetReference"][param]=valuelocation=locationorself.locationproject_id=project_idorself.project_idiflocation:dataset_reference["location"]=dataset_reference.get("location",location)dataset:Dataset=Dataset.from_api_repr(dataset_reference)self.log.info("Creating dataset: %s in project: %s ",dataset.dataset_id,dataset.project)dataset_object=self.get_client(project_id=project_id,location=location).create_dataset(dataset=dataset,exists_ok=exists_ok)self.log.info("Dataset created successfully.")returndataset_object.to_api_repr()
@GoogleBaseHook.fallback_to_default_project_id
[docs]defget_dataset_tables(self,dataset_id:str,project_id:str|None=None,max_results:int|None=None,retry:Retry=DEFAULT_RETRY,)->list[dict[str,Any]]:""" Get the list of tables for a given dataset. For more information, see: https://cloud.google.com/bigquery/docs/reference/rest/v2/tables/list :param dataset_id: the dataset ID of the requested dataset. :param project_id: (Optional) the project of the requested dataset. If None, self.project_id will be used. :param max_results: (Optional) the maximum number of tables to return. :param retry: How to retry the RPC. :return: List of tables associated with the dataset. """self.log.info("Start getting tables list from dataset: %s.%s",project_id,dataset_id)tables=self.get_client().list_tables(dataset=DatasetReference(project=project_id,dataset_id=dataset_id),max_results=max_results,retry=retry,)# Convert to a list (consumes all values)return[t.reference.to_api_repr()fortintables]
@GoogleBaseHook.fallback_to_default_project_id
[docs]defdelete_dataset(self,dataset_id:str,project_id:str|None=None,delete_contents:bool=False,retry:Retry=DEFAULT_RETRY,)->None:""" Delete a dataset of Big query in your project. :param project_id: The name of the project where we have the dataset. :param dataset_id: The dataset to be delete. :param delete_contents: If True, delete all the tables in the dataset. If False and the dataset contains tables, the request will fail. :param retry: How to retry the RPC. """self.log.info("Deleting from project: %s Dataset:%s",project_id,dataset_id)self.get_client(project_id=project_id).delete_dataset(dataset=DatasetReference(project=project_id,dataset_id=dataset_id),delete_contents=delete_contents,retry=retry,not_found_ok=True,
)@GoogleBaseHook.fallback_to_default_project_id
[docs]defcreate_external_table(self,external_project_dataset_table:str,schema_fields:list,source_uris:list,source_format:str="CSV",autodetect:bool=False,compression:str="NONE",ignore_unknown_values:bool=False,max_bad_records:int=0,skip_leading_rows:int=0,field_delimiter:str=",",quote_character:str|None=None,allow_quoted_newlines:bool=False,allow_jagged_rows:bool=False,encoding:str="UTF-8",src_fmt_configs:dict|None=None,labels:dict|None=None,description:str|None=None,encryption_configuration:dict|None=None,location:str|None=None,project_id:str|None=None,)->Table:""" Creates a new external table in the dataset with the data from Google Cloud Storage. See here: https://cloud.google.com/bigquery/docs/reference/rest/v2/tables#resource for more details about these parameters. :param external_project_dataset_table: The dotted ``(<project>.|<project>:)<dataset>.<table>($<partition>)`` BigQuery table name to create external table. If ``<project>`` is not included, project will be the project defined in the connection json. :param schema_fields: The schema field list as defined here: https://cloud.google.com/bigquery/docs/reference/rest/v2/tables#resource :param source_uris: The source Google Cloud Storage URI (e.g. gs://some-bucket/some-file.txt). A single wild per-object name can be used. :param source_format: File format to export. :param autodetect: Try to detect schema and format options automatically. Any option specified explicitly will be honored. :param compression: [Optional] The compression type of the data source. Possible values include GZIP and NONE. The default value is NONE. This setting is ignored for Google Cloud Bigtable, Google Cloud Datastore backups and Avro formats. :param ignore_unknown_values: [Optional] Indicates if BigQuery should allow extra values that are not represented in the table schema. If true, the extra values are ignored. If false, records with extra columns are treated as bad records, and if there are too many bad records, an invalid error is returned in the job result. :param max_bad_records: The maximum number of bad records that BigQuery can ignore when running the job. :param skip_leading_rows: Number of rows to skip when loading from a CSV. :param field_delimiter: The delimiter to use when loading from a CSV. :param quote_character: The value that is used to quote data sections in a CSV file. :param allow_quoted_newlines: Whether to allow quoted newlines (true) or not (false). :param allow_jagged_rows: Accept rows that are missing trailing optional columns. The missing values are treated as nulls. If false, records with missing trailing columns are treated as bad records, and if there are too many bad records, an invalid error is returned in the job result. Only applicable when source_format is CSV. :param encoding: The character encoding of the data. See: .. seealso:: https://cloud.google.com/bigquery/docs/reference/rest/v2/tables#externalDataConfiguration.csvOptions.encoding :param src_fmt_configs: configure optional fields specific to the source format :param labels: A dictionary containing labels for the BiqQuery table. :param description: A string containing the description for the BigQuery table. :param encryption_configuration: [Optional] Custom encryption configuration (e.g., Cloud KMS keys). **Example**: :: encryption_configuration = { "kmsKeyName": "projects/testp/locations/us/keyRings/test-kr/cryptoKeys/test-key" } """warnings.warn("This method is deprecated. Please use `BigQueryHook.create_empty_table` method with ""passing the `table_resource` object. This gives more flexibility than this method.",DeprecationWarning,)location=locationorself.locationsrc_fmt_configs=src_fmt_configsor{}source_format=source_format.upper()compression=compression.upper()external_config_api_repr={"autodetect":autodetect,"sourceFormat":source_format,"sourceUris":source_uris,"compression":compression,"ignoreUnknownValues":ignore_unknown_values,}# if following fields are not specified in src_fmt_configs,# honor the top-level params for backward-compatibilitybackward_compatibility_configs={"skipLeadingRows":skip_leading_rows,"fieldDelimiter":field_delimiter,"quote":quote_character,"allowQuotedNewlines":allow_quoted_newlines,"allowJaggedRows":allow_jagged_rows,"encoding":encoding,}src_fmt_to_param_mapping={"CSV":"csvOptions","GOOGLE_SHEETS":"googleSheetsOptions"}src_fmt_to_configs_mapping={"csvOptions":["allowJaggedRows","allowQuotedNewlines","fieldDelimiter","skipLeadingRows","quote","encoding",],"googleSheetsOptions":["skipLeadingRows"],}ifsource_formatinsrc_fmt_to_param_mapping.keys():valid_configs=src_fmt_to_configs_mapping[src_fmt_to_param_mapping[source_format]]src_fmt_configs=_validate_src_fmt_configs(source_format,src_fmt_configs,valid_configs,backward_compatibility_configs)external_config_api_repr[src_fmt_to_param_mapping[source_format]]=src_fmt_configs# build external configexternal_config=ExternalConfig.from_api_repr(external_config_api_repr)ifschema_fields:external_config.schema=[SchemaField.from_api_repr(f)forfinschema_fields]ifmax_bad_records:external_config.max_bad_records=max_bad_records# build table definitiontable=Table(table_ref=TableReference.from_string(external_project_dataset_table,project_id))table.external_data_configuration=external_configiflabels:table.labels=labelsifdescription:table.description=descriptionifencryption_configuration:table.encryption_configuration=EncryptionConfiguration.from_api_repr(encryption_configuration)self.log.info("Creating external table: %s",external_project_dataset_table)table_object=self.create_empty_table(table_resource=table.to_api_repr(),project_id=project_id,location=location,exists_ok=True)self.log.info("External table created successfully: %s",external_project_dataset_table)returntable_object
@GoogleBaseHook.fallback_to_default_project_id
[docs]defupdate_table(self,table_resource:dict[str,Any],fields:list[str]|None=None,dataset_id:str|None=None,table_id:str|None=None,project_id:str|None=None,)->dict[str,Any]:""" Change some fields of a table. Use ``fields`` to specify which fields to update. At least one field must be provided. If a field is listed in ``fields`` and is ``None`` in ``table``, the field value will be deleted. If ``table.etag`` is not ``None``, the update will only succeed if the table on the server has the same ETag. Thus reading a table with ``get_table``, changing its fields, and then passing it to ``update_table`` will ensure that the changes will only be saved if no modifications to the table occurred since the read. :param project_id: The project to create the table into. :param dataset_id: The dataset to create the table into. :param table_id: The Name of the table to be created. :param table_resource: Table resource as described in documentation: https://cloud.google.com/bigquery/docs/reference/rest/v2/tables#Table The table has to contain ``tableReference`` or ``project_id``, ``dataset_id`` and ``table_id`` have to be provided. :param fields: The fields of ``table`` to change, spelled as the Table properties (e.g. "friendly_name"). """fields=fieldsorlist(table_resource.keys())table_resource=self._resolve_table_reference(table_resource=table_resource,project_id=project_id,dataset_id=dataset_id,table_id=table_id)table=Table.from_api_repr(table_resource)self.log.info("Updating table: %s",table_resource["tableReference"])table_object=self.get_client(project_id=project_id).update_table(table=table,fields=fields)self.log.info("Table %s.%s.%s updated successfully",project_id,dataset_id,table_id)returntable_object.to_api_repr()
@GoogleBaseHook.fallback_to_default_project_id
[docs]defpatch_table(self,dataset_id:str,table_id:str,project_id:str|None=None,description:str|None=None,expiration_time:int|None=None,external_data_configuration:dict|None=None,friendly_name:str|None=None,labels:dict|None=None,schema:list|None=None,time_partitioning:dict|None=None,view:dict|None=None,require_partition_filter:bool|None=None,encryption_configuration:dict|None=None,)->None:""" Patch information in an existing table. It only updates fields that are provided in the request object. Reference: https://cloud.google.com/bigquery/docs/reference/rest/v2/tables/patch :param dataset_id: The dataset containing the table to be patched. :param table_id: The Name of the table to be patched. :param project_id: The project containing the table to be patched. :param description: [Optional] A user-friendly description of this table. :param expiration_time: [Optional] The time when this table expires, in milliseconds since the epoch. :param external_data_configuration: [Optional] A dictionary containing properties of a table stored outside of BigQuery. :param friendly_name: [Optional] A descriptive name for this table. :param labels: [Optional] A dictionary containing labels associated with this table. :param schema: [Optional] If set, the schema field list as defined here: https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs#configuration.load.schema The supported schema modifications and unsupported schema modification are listed here: https://cloud.google.com/bigquery/docs/managing-table-schemas **Example**: :: schema=[{"name": "emp_name", "type": "STRING", "mode": "REQUIRED"}, {"name": "salary", "type": "INTEGER", "mode": "NULLABLE"}] :param time_partitioning: [Optional] A dictionary containing time-based partitioning definition for the table. :param view: [Optional] A dictionary containing definition for the view. If set, it will patch a view instead of a table: https://cloud.google.com/bigquery/docs/reference/rest/v2/tables#ViewDefinition **Example**: :: view = { "query": "SELECT * FROM `test-project-id.test_dataset_id.test_table_prefix*` LIMIT 500", "useLegacySql": False } :param require_partition_filter: [Optional] If true, queries over the this table require a partition filter. If false, queries over the table :param encryption_configuration: [Optional] Custom encryption configuration (e.g., Cloud KMS keys). **Example**: :: encryption_configuration = { "kmsKeyName": "projects/testp/locations/us/keyRings/test-kr/cryptoKeys/test-key" } """warnings.warn("This method is deprecated, please use ``BigQueryHook.update_table`` method.",DeprecationWarning,)table_resource:dict[str,Any]={}ifdescriptionisnotNone:table_resource["description"]=descriptionifexpiration_timeisnotNone:table_resource["expirationTime"]=expiration_timeifexternal_data_configuration:table_resource["externalDataConfiguration"]=external_data_configurationiffriendly_nameisnotNone:table_resource["friendlyName"]=friendly_nameiflabels:table_resource["labels"]=labelsifschema:table_resource["schema"]={"fields":schema}iftime_partitioning:table_resource["timePartitioning"]=time_partitioningifview:table_resource["view"]=viewifrequire_partition_filterisnotNone:table_resource["requirePartitionFilter"]=require_partition_filterifencryption_configuration:table_resource["encryptionConfiguration"]=encryption_configurationself.update_table(table_resource=table_resource,fields=list(table_resource.keys()),project_id=project_id,dataset_id=dataset_id,table_id=table_id,
)@GoogleBaseHook.fallback_to_default_project_id
[docs]definsert_all(self,project_id:str,dataset_id:str,table_id:str,rows:list,ignore_unknown_values:bool=False,skip_invalid_rows:bool=False,fail_on_error:bool=False,)->None:""" Method to stream data into BigQuery one record at a time without needing to run a load job .. seealso:: For more information, see: https://cloud.google.com/bigquery/docs/reference/rest/v2/tabledata/insertAll :param project_id: The name of the project where we have the table :param dataset_id: The name of the dataset where we have the table :param table_id: The name of the table :param rows: the rows to insert **Example or rows**: rows=[{"json": {"a_key": "a_value_0"}}, {"json": {"a_key": "a_value_1"}}] :param ignore_unknown_values: [Optional] Accept rows that contain values that do not match the schema. The unknown values are ignored. The default value is false, which treats unknown values as errors. :param skip_invalid_rows: [Optional] Insert all valid rows of a request, even if invalid rows exist. The default value is false, which causes the entire request to fail if any invalid rows exist. :param fail_on_error: [Optional] Force the task to fail if any errors occur. The default value is false, which indicates the task should not fail even if any insertion errors occur. """self.log.info("Inserting %s row(s) into table %s:%s.%s",len(rows),project_id,dataset_id,table_id)table_ref=TableReference(dataset_ref=DatasetReference(project_id,dataset_id),table_id=table_id)bq_client=self.get_client(project_id=project_id)table=bq_client.get_table(table_ref)errors=bq_client.insert_rows(table=table,rows=rows,ignore_unknown_values=ignore_unknown_values,skip_invalid_rows=skip_invalid_rows,)iferrors:error_msg=f"{len(errors)} insert error(s) occurred. Details: {errors}"self.log.error(error_msg)iffail_on_error:raiseAirflowException(f"BigQuery job failed. Error was: {error_msg}")else:self.log.info("All row(s) inserted successfully: %s:%s.%s",project_id,dataset_id,table_id)
@GoogleBaseHook.fallback_to_default_project_id
[docs]defupdate_dataset(self,fields:Sequence[str],dataset_resource:dict[str,Any],dataset_id:str|None=None,project_id:str|None=None,retry:Retry=DEFAULT_RETRY,)->Dataset:""" Change some fields of a dataset. Use ``fields`` to specify which fields to update. At least one field must be provided. If a field is listed in ``fields`` and is ``None`` in ``dataset``, it will be deleted. If ``dataset.etag`` is not ``None``, the update will only succeed if the dataset on the server has the same ETag. Thus reading a dataset with ``get_dataset``, changing its fields, and then passing it to ``update_dataset`` will ensure that the changes will only be saved if no modifications to the dataset occurred since the read. :param dataset_resource: Dataset resource that will be provided in request body. https://cloud.google.com/bigquery/docs/reference/rest/v2/datasets#resource :param dataset_id: The id of the dataset. :param fields: The properties of ``dataset`` to change (e.g. "friendly_name"). :param project_id: The Google Cloud Project ID :param retry: How to retry the RPC. """dataset_resource["datasetReference"]=dataset_resource.get("datasetReference",{})forkey,valueinzip(["datasetId","projectId"],[dataset_id,project_id]):spec_value=dataset_resource["datasetReference"].get(key)ifvalueandnotspec_value:dataset_resource["datasetReference"][key]=valueself.log.info("Start updating dataset")dataset=self.get_client(project_id=project_id).update_dataset(dataset=Dataset.from_api_repr(dataset_resource),fields=fields,retry=retry,)self.log.info("Dataset successfully updated: %s",dataset)returndataset
[docs]defpatch_dataset(self,dataset_id:str,dataset_resource:dict,project_id:str|None=None)->dict:""" Patches information in an existing dataset. It only replaces fields that are provided in the submitted dataset resource. More info: https://cloud.google.com/bigquery/docs/reference/rest/v2/datasets/patch :param dataset_id: The BigQuery Dataset ID :param dataset_resource: Dataset resource that will be provided in request body. https://cloud.google.com/bigquery/docs/reference/rest/v2/datasets#resource :param project_id: The Google Cloud Project ID """warnings.warn("This method is deprecated. Please use ``update_dataset``.",DeprecationWarning)project_id=project_idorself.project_idifnotdataset_idornotisinstance(dataset_id,str):raiseValueError(f"dataset_id argument must be provided and has a type 'str'. You provided: {dataset_id}")service=self.get_service()dataset_project_id=project_idorself.project_idself.log.info("Start patching dataset: %s:%s",dataset_project_id,dataset_id)dataset=(service.datasets().patch(datasetId=dataset_id,projectId=dataset_project_id,body=dataset_resource,).execute(num_retries=self.num_retries))self.log.info("Dataset successfully patched: %s",dataset)returndataset
[docs]defget_dataset_tables_list(self,dataset_id:str,project_id:str|None=None,table_prefix:str|None=None,max_results:int|None=None,)->list[dict[str,Any]]:""" Method returns tables list of a BigQuery tables. If table prefix is specified, only tables beginning by it are returned. For more information, see: https://cloud.google.com/bigquery/docs/reference/rest/v2/tables/list :param dataset_id: The BigQuery Dataset ID :param project_id: The Google Cloud Project ID :param table_prefix: Tables must begin by this prefix to be returned (case sensitive) :param max_results: The maximum number of results to return in a single response page. Leverage the page tokens to iterate through the entire collection. :return: List of tables associated with the dataset """warnings.warn("This method is deprecated. Please use ``get_dataset_tables``.",DeprecationWarning)project_id=project_idorself.project_idtables=self.get_client().list_tables(dataset=DatasetReference(project=project_id,dataset_id=dataset_id),max_results=max_results,)iftable_prefix:result=[t.reference.to_api_repr()fortintablesift.table_id.startswith(table_prefix)]else:result=[t.reference.to_api_repr()fortintables]self.log.info("%s tables found",len(result))returnresult
@GoogleBaseHook.fallback_to_default_project_id
[docs]defget_datasets_list(self,project_id:str|None=None,include_all:bool=False,filter_:str|None=None,max_results:int|None=None,page_token:str|None=None,retry:Retry=DEFAULT_RETRY,return_iterator:bool=False,)->list[DatasetListItem]|HTTPIterator:""" Method returns full list of BigQuery datasets in the current project For more information, see: https://cloud.google.com/bigquery/docs/reference/rest/v2/datasets/list :param project_id: Google Cloud Project for which you try to get all datasets :param include_all: True if results include hidden datasets. Defaults to False. :param filter_: An expression for filtering the results by label. For syntax, see https://cloud.google.com/bigquery/docs/reference/rest/v2/datasets/list#filter. :param filter_: str :param max_results: Maximum number of datasets to return. :param max_results: int :param page_token: Token representing a cursor into the datasets. If not passed, the API will return the first page of datasets. The token marks the beginning of the iterator to be returned and the value of the ``page_token`` can be accessed at ``next_page_token`` of the :class:`~google.api_core.page_iterator.HTTPIterator`. :param page_token: str :param retry: How to retry the RPC. :param return_iterator: Instead of returning a list[Row], returns a HTTPIterator which can be used to obtain the next_page_token property. """iterator=self.get_client(project_id=project_id).list_datasets(project=project_id,include_all=include_all,filter=filter_,max_results=max_results,page_token=page_token,retry=retry,)# If iterator is requested, we cannot perform a list() on it to log the number# of datasets because we will have started iterationifreturn_iterator:returniteratordatasets_list=list(iterator)self.log.info("Datasets List: %s",len(datasets_list))returndatasets_list
@GoogleBaseHook.fallback_to_default_project_id
[docs]defget_dataset(self,dataset_id:str,project_id:str|None=None)->Dataset:""" Fetch the dataset referenced by dataset_id. :param dataset_id: The BigQuery Dataset ID :param project_id: The Google Cloud Project ID :return: dataset_resource .. seealso:: For more information, see Dataset Resource content: https://cloud.google.com/bigquery/docs/reference/rest/v2/datasets#resource """dataset=self.get_client(project_id=project_id).get_dataset(dataset_ref=DatasetReference(project_id,dataset_id))self.log.info("Dataset Resource: %s",dataset)returndataset
@GoogleBaseHook.fallback_to_default_project_id
[docs]defrun_grant_dataset_view_access(self,source_dataset:str,view_dataset:str,view_table:str,view_project:str|None=None,project_id:str|None=None,)->dict[str,Any]:""" Grant authorized view access of a dataset to a view table. If this view has already been granted access to the dataset, do nothing. This method is not atomic. Running it may clobber a simultaneous update. :param source_dataset: the source dataset :param view_dataset: the dataset that the view is in :param view_table: the table of the view :param project_id: the project of the source dataset. If None, self.project_id will be used. :param view_project: the project that the view is in. If None, self.project_id will be used. :return: the datasets resource of the source dataset. """view_project=view_projectorproject_idview_access=AccessEntry(role=None,entity_type="view",entity_id={"projectId":view_project,"datasetId":view_dataset,"tableId":view_table},)dataset=self.get_dataset(project_id=project_id,dataset_id=source_dataset)# Check to see if the view we want to add already exists.ifview_accessnotindataset.access_entries:self.log.info("Granting table %s:%s.%s authorized view access to %s:%s dataset.",view_project,view_dataset,view_table,project_id,source_dataset,)dataset.access_entries+=[view_access]dataset=self.update_dataset(fields=["access"],dataset_resource=dataset.to_api_repr(),project_id=project_id)else:self.log.info("Table %s:%s.%s already has authorized view access to %s:%s dataset.",view_project,view_dataset,view_table,project_id,source_dataset,)returndataset.to_api_repr()
@GoogleBaseHook.fallback_to_default_project_id
[docs]defrun_table_upsert(self,dataset_id:str,table_resource:dict[str,Any],project_id:str|None=None)->dict[str,Any]:""" If the table already exists, update the existing table if not create new. Since BigQuery does not natively allow table upserts, this is not an atomic operation. :param dataset_id: the dataset to upsert the table into. :param table_resource: a table resource. see https://cloud.google.com/bigquery/docs/reference/v2/tables#resource :param project_id: the project to upsert the table into. If None, project will be self.project_id. :return: """table_id=table_resource["tableReference"]["tableId"]table_resource=self._resolve_table_reference(table_resource=table_resource,project_id=project_id,dataset_id=dataset_id,table_id=table_id)tables_list_resp=self.get_dataset_tables(dataset_id=dataset_id,project_id=project_id)ifany(table["tableId"]==table_idfortableintables_list_resp):self.log.info("Table %s:%s.%s exists, updating.",project_id,dataset_id,table_id)table=self.update_table(table_resource=table_resource)else:self.log.info("Table %s:%s.%s does not exist. creating.",project_id,dataset_id,table_id)table=self.create_empty_table(table_resource=table_resource,project_id=project_id).to_api_repr()returntable
[docs]defrun_table_delete(self,deletion_dataset_table:str,ignore_if_missing:bool=False)->None:""" Delete an existing table from the dataset; If the table does not exist, return an error unless ignore_if_missing is set to True. :param deletion_dataset_table: A dotted ``(<project>.|<project>:)<dataset>.<table>`` that indicates which table will be deleted. :param ignore_if_missing: if True, then return success even if the requested table does not exist. :return: """warnings.warn("This method is deprecated. Please use `delete_table`.",DeprecationWarning)returnself.delete_table(table_id=deletion_dataset_table,not_found_ok=ignore_if_missing)
@GoogleBaseHook.fallback_to_default_project_id
[docs]defdelete_table(self,table_id:str,not_found_ok:bool=True,project_id:str|None=None,)->None:""" Delete an existing table from the dataset. If the table does not exist, return an error unless not_found_ok is set to True. :param table_id: A dotted ``(<project>.|<project>:)<dataset>.<table>`` that indicates which table will be deleted. :param not_found_ok: if True, then return success even if the requested table does not exist. :param project_id: the project used to perform the request """self.get_client(project_id=project_id).delete_table(table=table_id,not_found_ok=not_found_ok,)self.log.info("Deleted table %s",table_id)
[docs]defget_tabledata(self,dataset_id:str,table_id:str,max_results:int|None=None,selected_fields:str|None=None,page_token:str|None=None,start_index:int|None=None,)->list[dict]:""" Get the data of a given dataset.table and optionally with selected columns. see https://cloud.google.com/bigquery/docs/reference/v2/tabledata/list :param dataset_id: the dataset ID of the requested table. :param table_id: the table ID of the requested table. :param max_results: the maximum results to return. :param selected_fields: List of fields to return (comma-separated). If unspecified, all fields are returned. :param page_token: page token, returned from a previous call, identifying the result set. :param start_index: zero based index of the starting row to read. :return: list of rows """warnings.warn("This method is deprecated. Please use `list_rows`.",DeprecationWarning)rows=self.list_rows(dataset_id=dataset_id,table_id=table_id,max_results=max_results,selected_fields=selected_fields,page_token=page_token,start_index=start_index,)return[dict(r)forrinrows]
@GoogleBaseHook.fallback_to_default_project_id
[docs]deflist_rows(self,dataset_id:str,table_id:str,max_results:int|None=None,selected_fields:list[str]|str|None=None,page_token:str|None=None,start_index:int|None=None,project_id:str|None=None,location:str|None=None,retry:Retry=DEFAULT_RETRY,return_iterator:bool=False,)->list[Row]|RowIterator:""" List the rows of the table. See https://cloud.google.com/bigquery/docs/reference/rest/v2/tabledata/list :param dataset_id: the dataset ID of the requested table. :param table_id: the table ID of the requested table. :param max_results: the maximum results to return. :param selected_fields: List of fields to return (comma-separated). If unspecified, all fields are returned. :param page_token: page token, returned from a previous call, identifying the result set. :param start_index: zero based index of the starting row to read. :param project_id: Project ID for the project which the client acts on behalf of. :param location: Default location for job. :param retry: How to retry the RPC. :param return_iterator: Instead of returning a list[Row], returns a RowIterator which can be used to obtain the next_page_token property. :return: list of rows """location=locationorself.locationifisinstance(selected_fields,str):selected_fields=selected_fields.split(",")ifselected_fields:selected_fields=[SchemaField(n,"")forninselected_fields]else:selected_fields=Nonetable=self._resolve_table_reference(table_resource={},project_id=project_id,dataset_id=dataset_id,table_id=table_id,)iterator=self.get_client(project_id=project_id,location=location).list_rows(table=Table.from_api_repr(table),selected_fields=selected_fields,max_results=max_results,page_token=page_token,start_index=start_index,retry=retry,)ifreturn_iterator:returniteratorreturnlist(iterator)
@GoogleBaseHook.fallback_to_default_project_id
[docs]defget_schema(self,dataset_id:str,table_id:str,project_id:str|None=None)->dict:""" Get the schema for a given dataset and table. see https://cloud.google.com/bigquery/docs/reference/v2/tables#resource :param dataset_id: the dataset ID of the requested table :param table_id: the table ID of the requested table :param project_id: the optional project ID of the requested table. If not provided, the connector's configured project will be used. :return: a table schema """table_ref=TableReference(dataset_ref=DatasetReference(project_id,dataset_id),table_id=table_id)table=self.get_client(project_id=project_id).get_table(table_ref)return{"fields":[s.to_api_repr()forsintable.schema]}
@GoogleBaseHook.fallback_to_default_project_id
[docs]defupdate_table_schema(self,schema_fields_updates:list[dict[str,Any]],include_policy_tags:bool,dataset_id:str,table_id:str,project_id:str|None=None,)->dict[str,Any]:""" Update fields within a schema for a given dataset and table. Note that some fields in schemas are immutable and trying to change them will cause an exception. If a new field is included it will be inserted which requires all required fields to be set. See https://cloud.google.com/bigquery/docs/reference/rest/v2/tables#TableSchema :param include_policy_tags: If set to True policy tags will be included in the update request which requires special permissions even if unchanged see https://cloud.google.com/bigquery/docs/column-level-security#roles :param dataset_id: the dataset ID of the requested table to be updated :param table_id: the table ID of the table to be updated :param schema_fields_updates: a partial schema resource. see https://cloud.google.com/bigquery/docs/reference/rest/v2/tables#TableSchema **Example**: :: schema_fields_updates=[ {"name": "emp_name", "description": "Some New Description"}, {"name": "salary", "description": "Some New Description"}, {"name": "departments", "fields": [ {"name": "name", "description": "Some New Description"}, {"name": "type", "description": "Some New Description"} ]}, ] :param project_id: The name of the project where we want to update the table. """def_build_new_schema(current_schema:list[dict[str,Any]],schema_fields_updates:list[dict[str,Any]])->list[dict[str,Any]]:# Turn schema_field_updates into a dict keyed on field namesschema_fields_updates_dict={field["name"]:fieldforfieldindeepcopy(schema_fields_updates)}# Create a new dict for storing the new schema, initiated based on the current_schema# as of Python 3.6, dicts retain order.new_schema={field["name"]:fieldforfieldindeepcopy(current_schema)}# Each item in schema_fields_updates contains a potential patch# to a schema field, iterate over themforfield_name,patched_valueinschema_fields_updates_dict.items():# If this field already exists, update itiffield_nameinnew_schema:# If this field is of type RECORD and has a fields key we need to patch it recursivelyif"fields"inpatched_value:patched_value["fields"]=_build_new_schema(new_schema[field_name]["fields"],patched_value["fields"])# Update the new_schema with the patched valuenew_schema[field_name].update(patched_value)# This is a new field, just include the whole configuration for itelse:new_schema[field_name]=patched_valuereturnlist(new_schema.values())def_remove_policy_tags(schema:list[dict[str,Any]]):forfieldinschema:if"policyTags"infield:delfield["policyTags"]if"fields"infield:_remove_policy_tags(field["fields"])current_table_schema=self.get_schema(dataset_id=dataset_id,table_id=table_id,project_id=project_id)["fields"]new_schema=_build_new_schema(current_table_schema,schema_fields_updates)ifnotinclude_policy_tags:_remove_policy_tags(new_schema)table=self.update_table(table_resource={"schema":{"fields":new_schema}},fields=["schema"],project_id=project_id,dataset_id=dataset_id,table_id=table_id,)returntable
@GoogleBaseHook.fallback_to_default_project_id
[docs]defpoll_job_complete(self,job_id:str,project_id:str|None=None,location:str|None=None,retry:Retry=DEFAULT_RETRY,)->bool:""" Check if jobs completed. :param job_id: id of the job. :param project_id: Google Cloud Project where the job is running :param location: location the job is running :param retry: How to retry the RPC. """location=locationorself.locationjob=self.get_client(project_id=project_id,location=location).get_job(job_id=job_id)returnjob.done(retry=retry)
[docs]defcancel_query(self)->None:"""Cancel all started queries that have not yet completed"""warnings.warn("This method is deprecated. Please use `BigQueryHook.cancel_job`.",DeprecationWarning,)ifself.running_job_id:self.cancel_job(job_id=self.running_job_id)else:self.log.info("No running BigQuery jobs to cancel.")
@GoogleBaseHook.fallback_to_default_project_id
[docs]defcancel_job(self,job_id:str,project_id:str|None=None,location:str|None=None,)->None:""" Cancel a job and wait for cancellation to complete :param job_id: id of the job. :param project_id: Google Cloud Project where the job is running :param location: location the job is running """project_id=project_idorself.project_idlocation=locationorself.locationifself.poll_job_complete(job_id=job_id,project_id=project_id,location=location):self.log.info("No running BigQuery jobs to cancel.")returnself.log.info("Attempting to cancel job : %s, %s",project_id,job_id)self.get_client(location=location,project_id=project_id).cancel_job(job_id=job_id)# Wait for all the calls to cancel to finishmax_polling_attempts=12polling_attempts=0job_complete=Falsewhilepolling_attempts<max_polling_attemptsandnotjob_complete:polling_attempts+=1job_complete=self.poll_job_complete(job_id=job_id,project_id=project_id,location=location)ifjob_complete:self.log.info("Job successfully canceled: %s, %s",project_id,job_id)elifpolling_attempts==max_polling_attempts:self.log.info("Stopping polling due to timeout. Job %s, %s ""has not completed cancel and may or may not finish.",project_id,job_id,)else:self.log.info("Waiting for canceled job %s, %s to finish.",project_id,job_id)time.sleep(5)
@GoogleBaseHook.fallback_to_default_project_id
[docs]defget_job(self,job_id:str|None=None,project_id:str|None=None,location:str|None=None,)->CopyJob|QueryJob|LoadJob|ExtractJob:""" Retrieves a BigQuery job. For more information see: https://cloud.google.com/bigquery/docs/reference/v2/jobs :param job_id: The ID of the job. The ID must contain only letters (a-z, A-Z), numbers (0-9), underscores (_), or dashes (-). The maximum length is 1,024 characters. If not provided then uuid will be generated. :param project_id: Google Cloud Project where the job is running :param location: location the job is running """client=self.get_client(project_id=project_id,location=location)job=client.get_job(job_id=job_id,project=project_id,location=location)returnjob
[docs]definsert_job(self,configuration:dict,job_id:str|None=None,project_id:str|None=None,location:str|None=None,nowait:bool=False,retry:Retry=DEFAULT_RETRY,timeout:float|None=None,)->BigQueryJob:""" Executes a BigQuery job. Waits for the job to complete and returns job id. See here: https://cloud.google.com/bigquery/docs/reference/v2/jobs :param configuration: The configuration parameter maps directly to BigQuery's configuration field in the job object. See https://cloud.google.com/bigquery/docs/reference/v2/jobs for details. :param job_id: The ID of the job. The ID must contain only letters (a-z, A-Z), numbers (0-9), underscores (_), or dashes (-). The maximum length is 1,024 characters. If not provided then uuid will be generated. :param project_id: Google Cloud Project where the job is running :param location: location the job is running :param nowait: specify whether to insert job without waiting for the result :param retry: How to retry the RPC. :param timeout: The number of seconds to wait for the underlying HTTP transport before using ``retry``. """location=locationorself.locationjob_id=job_idorself._custom_job_id(configuration)client=self.get_client(project_id=project_id,location=location)job_data={"configuration":configuration,"jobReference":{"jobId":job_id,"projectId":project_id,"location":location},}supported_jobs={LoadJob._JOB_TYPE:LoadJob,CopyJob._JOB_TYPE:CopyJob,ExtractJob._JOB_TYPE:ExtractJob,QueryJob._JOB_TYPE:QueryJob,}job=Noneforjob_type,job_objectinsupported_jobs.items():ifjob_typeinconfiguration:job=job_objectbreakifnotjob:raiseAirflowException(f"Unknown job type. Supported types: {supported_jobs.keys()}")job=job.from_api_repr(job_data,client)self.log.info("Inserting job %s",job.job_id)ifnowait:# Initiate the job and don't wait for it to complete.job._begin()else:# Start the job and wait for it to complete and get the result.job.result(timeout=timeout,retry=retry)returnjob
[docs]defrun_with_configuration(self,configuration:dict)->str:""" Executes a BigQuery SQL query. See here: https://cloud.google.com/bigquery/docs/reference/v2/jobs For more details about the configuration parameter. :param configuration: The configuration parameter maps directly to BigQuery's configuration field in the job object. See https://cloud.google.com/bigquery/docs/reference/v2/jobs for details. """warnings.warn("This method is deprecated. Please use `BigQueryHook.insert_job`",DeprecationWarning)job=self.insert_job(configuration=configuration,project_id=self.project_id)self.running_job_id=job.job_idreturnjob.job_id
[docs]defrun_load(self,destination_project_dataset_table:str,source_uris:list,schema_fields:list|None=None,source_format:str="CSV",create_disposition:str="CREATE_IF_NEEDED",skip_leading_rows:int=0,write_disposition:str="WRITE_EMPTY",field_delimiter:str=",",max_bad_records:int=0,quote_character:str|None=None,ignore_unknown_values:bool=False,allow_quoted_newlines:bool=False,allow_jagged_rows:bool=False,encoding:str="UTF-8",schema_update_options:Iterable|None=None,src_fmt_configs:dict|None=None,time_partitioning:dict|None=None,cluster_fields:list|None=None,autodetect:bool=False,encryption_configuration:dict|None=None,labels:dict|None=None,description:str|None=None,)->str:""" Executes a BigQuery load command to load data from Google Cloud Storage to BigQuery. See here: https://cloud.google.com/bigquery/docs/reference/v2/jobs For more details about these parameters. :param destination_project_dataset_table: The dotted ``(<project>.|<project>:)<dataset>.<table>($<partition>)`` BigQuery table to load data into. If ``<project>`` is not included, project will be the project defined in the connection json. If a partition is specified the operator will automatically append the data, create a new partition or create a new DAY partitioned table. :param schema_fields: The schema field list as defined here: https://cloud.google.com/bigquery/docs/reference/v2/jobs#configuration.load Required if autodetect=False; optional if autodetect=True. :param autodetect: Attempt to autodetect the schema for CSV and JSON source files. :param source_uris: The source Google Cloud Storage URI (e.g. gs://some-bucket/some-file.txt). A single wild per-object name can be used. :param source_format: File format to export. :param create_disposition: The create disposition if the table doesn't exist. :param skip_leading_rows: Number of rows to skip when loading from a CSV. :param write_disposition: The write disposition if the table already exists. :param field_delimiter: The delimiter to use when loading from a CSV. :param max_bad_records: The maximum number of bad records that BigQuery can ignore when running the job. :param quote_character: The value that is used to quote data sections in a CSV file. :param ignore_unknown_values: [Optional] Indicates if BigQuery should allow extra values that are not represented in the table schema. If true, the extra values are ignored. If false, records with extra columns are treated as bad records, and if there are too many bad records, an invalid error is returned in the job result. :param allow_quoted_newlines: Whether to allow quoted newlines (true) or not (false). :param allow_jagged_rows: Accept rows that are missing trailing optional columns. The missing values are treated as nulls. If false, records with missing trailing columns are treated as bad records, and if there are too many bad records, an invalid error is returned in the job result. Only applicable when source_format is CSV. :param encoding: The character encoding of the data. .. seealso:: https://cloud.google.com/bigquery/docs/reference/rest/v2/tables#externalDataConfiguration.csvOptions.encoding :param schema_update_options: Allows the schema of the destination table to be updated as a side effect of the load job. :param src_fmt_configs: configure optional fields specific to the source format :param time_partitioning: configure optional time partitioning fields i.e. partition by field, type and expiration as per API specifications. :param cluster_fields: Request that the result of this load be stored sorted by one or more columns. BigQuery supports clustering for both partitioned and non-partitioned tables. The order of columns given determines the sort order. :param encryption_configuration: [Optional] Custom encryption configuration (e.g., Cloud KMS keys). **Example**: :: encryption_configuration = { "kmsKeyName": "projects/testp/locations/us/keyRings/test-kr/cryptoKeys/test-key" } :param labels: A dictionary containing labels for the BiqQuery table. :param description: A string containing the description for the BigQuery table. """warnings.warn("This method is deprecated. Please use `BigQueryHook.insert_job` method.",DeprecationWarning)ifnotself.project_id:raiseValueError("The project_id should be set")# To provide backward compatibilityschema_update_options=list(schema_update_optionsor[])# bigquery only allows certain source formats# we check to make sure the passed source format is valid# if it's not, we raise a ValueError# Refer to this link for more details:# https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs#configuration.query.tableDefinitions.(key).sourceFormat # noqaifschema_fieldsisNoneandnotautodetect:raiseValueError("You must either pass a schema or autodetect=True.")ifsrc_fmt_configsisNone:src_fmt_configs={}source_format=source_format.upper()allowed_formats=["CSV","NEWLINE_DELIMITED_JSON","AVRO","GOOGLE_SHEETS","DATASTORE_BACKUP","PARQUET",]ifsource_formatnotinallowed_formats:raiseValueError(f"{source_format} is not a valid source format. "f"Please use one of the following types: {allowed_formats}.")# bigquery also allows you to define how you want a table's schema to change# as a side effect of a load# for more details:# https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs#configuration.load.schemaUpdateOptionsallowed_schema_update_options=["ALLOW_FIELD_ADDITION","ALLOW_FIELD_RELAXATION"]ifnotset(allowed_schema_update_options).issuperset(set(schema_update_options)):raiseValueError(f"{schema_update_options} contains invalid schema update options. "f"Please only use one or more of the following options: {allowed_schema_update_options}")destination_project,destination_dataset,destination_table=self.split_tablename(table_input=destination_project_dataset_table,default_project_id=self.project_id,var_name="destination_project_dataset_table",)configuration:dict[str,Any]={"load":{"autodetect":autodetect,"createDisposition":create_disposition,"destinationTable":{"projectId":destination_project,"datasetId":destination_dataset,"tableId":destination_table,},"sourceFormat":source_format,"sourceUris":source_uris,"writeDisposition":write_disposition,"ignoreUnknownValues":ignore_unknown_values,}}time_partitioning=_cleanse_time_partitioning(destination_project_dataset_table,time_partitioning)iftime_partitioning:configuration["load"].update({"timePartitioning":time_partitioning})ifcluster_fields:configuration["load"].update({"clustering":{"fields":cluster_fields}})ifschema_fields:configuration["load"]["schema"]={"fields":schema_fields}ifschema_update_options:ifwrite_dispositionnotin["WRITE_APPEND","WRITE_TRUNCATE"]:raiseValueError("schema_update_options is only ""allowed if write_disposition is ""'WRITE_APPEND' or 'WRITE_TRUNCATE'.")else:self.log.info("Adding experimental 'schemaUpdateOptions': %s",schema_update_options)configuration["load"]["schemaUpdateOptions"]=schema_update_optionsifmax_bad_records:configuration["load"]["maxBadRecords"]=max_bad_recordsifencryption_configuration:configuration["load"]["destinationEncryptionConfiguration"]=encryption_configurationiflabelsordescription:configuration["load"].update({"destinationTableProperties":{}})iflabels:configuration["load"]["destinationTableProperties"]["labels"]=labelsifdescription:configuration["load"]["destinationTableProperties"]["description"]=descriptionsrc_fmt_to_configs_mapping={"CSV":["allowJaggedRows","allowQuotedNewlines","autodetect","fieldDelimiter","skipLeadingRows","ignoreUnknownValues","nullMarker","quote","encoding","preserveAsciiControlCharacters",],"DATASTORE_BACKUP":["projectionFields"],"NEWLINE_DELIMITED_JSON":["autodetect","ignoreUnknownValues"],"PARQUET":["autodetect","ignoreUnknownValues"],"AVRO":["useAvroLogicalTypes"],}valid_configs=src_fmt_to_configs_mapping[source_format]# if following fields are not specified in src_fmt_configs,# honor the top-level params for backward-compatibilitybackward_compatibility_configs={"skipLeadingRows":skip_leading_rows,"fieldDelimiter":field_delimiter,"ignoreUnknownValues":ignore_unknown_values,"quote":quote_character,"allowQuotedNewlines":allow_quoted_newlines,"encoding":encoding,}src_fmt_configs=_validate_src_fmt_configs(source_format,src_fmt_configs,valid_configs,backward_compatibility_configs)configuration["load"].update(src_fmt_configs)ifallow_jagged_rows:configuration["load"]["allowJaggedRows"]=allow_jagged_rowsjob=self.insert_job(configuration=configuration,project_id=self.project_id)self.running_job_id=job.job_idreturnjob.job_id
[docs]defrun_copy(self,source_project_dataset_tables:list|str,destination_project_dataset_table:str,write_disposition:str="WRITE_EMPTY",create_disposition:str="CREATE_IF_NEEDED",labels:dict|None=None,encryption_configuration:dict|None=None,)->str:""" Executes a BigQuery copy command to copy data from one BigQuery table to another. See here: https://cloud.google.com/bigquery/docs/reference/v2/jobs#configuration.copy For more details about these parameters. :param source_project_dataset_tables: One or more dotted ``(project:|project.)<dataset>.<table>`` BigQuery tables to use as the source data. Use a list if there are multiple source tables. If ``<project>`` is not included, project will be the project defined in the connection json. :param destination_project_dataset_table: The destination BigQuery table. Format is: ``(project:|project.)<dataset>.<table>`` :param write_disposition: The write disposition if the table already exists. :param create_disposition: The create disposition if the table doesn't exist. :param labels: a dictionary containing labels for the job/query, passed to BigQuery :param encryption_configuration: [Optional] Custom encryption configuration (e.g., Cloud KMS keys). **Example**: :: encryption_configuration = { "kmsKeyName": "projects/testp/locations/us/keyRings/test-kr/cryptoKeys/test-key" } """warnings.warn("This method is deprecated. Please use `BigQueryHook.insert_job` method.",DeprecationWarning)ifnotself.project_id:raiseValueError("The project_id should be set")source_project_dataset_tables=([source_project_dataset_tables]ifnotisinstance(source_project_dataset_tables,list)elsesource_project_dataset_tables)source_project_dataset_tables_fixup=[]forsource_project_dataset_tableinsource_project_dataset_tables:source_project,source_dataset,source_table=self.split_tablename(table_input=source_project_dataset_table,default_project_id=self.project_id,var_name="source_project_dataset_table",)source_project_dataset_tables_fixup.append({"projectId":source_project,"datasetId":source_dataset,"tableId":source_table})destination_project,destination_dataset,destination_table=self.split_tablename(table_input=destination_project_dataset_table,default_project_id=self.project_id)configuration={"copy":{"createDisposition":create_disposition,"writeDisposition":write_disposition,"sourceTables":source_project_dataset_tables_fixup,"destinationTable":{"projectId":destination_project,"datasetId":destination_dataset,"tableId":destination_table,},}}iflabels:configuration["labels"]=labelsifencryption_configuration:configuration["copy"]["destinationEncryptionConfiguration"]=encryption_configurationjob=self.insert_job(configuration=configuration,project_id=self.project_id)self.running_job_id=job.job_idreturnjob.job_id
[docs]defrun_extract(self,source_project_dataset_table:str,destination_cloud_storage_uris:list[str],compression:str="NONE",export_format:str="CSV",field_delimiter:str=",",print_header:bool=True,labels:dict|None=None,return_full_job:bool=False,)->str|BigQueryJob:""" Executes a BigQuery extract command to copy data from BigQuery to Google Cloud Storage. See here: https://cloud.google.com/bigquery/docs/reference/v2/jobs For more details about these parameters. :param source_project_dataset_table: The dotted ``<dataset>.<table>`` BigQuery table to use as the source data. :param destination_cloud_storage_uris: The destination Google Cloud Storage URI (e.g. gs://some-bucket/some-file.txt). Follows convention defined here: https://cloud.google.com/bigquery/exporting-data-from-bigquery#exportingmultiple :param compression: Type of compression to use. :param export_format: File format to export. :param field_delimiter: The delimiter to use when extracting to a CSV. :param print_header: Whether to print a header for a CSV file extract. :param labels: a dictionary containing labels for the job/query, passed to BigQuery :param return_full_job: return full job instead of job id only """warnings.warn("This method is deprecated. Please use `BigQueryHook.insert_job` method.",DeprecationWarning)ifnotself.project_id:raiseValueError("The project_id should be set")source_project,source_dataset,source_table=self.split_tablename(table_input=source_project_dataset_table,default_project_id=self.project_id,var_name="source_project_dataset_table",)configuration:dict[str,Any]={"extract":{"sourceTable":{"projectId":source_project,"datasetId":source_dataset,"tableId":source_table,},"compression":compression,"destinationUris":destination_cloud_storage_uris,"destinationFormat":export_format,}}iflabels:configuration["labels"]=labelsifexport_format=="CSV":# Only set fieldDelimiter and printHeader fields if using CSV.# Google does not like it if you set these fields for other export# formats.configuration["extract"]["fieldDelimiter"]=field_delimiterconfiguration["extract"]["printHeader"]=print_headerjob=self.insert_job(configuration=configuration,project_id=self.project_id)self.running_job_id=job.job_idifreturn_full_job:returnjobreturnjob.job_id
[docs]defrun_query(self,sql:str,destination_dataset_table:str|None=None,write_disposition:str="WRITE_EMPTY",allow_large_results:bool=False,flatten_results:bool|None=None,udf_config:list|None=None,use_legacy_sql:bool|None=None,maximum_billing_tier:int|None=None,maximum_bytes_billed:float|None=None,create_disposition:str="CREATE_IF_NEEDED",query_params:list|None=None,labels:dict|None=None,schema_update_options:Iterable|None=None,priority:str="INTERACTIVE",time_partitioning:dict|None=None,api_resource_configs:dict|None=None,cluster_fields:list[str]|None=None,location:str|None=None,encryption_configuration:dict|None=None,)->str:""" Executes a BigQuery SQL query. Optionally persists results in a BigQuery table. See here: https://cloud.google.com/bigquery/docs/reference/v2/jobs For more details about these parameters. :param sql: The BigQuery SQL to execute. :param destination_dataset_table: The dotted ``<dataset>.<table>`` BigQuery table to save the query results. :param write_disposition: What to do if the table already exists in BigQuery. :param allow_large_results: Whether to allow large results. :param flatten_results: If true and query uses legacy SQL dialect, flattens all nested and repeated fields in the query results. ``allowLargeResults`` must be true if this is set to false. For standard SQL queries, this flag is ignored and results are never flattened. :param udf_config: The User Defined Function configuration for the query. See https://cloud.google.com/bigquery/user-defined-functions for details. :param use_legacy_sql: Whether to use legacy SQL (true) or standard SQL (false). If `None`, defaults to `self.use_legacy_sql`. :param api_resource_configs: a dictionary that contain params 'configuration' applied for Google BigQuery Jobs API: https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs for example, {'query': {'useQueryCache': False}}. You could use it if you need to provide some params that are not supported by the BigQueryHook like args. :param maximum_billing_tier: Positive integer that serves as a multiplier of the basic price. :param maximum_bytes_billed: Limits the bytes billed for this job. Queries that will have bytes billed beyond this limit will fail (without incurring a charge). If unspecified, this will be set to your project default. :param create_disposition: Specifies whether the job is allowed to create new tables. :param query_params: a list of dictionary containing query parameter types and values, passed to BigQuery :param labels: a dictionary containing labels for the job/query, passed to BigQuery :param schema_update_options: Allows the schema of the destination table to be updated as a side effect of the query job. :param priority: Specifies a priority for the query. Possible values include INTERACTIVE and BATCH. The default value is INTERACTIVE. :param time_partitioning: configure optional time partitioning fields i.e. partition by field, type and expiration as per API specifications. :param cluster_fields: Request that the result of this query be stored sorted by one or more columns. BigQuery supports clustering for both partitioned and non-partitioned tables. The order of columns given determines the sort order. :param location: The geographic location of the job. Required except for US and EU. See details at https://cloud.google.com/bigquery/docs/locations#specifying_your_location :param encryption_configuration: [Optional] Custom encryption configuration (e.g., Cloud KMS keys). **Example**: :: encryption_configuration = { "kmsKeyName": "projects/testp/locations/us/keyRings/test-kr/cryptoKeys/test-key" } """warnings.warn("This method is deprecated. Please use `BigQueryHook.insert_job` method.",DeprecationWarning)ifnotself.project_id:raiseValueError("The project_id should be set")labels=labelsorself.labelsschema_update_options=list(schema_update_optionsor[])iftime_partitioningisNone:time_partitioning={}iflocation:self.location=locationifnotapi_resource_configs:api_resource_configs=self.api_resource_configselse:_validate_value("api_resource_configs",api_resource_configs,dict)configuration=deepcopy(api_resource_configs)if"query"notinconfiguration:configuration["query"]={}else:_validate_value("api_resource_configs['query']",configuration["query"],dict)ifsqlisNoneandnotconfiguration["query"].get("query",None):raiseTypeError("`BigQueryBaseCursor.run_query` missing 1 required positional argument: `sql`")# BigQuery also allows you to define how you want a table's schema to change# as a side effect of a query job# for more details:# https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs#configuration.query.schemaUpdateOptions # noqaallowed_schema_update_options=["ALLOW_FIELD_ADDITION","ALLOW_FIELD_RELAXATION"]ifnotset(allowed_schema_update_options).issuperset(set(schema_update_options)):raiseValueError(f"{schema_update_options} contains invalid schema update options."f" Please only use one or more of the following options: {allowed_schema_update_options}")ifschema_update_options:ifwrite_dispositionnotin["WRITE_APPEND","WRITE_TRUNCATE"]:raiseValueError("schema_update_options is only ""allowed if write_disposition is ""'WRITE_APPEND' or 'WRITE_TRUNCATE'.")ifdestination_dataset_table:destination_project,destination_dataset,destination_table=self.split_tablename(table_input=destination_dataset_table,default_project_id=self.project_id)destination_dataset_table={# type: ignore"projectId":destination_project,"datasetId":destination_dataset,"tableId":destination_table,}ifcluster_fields:cluster_fields={"fields":cluster_fields}# type: ignorequery_param_list:list[tuple[Any,str,str|bool|None|dict,type|tuple[type]]]=[(sql,"query",None,(str,)),(priority,"priority","INTERACTIVE",(str,)),(use_legacy_sql,"useLegacySql",self.use_legacy_sql,bool),(query_params,"queryParameters",None,list),(udf_config,"userDefinedFunctionResources",None,list),(maximum_billing_tier,"maximumBillingTier",None,int),(maximum_bytes_billed,"maximumBytesBilled",None,float),(time_partitioning,"timePartitioning",{},dict),(schema_update_options,"schemaUpdateOptions",None,list),(destination_dataset_table,"destinationTable",None,dict),(cluster_fields,"clustering",None,dict),]forparam,param_name,param_default,param_typeinquery_param_list:ifparam_namenotinconfiguration["query"]andparamin[None,{},()]:ifparam_name=="timePartitioning":param_default=_cleanse_time_partitioning(destination_dataset_table,time_partitioning)param=param_defaultifparamin[None,{},()]:continue_api_resource_configs_duplication_check(param_name,param,configuration["query"])configuration["query"][param_name]=param# check valid type of provided param,# it last step because we can get param from 2 sources,# and first of all need to find it_validate_value(param_name,configuration["query"][param_name],param_type)ifparam_name=="schemaUpdateOptions"andparam:self.log.info("Adding experimental 'schemaUpdateOptions': %s",schema_update_options)ifparam_name!="destinationTable":continueforkeyin["projectId","datasetId","tableId"]:ifkeynotinconfiguration["query"]["destinationTable"]:raiseValueError("Not correct 'destinationTable' in ""api_resource_configs. 'destinationTable' ""must be a dict with {'projectId':'', ""'datasetId':'', 'tableId':''}")configuration["query"].update({"allowLargeResults":allow_large_results,"flattenResults":flatten_results,"writeDisposition":write_disposition,"createDisposition":create_disposition,})if("useLegacySql"inconfiguration["query"]andconfiguration["query"]["useLegacySql"]and"queryParameters"inconfiguration["query"]):raiseValueError("Query parameters are not allowed when using legacy SQL")iflabels:_api_resource_configs_duplication_check("labels",labels,configuration)configuration["labels"]=labelsifencryption_configuration:configuration["query"]["destinationEncryptionConfiguration"]=encryption_configurationjob=self.insert_job(configuration=configuration,project_id=self.project_id)self.running_job_id=job.job_idreturnjob.job_id
[docs]defsplit_tablename(self,table_input:str,default_project_id:str,var_name:str|None=None)->tuple[str,str,str]:if"."notintable_input:raiseValueError(f"Expected table name in the format of <dataset>.<table>. Got: {table_input}")ifnotdefault_project_id:raiseValueError("INTERNAL: No default project is specified")defvar_print(var_name):ifvar_nameisNone:return""else:returnf"Format exception for {var_name}: "iftable_input.count(".")+table_input.count(":")>3:raiseException(f"{var_print(var_name)}Use either : or . to specify project got {table_input}")cmpt=table_input.rsplit(":",1)project_id=Nonerest=table_inputiflen(cmpt)==1:project_id=Nonerest=cmpt[0]eliflen(cmpt)==2andcmpt[0].count(":")<=1:ifcmpt[-1].count(".")!=2:project_id=cmpt[0]rest=cmpt[1]else:raiseException(f"{var_print(var_name)}Expect format of (<project:)<dataset>.<table>, got {table_input}")cmpt=rest.split(".")iflen(cmpt)==3:ifproject_id:raiseValueError(f"{var_print(var_name)}Use either : or . to specify project")project_id=cmpt[0]dataset_id=cmpt[1]table_id=cmpt[2]eliflen(cmpt)==2:dataset_id=cmpt[0]table_id=cmpt[1]else:raiseException(f"{var_print(var_name)} Expect format of (<project.|<project:)<dataset>.<table>, "f"got {table_input}")ifproject_idisNone:ifvar_nameisnotNone:self.log.info('Project is not included in %s: %s; using project "%s"',var_name,table_input,default_project_id,)project_id=default_project_idreturnproject_id,dataset_id,table_id
[docs]classBigQueryConnection:""" BigQuery does not have a notion of a persistent connection. Thus, these objects are small stateless factories for cursors, which do all the real work. """def__init__(self,*args,**kwargs)->None:self._args=argsself._kwargs=kwargs
[docs]defclose(self)->None:"""The BigQueryConnection does not have anything to close"""
[docs]defcommit(self)->None:"""The BigQueryConnection does not support transactions"""
[docs]defcursor(self)->BigQueryCursor:"""Return a new :py:class:`Cursor` object using the connection"""returnBigQueryCursor(*self._args,**self._kwargs)
[docs]defrollback(self)->NoReturn:"""The BigQueryConnection does not have transactions"""raiseNotImplementedError("BigQueryConnection does not have transactions")
[docs]classBigQueryBaseCursor(LoggingMixin):""" The BigQuery base cursor contains helper methods to execute queries against BigQuery. The methods can be used directly by operators, in cases where a PEP 249 cursor isn't needed. """def__init__(self,service:Any,project_id:str,hook:BigQueryHook,use_legacy_sql:bool=True,api_resource_configs:dict|None=None,location:str|None=None,num_retries:int=5,labels:dict|None=None,)->None:super().__init__()self.service=serviceself.project_id=project_idself.use_legacy_sql=use_legacy_sqlifapi_resource_configs:_validate_value("api_resource_configs",api_resource_configs,dict)self.api_resource_configs:dict=api_resource_configsifapi_resource_configselse{}self.running_job_id:str|None=Noneself.location=locationself.num_retries=num_retriesself.labels=labelsself.hook=hook
[docs]defcreate_empty_table(self,*args,**kwargs):""" This method is deprecated. Please use `airflow.providers.google.cloud.hooks.bigquery.BigQueryHook.create_empty_table` """warnings.warn("This method is deprecated. ""Please use `airflow.providers.google.cloud.hooks.bigquery.BigQueryHook.create_empty_table`",DeprecationWarning,stacklevel=3,)returnself.hook.create_empty_table(*args,**kwargs)
[docs]defcreate_empty_dataset(self,*args,**kwargs)->dict[str,Any]:""" This method is deprecated. Please use `airflow.providers.google.cloud.hooks.bigquery.BigQueryHook.create_empty_dataset` """warnings.warn("This method is deprecated. ""Please use `airflow.providers.google.cloud.hooks.bigquery.BigQueryHook.create_empty_dataset`",DeprecationWarning,stacklevel=3,)returnself.hook.create_empty_dataset(*args,**kwargs)
[docs]defget_dataset_tables(self,*args,**kwargs)->list[dict[str,Any]]:""" This method is deprecated. Please use `airflow.providers.google.cloud.hooks.bigquery.BigQueryHook.get_dataset_tables` """warnings.warn("This method is deprecated. ""Please use `airflow.providers.google.cloud.hooks.bigquery.BigQueryHook.get_dataset_tables`",DeprecationWarning,stacklevel=3,)returnself.hook.get_dataset_tables(*args,**kwargs)
[docs]defdelete_dataset(self,*args,**kwargs)->None:""" This method is deprecated. Please use `airflow.providers.google.cloud.hooks.bigquery.BigQueryHook.delete_dataset` """warnings.warn("This method is deprecated. ""Please use `airflow.providers.google.cloud.hooks.bigquery.BigQueryHook.delete_dataset`",DeprecationWarning,stacklevel=3,)returnself.hook.delete_dataset(*args,**kwargs)
[docs]defcreate_external_table(self,*args,**kwargs):""" This method is deprecated. Please use `airflow.providers.google.cloud.hooks.bigquery.BigQueryHook.create_external_table` """warnings.warn("This method is deprecated. ""Please use `airflow.providers.google.cloud.hooks.bigquery.BigQueryHook.create_external_table`",DeprecationWarning,stacklevel=3,)returnself.hook.create_external_table(*args,**kwargs)
[docs]defpatch_table(self,*args,**kwargs)->None:""" This method is deprecated. Please use `airflow.providers.google.cloud.hooks.bigquery.BigQueryHook.patch_table` """warnings.warn("This method is deprecated. ""Please use `airflow.providers.google.cloud.hooks.bigquery.BigQueryHook.patch_table`",DeprecationWarning,stacklevel=3,)returnself.hook.patch_table(*args,**kwargs)
[docs]definsert_all(self,*args,**kwargs)->None:""" This method is deprecated. Please use `airflow.providers.google.cloud.hooks.bigquery.BigQueryHook.insert_all` """warnings.warn("This method is deprecated. ""Please use `airflow.providers.google.cloud.hooks.bigquery.BigQueryHook.insert_all`",DeprecationWarning,stacklevel=3,)returnself.hook.insert_all(*args,**kwargs)
[docs]defupdate_dataset(self,*args,**kwargs)->dict:""" This method is deprecated. Please use `airflow.providers.google.cloud.hooks.bigquery.BigQueryHook.update_dataset` """warnings.warn("This method is deprecated. ""Please use `airflow.providers.google.cloud.hooks.bigquery.BigQueryHook.update_dataset`",DeprecationWarning,stacklevel=3,)returnDataset.to_api_repr(self.hook.update_dataset(*args,**kwargs))
[docs]defpatch_dataset(self,*args,**kwargs)->dict:""" This method is deprecated. Please use `airflow.providers.google.cloud.hooks.bigquery.BigQueryHook.patch_dataset` """warnings.warn("This method is deprecated. ""Please use `airflow.providers.google.cloud.hooks.bigquery.BigQueryHook.patch_dataset`",DeprecationWarning,stacklevel=3,)returnself.hook.patch_dataset(*args,**kwargs)
[docs]defget_dataset_tables_list(self,*args,**kwargs)->list[dict[str,Any]]:""" This method is deprecated. Please use `airflow.providers.google.cloud.hooks.bigquery.BigQueryHook.get_dataset_tables_list` """warnings.warn("This method is deprecated. ""Please use `airflow.providers.google.cloud.hooks.bigquery.BigQueryHook.get_dataset_tables_list`",DeprecationWarning,stacklevel=3,)returnself.hook.get_dataset_tables_list(*args,**kwargs)
[docs]defget_datasets_list(self,*args,**kwargs)->list|HTTPIterator:""" This method is deprecated. Please use `airflow.providers.google.cloud.hooks.bigquery.BigQueryHook.get_datasets_list` """warnings.warn("This method is deprecated. ""Please use `airflow.providers.google.cloud.hooks.bigquery.BigQueryHook.get_datasets_list`",DeprecationWarning,stacklevel=3,)returnself.hook.get_datasets_list(*args,**kwargs)
[docs]defget_dataset(self,*args,**kwargs)->dict:""" This method is deprecated. Please use `airflow.providers.google.cloud.hooks.bigquery.BigQueryHook.get_dataset` """warnings.warn("This method is deprecated. ""Please use `airflow.providers.google.cloud.hooks.bigquery.BigQueryHook.get_dataset`",DeprecationWarning,stacklevel=3,)returnself.hook.get_dataset(*args,**kwargs)
[docs]defrun_grant_dataset_view_access(self,*args,**kwargs)->dict:""" This method is deprecated. Please use `airflow.providers.google.cloud.hooks.bigquery.BigQueryHook.run_grant_dataset_view_access` """warnings.warn("This method is deprecated. ""Please use `airflow.providers.google.cloud.hooks"".bigquery.BigQueryHook.run_grant_dataset_view_access`",DeprecationWarning,stacklevel=3,)returnself.hook.run_grant_dataset_view_access(*args,**kwargs)
[docs]defrun_table_upsert(self,*args,**kwargs)->dict:""" This method is deprecated. Please use `airflow.providers.google.cloud.hooks.bigquery.BigQueryHook.run_table_upsert` """warnings.warn("This method is deprecated. ""Please use `airflow.providers.google.cloud.hooks.bigquery.BigQueryHook.run_table_upsert`",DeprecationWarning,stacklevel=3,)returnself.hook.run_table_upsert(*args,**kwargs)
[docs]defrun_table_delete(self,*args,**kwargs)->None:""" This method is deprecated. Please use `airflow.providers.google.cloud.hooks.bigquery.BigQueryHook.run_table_delete` """warnings.warn("This method is deprecated. ""Please use `airflow.providers.google.cloud.hooks.bigquery.BigQueryHook.run_table_delete`",DeprecationWarning,stacklevel=3,)returnself.hook.run_table_delete(*args,**kwargs)
[docs]defget_tabledata(self,*args,**kwargs)->list[dict]:""" This method is deprecated. Please use `airflow.providers.google.cloud.hooks.bigquery.BigQueryHook.get_tabledata` """warnings.warn("This method is deprecated. ""Please use `airflow.providers.google.cloud.hooks.bigquery.BigQueryHook.get_tabledata`",DeprecationWarning,stacklevel=3,)returnself.hook.get_tabledata(*args,**kwargs)
[docs]defget_schema(self,*args,**kwargs)->dict:""" This method is deprecated. Please use `airflow.providers.google.cloud.hooks.bigquery.BigQueryHook.get_schema` """warnings.warn("This method is deprecated. ""Please use `airflow.providers.google.cloud.hooks.bigquery.BigQueryHook.get_schema`",DeprecationWarning,stacklevel=3,)returnself.hook.get_schema(*args,**kwargs)
[docs]defpoll_job_complete(self,*args,**kwargs)->bool:""" This method is deprecated. Please use `airflow.providers.google.cloud.hooks.bigquery.BigQueryHook.poll_job_complete` """warnings.warn("This method is deprecated. ""Please use `airflow.providers.google.cloud.hooks.bigquery.BigQueryHook.poll_job_complete`",DeprecationWarning,stacklevel=3,)returnself.hook.poll_job_complete(*args,**kwargs)
[docs]defcancel_query(self,*args,**kwargs)->None:""" This method is deprecated. Please use `airflow.providers.google.cloud.hooks.bigquery.BigQueryHook.cancel_query` """warnings.warn("This method is deprecated. ""Please use `airflow.providers.google.cloud.hooks.bigquery.BigQueryHook.cancel_query`",DeprecationWarning,stacklevel=3,)returnself.hook.cancel_query(*args,**kwargs)# type: ignore
[docs]defrun_with_configuration(self,*args,**kwargs)->str:""" This method is deprecated. Please use `airflow.providers.google.cloud.hooks.bigquery.BigQueryHook.run_with_configuration` """warnings.warn("This method is deprecated. ""Please use `airflow.providers.google.cloud.hooks.bigquery.BigQueryHook.run_with_configuration`",DeprecationWarning,stacklevel=3,)returnself.hook.run_with_configuration(*args,**kwargs)
[docs]defrun_load(self,*args,**kwargs)->str:""" This method is deprecated. Please use `airflow.providers.google.cloud.hooks.bigquery.BigQueryHook.run_load` """warnings.warn("This method is deprecated. ""Please use `airflow.providers.google.cloud.hooks.bigquery.BigQueryHook.run_load`",DeprecationWarning,stacklevel=3,)returnself.hook.run_load(*args,**kwargs)
[docs]defrun_copy(self,*args,**kwargs)->str:""" This method is deprecated. Please use `airflow.providers.google.cloud.hooks.bigquery.BigQueryHook.run_copy` """warnings.warn("This method is deprecated. ""Please use `airflow.providers.google.cloud.hooks.bigquery.BigQueryHook.run_copy`",DeprecationWarning,stacklevel=3,)returnself.hook.run_copy(*args,**kwargs)
[docs]defrun_extract(self,*args,**kwargs)->str:""" This method is deprecated. Please use `airflow.providers.google.cloud.hooks.bigquery.BigQueryHook.run_extract` """warnings.warn("This method is deprecated. ""Please use `airflow.providers.google.cloud.hooks.bigquery.BigQueryHook.run_extract`",DeprecationWarning,stacklevel=3,)returnself.hook.run_extract(*args,**kwargs)
[docs]defrun_query(self,*args,**kwargs)->str:""" This method is deprecated. Please use `airflow.providers.google.cloud.hooks.bigquery.BigQueryHook.run_query` """warnings.warn("This method is deprecated. ""Please use `airflow.providers.google.cloud.hooks.bigquery.BigQueryHook.run_query`",DeprecationWarning,stacklevel=3,)returnself.hook.run_query(*args,**kwargs)
[docs]classBigQueryCursor(BigQueryBaseCursor):""" A very basic BigQuery PEP 249 cursor implementation. The PyHive PEP 249 implementation was used as a reference: https://github.com/dropbox/PyHive/blob/master/pyhive/presto.py https://github.com/dropbox/PyHive/blob/master/pyhive/common.py """def__init__(self,service:Any,project_id:str,hook:BigQueryHook,use_legacy_sql:bool=True,location:str|None=None,num_retries:int=5,)->None:super().__init__(service=service,project_id=project_id,hook=hook,use_legacy_sql=use_legacy_sql,location=location,num_retries=num_retries,)self.buffersize:int|None=Noneself.page_token:str|None=Noneself.job_id:str|None=Noneself.buffer:list=[]self.all_pages_loaded:bool=Falseself._description:list=[]@property
[docs]defdescription(self)->list:"""Return the cursor description"""returnself._description
[docs]defclose(self)->None:"""By default, do nothing"""
@property
[docs]defrowcount(self)->int:"""By default, return -1 to indicate that this is not supported"""return-1
[docs]defexecute(self,operation:str,parameters:dict|None=None)->None:""" Executes a BigQuery query, and returns the job ID. :param operation: The query to execute. :param parameters: Parameters to substitute into the query. """sql=_bind_parameters(operation,parameters)ifparameterselseoperationself.flush_results()self.job_id=self.hook.run_query(sql)query_results=self._get_query_result()if"schema"inquery_results:self.description=_format_schema_for_description(query_results["schema"])else:self.description=[]
[docs]defexecutemany(self,operation:str,seq_of_parameters:list)->None:""" Execute a BigQuery query multiple times with different parameters. :param operation: The query to execute. :param seq_of_parameters: List of dictionary parameters to substitute into the query. """forparametersinseq_of_parameters:self.execute(operation,parameters)
[docs]defflush_results(self)->None:"""Flush results related cursor attributes"""self.page_token=Noneself.job_id=Noneself.all_pages_loaded=Falseself.buffer=[]
[docs]deffetchone(self)->list|None:"""Fetch the next row of a query result set"""returnself.next()
[docs]defnext(self)->list|None:""" Helper method for fetchone, which returns the next row from a buffer. If the buffer is empty, attempts to paginate through the result set for the next page, and load it into the buffer. """ifnotself.job_id:returnNoneifnotself.buffer:ifself.all_pages_loaded:returnNonequery_results=self._get_query_result()if"rows"inquery_resultsandquery_results["rows"]:self.page_token=query_results.get("pageToken")fields=query_results["schema"]["fields"]col_types=[field["type"]forfieldinfields]rows=query_results["rows"]fordict_rowinrows:typed_row=[bq_cast(vs["v"],col_types[idx])foridx,vsinenumerate(dict_row["f"])]self.buffer.append(typed_row)ifnotself.page_token:self.all_pages_loaded=Trueelse:# Reset all state since we've exhausted the results.self.flush_results()returnNonereturnself.buffer.pop(0)
[docs]deffetchmany(self,size:int|None=None)->list:""" Fetch the next set of rows of a query result, returning a sequence of sequences (e.g. a list of tuples). An empty sequence is returned when no more rows are available. The number of rows to fetch per call is specified by the parameter. If it is not given, the cursor's arraysize determines the number of rows to be fetched. The method should try to fetch as many rows as indicated by the size parameter. If this is not possible due to the specified number of rows not being available, fewer rows may be returned. An :py:class:`~pyhive.exc.Error` (or subclass) exception is raised if the previous call to :py:meth:`execute` did not produce any result set or no call was issued yet. """ifsizeisNone:size=self.arraysizeresult=[]for_inrange(size):one=self.fetchone()ifoneisNone:breakresult.append(one)returnresult
[docs]deffetchall(self)->list[list]:""" Fetch all (remaining) rows of a query result, returning them as a sequence of sequences (e.g. a list of tuples). """result=[]whileTrue:one=self.fetchone()ifoneisNone:breakresult.append(one)returnresult
[docs]defget_arraysize(self)->int:"""Specifies the number of rows to fetch at a time with .fetchmany()"""returnself.buffersizeor1
[docs]defset_arraysize(self,arraysize:int)->None:"""Specifies the number of rows to fetch at a time with .fetchmany()"""self.buffersize=arraysize
[docs]defsetinputsizes(self,sizes:Any)->None:"""Does nothing by default"""
[docs]defsetoutputsize(self,size:Any,column:Any=None)->None:"""Does nothing by default"""
def_get_query_result(self)->dict:"""Get job query results like data, schema, job type..."""query_results=(self.service.jobs().getQueryResults(projectId=self.project_id,jobId=self.job_id,location=self.location,pageToken=self.page_token,).execute(num_retries=self.num_retries))returnquery_results
def_bind_parameters(operation:str,parameters:dict)->str:"""Helper method that binds parameters to a SQL query"""# inspired by MySQL Python Connector (conversion.py)string_parameters={}# type dict[str, str]for(name,value)inparameters.items():ifvalueisNone:string_parameters[name]="NULL"elifisinstance(value,str):string_parameters[name]="'"+_escape(value)+"'"else:string_parameters[name]=str(value)returnoperation%string_parametersdef_escape(s:str)->str:"""Helper method that escapes parameters to a SQL query"""e=se=e.replace("\\","\\\\")e=e.replace("\n","\\n")e=e.replace("\r","\\r")e=e.replace("'","\\'")e=e.replace('"','\\"')returne
[docs]defsplit_tablename(table_input:str,default_project_id:str,var_name:str|None=None)->tuple[str,str,str]:if"."notintable_input:raiseValueError(f"Expected table name in the format of <dataset>.<table>. Got: {table_input}")ifnotdefault_project_id:raiseValueError("INTERNAL: No default project is specified")defvar_print(var_name):ifvar_nameisNone:return""else:returnf"Format exception for {var_name}: "iftable_input.count(".")+table_input.count(":")>3:raiseException(f"{var_print(var_name)}Use either : or . to specify project got {table_input}")cmpt=table_input.rsplit(":",1)project_id=Nonerest=table_inputiflen(cmpt)==1:project_id=Nonerest=cmpt[0]eliflen(cmpt)==2andcmpt[0].count(":")<=1:ifcmpt[-1].count(".")!=2:project_id=cmpt[0]rest=cmpt[1]else:raiseException(f"{var_print(var_name)}Expect format of (<project:)<dataset>.<table>, got {table_input}")cmpt=rest.split(".")iflen(cmpt)==3:ifproject_id:raiseValueError(f"{var_print(var_name)}Use either : or . to specify project")project_id=cmpt[0]dataset_id=cmpt[1]table_id=cmpt[2]eliflen(cmpt)==2:dataset_id=cmpt[0]table_id=cmpt[1]else:raiseException(f"{var_print(var_name)}Expect format of (<project.|<project:)<dataset>.<table>, got {table_input}")ifproject_idisNone:ifvar_nameisnotNone:log.info('Project is not included in %s: %s; using project "%s"',var_name,table_input,default_project_id,)project_id=default_project_idreturnproject_id,dataset_id,table_id
def_cleanse_time_partitioning(destination_dataset_table:str|None,time_partitioning_in:dict|None)->dict:# if it is a partitioned table ($ is in the table name) add partition load optioniftime_partitioning_inisNone:time_partitioning_in={}time_partitioning_out={}ifdestination_dataset_tableand"$"indestination_dataset_table:time_partitioning_out["type"]="DAY"time_partitioning_out.update(time_partitioning_in)returntime_partitioning_outdef_validate_value(key:Any,value:Any,expected_type:type|tuple[type])->None:"""Function to check expected type and raise error if type is not correct"""ifnotisinstance(value,expected_type):raiseTypeError(f"{key} argument must have a type {expected_type} not {type(value)}")def_api_resource_configs_duplication_check(key:Any,value:Any,config_dict:dict,config_dict_name="api_resource_configs")->None:ifkeyinconfig_dictandvalue!=config_dict[key]:raiseValueError("Values of {param_name} param are duplicated. ""{dict_name} contained {param_name} param ""in `query` config and {param_name} was also provided ""with arg to run_query() method. Please remove duplicates.".format(param_name=key,dict_name=config_dict_name))def_validate_src_fmt_configs(source_format:str,src_fmt_configs:dict,valid_configs:list[str],backward_compatibility_configs:dict|None=None,)->dict:""" Validates the given src_fmt_configs against a valid configuration for the source format. Adds the backward compatibility config to the src_fmt_configs. :param source_format: File format to export. :param src_fmt_configs: Configure optional fields specific to the source format. :param valid_configs: Valid configuration specific to the source format :param backward_compatibility_configs: The top-level params for backward-compatibility """ifbackward_compatibility_configsisNone:backward_compatibility_configs={}fork,vinbackward_compatibility_configs.items():ifknotinsrc_fmt_configsandkinvalid_configs:src_fmt_configs[k]=vfork,vinsrc_fmt_configs.items():ifknotinvalid_configs:raiseValueError(f"{k} is not a valid src_fmt_configs for type {source_format}.")returnsrc_fmt_configsdef_format_schema_for_description(schema:dict)->list:""" Reformat the schema to match cursor description standard which is a tuple of 7 elemenbts (name, type, display_size, internal_size, precision, scale, null_ok) """description=[]forfieldinschema["fields"]:mode=field.get("mode","NULLABLE")field_description=(field["name"],field["type"],None,None,None,None,mode=="NULLABLE",)description.append(field_description)returndescription
[docs]classBigQueryAsyncHook(GoogleBaseAsyncHook):"""Uses gcloud-aio library to retrieve Job details"""
[docs]asyncdefget_job_instance(self,project_id:str|None,job_id:str|None,session:ClientSession)->Job:"""Get the specified job resource by job ID and project ID."""withawaitself.service_file_as_context()asf:returnJob(job_id=job_id,project=project_id,service_file=f,session=cast(Session,session))
[docs]asyncdefget_job_status(self,job_id:str|None,project_id:str|None=None,)->str|None:""" Polls for job status asynchronously using gcloud-aio. Note that an OSError is raised when Job results are still pending. Exception means that Job finished with errors """asyncwithClientSession()ass:try:self.log.info("Executing get_job_status...")job_client=awaitself.get_job_instance(project_id,job_id,s)job_status_response=awaitjob_client.result(cast(Session,s))ifjob_status_response:job_status="success"exceptOSError:job_status="pending"exceptExceptionase:self.log.info("Query execution finished with errors...")job_status=str(e)returnjob_status
[docs]asyncdefget_job_output(self,job_id:str|None,project_id:str|None=None,)->dict[str,Any]:"""Get the big query job output for the given job id asynchronously using gcloud-aio."""asyncwithClientSession()assession:self.log.info("Executing get_job_output..")job_client=awaitself.get_job_instance(project_id,job_id,session)job_query_response=awaitjob_client.get_query_results(cast(Session,session))returnjob_query_response
[docs]asyncdefcreate_job_for_partition_get(self,dataset_id:str|None,project_id:str|None=None,):"""Create a new job and get the job_id using gcloud-aio."""asyncwithClientSession()assession:self.log.info("Executing create_job..")job_client=awaitself.get_job_instance(project_id,"",session)query_request={"query":"SELECT partition_id "f"FROM `{project_id}.{dataset_id}.INFORMATION_SCHEMA.PARTITIONS`","useLegacySql":False,}job_query_resp=awaitjob_client.query(query_request,cast(Session,session))returnjob_query_resp["jobReference"]["jobId"]
[docs]defget_records(self,query_results:dict[str,Any])->list[Any]:""" Given the output query response from gcloud-aio bigquery, convert the response to records. :param query_results: the results from a SQL query """buffer=[]if"rows"inquery_resultsandquery_results["rows"]:rows=query_results["rows"]fields=query_results["schema"]["fields"]col_types=[field["type"]forfieldinfields]fordict_rowinrows:typed_row=[bq_cast(vs["v"],col_types[idx])foridx,vsinenumerate(dict_row["f"])]buffer.append(typed_row)returnbuffer
[docs]defvalue_check(self,sql:str,pass_value:Any,records:list[Any],tolerance:float|None=None,)->None:""" Match a single query resulting row and tolerance with pass_value :return: If Match fail, we throw an AirflowException. """ifnotrecords:raiseAirflowException("The query returned None")pass_value_conv=self._convert_to_float_if_possible(pass_value)is_numeric_value_check=isinstance(pass_value_conv,float)tolerance_pct_str=str(tolerance*100)+"%"iftoleranceelseNoneerror_msg=("Test failed.\nPass value:{pass_value_conv}\n""Tolerance:{tolerance_pct_str}\n""Query:\n{sql}\nResults:\n{records!s}").format(pass_value_conv=pass_value_conv,tolerance_pct_str=tolerance_pct_str,sql=sql,records=records,)ifnotis_numeric_value_check:tests=[str(record)==pass_value_convforrecordinrecords]else:try:numeric_records=[float(record)forrecordinrecords]except(ValueError,TypeError):raiseAirflowException(f"Converting a result to float failed.\n{error_msg}")tests=self._get_numeric_matches(numeric_records,pass_value_conv,tolerance)ifnotall(tests):raiseAirflowException(error_msg)
@staticmethoddef_get_numeric_matches(records:list[float],pass_value:Any,tolerance:float|None=None)->list[bool]:""" A helper function to match numeric pass_value, tolerance with records value :param records: List of value to match against :param pass_value: Expected value :param tolerance: Allowed tolerance for match to succeed """iftolerance:return[pass_value*(1-tolerance)<=record<=pass_value*(1+tolerance)forrecordinrecords]return[record==pass_valueforrecordinrecords]@staticmethoddef_convert_to_float_if_possible(s:Any)->Any:""" A small helper function to convert a string to a numeric value if appropriate :param s: the string to be converted """try:returnfloat(s)except(ValueError,TypeError):returns
[docs]definterval_check(self,row1:str|None,row2:str|None,metrics_thresholds:dict[str,Any],ignore_zero:bool,ratio_formula:str,)->None:""" Checks that the values of metrics given as SQL expressions are within a certain tolerance :param row1: first resulting row of a query execution job for first SQL query :param row2: first resulting row of a query execution job for second SQL query :param metrics_thresholds: a dictionary of ratios indexed by metrics, for example 'COUNT(*)': 1.5 would require a 50 percent or less difference between the current day, and the prior days_back. :param ignore_zero: whether we should ignore zero metrics :param ratio_formula: which formula to use to compute the ratio between the two metrics. Assuming cur is the metric of today and ref is the metric to today - days_back. max_over_min: computes max(cur, ref) / min(cur, ref) relative_diff: computes abs(cur-ref) / ref """ifnotrow2:raiseAirflowException("The second SQL query returned None")ifnotrow1:raiseAirflowException("The first SQL query returned None")ratio_formulas={"max_over_min":lambdacur,ref:float(max(cur,ref))/min(cur,ref),"relative_diff":lambdacur,ref:float(abs(cur-ref))/ref,}metrics_sorted=sorted(metrics_thresholds.keys())current=dict(zip(metrics_sorted,row1))reference=dict(zip(metrics_sorted,row2))ratios:dict[str,Any]={}test_results:dict[str,Any]={}formetricinmetrics_sorted:cur=float(current[metric])ref=float(reference[metric])threshold=float(metrics_thresholds[metric])ifcur==0orref==0:ratios[metric]=Nonetest_results[metric]=ignore_zeroelse:ratios[metric]=ratio_formulas[ratio_formula](float(current[metric]),float(reference[metric]))test_results[metric]=float(ratios[metric])<thresholdself.log.info(("Current metric for %s: %s\n""Past metric for %s: %s\n""Ratio for %s: %s\n""Threshold: %s\n"),metric,cur,metric,ref,metric,ratios[metric],threshold,)ifnotall(test_results.values()):failed_tests=[metricformetric,valueintest_results.items()ifnotvalue]self.log.warning("The following %s tests out of %s failed:",len(failed_tests),len(metrics_sorted),)forkinfailed_tests:self.log.warning("'%s' check failed. %s is above %s",k,ratios[k],metrics_thresholds[k],)raiseAirflowException(f"The following tests have failed:\n{', '.join(sorted(failed_tests))}")self.log.info("All tests have passed")
[docs]classBigQueryTableAsyncHook(GoogleBaseAsyncHook):"""Class to get async hook for Bigquery Table Async"""
[docs]asyncdefget_table_client(self,dataset:str,table_id:str,project_id:str,session:ClientSession)->Table_async:""" Returns a Google Big Query Table object. :param dataset: The name of the dataset in which to look for the table storage bucket. :param table_id: The name of the table to check the existence of. :param project_id: The Google cloud project in which to look for the table. The connection supplied to the hook must provide access to the specified project. :param session: aiohttp ClientSession """withawaitself.service_file_as_context()asfile:returnTable_async(dataset_name=dataset,table_name=table_id,project=project_id,service_file=file,session=cast(Session,session),