apache-airflow-providers-google 10.20.0rc1__py3-none-any.whl → 10.21.0__py3-none-any.whl

airflow/providers/google/cloud/operators/bigquery.py

@@ -167,7 +167,8 @@ class _BigQueryOperatorsEncryptionConfigurationMixin:
 class BigQueryCheckOperator(
     _BigQueryDbHookMixin, SQLCheckOperator, _BigQueryOperatorsEncryptionConfigurationMixin
 ):
-    """Performs checks against BigQuery.
+    """
+    Performs checks against BigQuery.
 
     This operator expects a SQL query that returns a single row. Each value on
     that row is evaluated using a Python ``bool`` cast. If any of the values
@@ -221,6 +222,12 @@ class BigQueryCheckOperator(
     :param deferrable: Run operator in the deferrable mode.
     :param poll_interval: (Deferrable mode only) polling period in seconds to
         check for the status of job.
+    :param query_params: a list of dictionary containing query parameter types and
+        values, passed to BigQuery. The structure of dictionary should look like
+        'queryParameters' in Google BigQuery Jobs API:
+        https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs.
+        For example, [{ 'name': 'corpus', 'parameterType': { 'type': 'STRING' },
+        'parameterValue': { 'value': 'romeoandjuliet' } }]. (templated)
     """
 
     template_fields: Sequence[str] = (
@@ -228,6 +235,7 @@ class BigQueryCheckOperator(
         "gcp_conn_id",
         "impersonation_chain",
         "labels",
+        "query_params",
     )
     template_ext: Sequence[str] = (".sql",)
     ui_color = BigQueryUIColors.CHECK.value
@@ -245,6 +253,7 @@ class BigQueryCheckOperator(
         encryption_configuration: dict | None = None,
         deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
         poll_interval: float = 4.0,
+        query_params: list | None = None,
         **kwargs,
     ) -> None:
         super().__init__(sql=sql, **kwargs)
@@ -256,6 +265,7 @@ class BigQueryCheckOperator(
         self.encryption_configuration = encryption_configuration
         self.deferrable = deferrable
         self.poll_interval = poll_interval
+        self.query_params = query_params
 
     def _submit_job(
         self,
@@ -264,6 +274,8 @@ class BigQueryCheckOperator(
     ) -> BigQueryJob:
         """Submit a new job and get the job id for polling the status using Trigger."""
         configuration = {"query": {"query": self.sql, "useLegacySql": self.use_legacy_sql}}
+        if self.query_params:
+            configuration["query"]["queryParameters"] = self.query_params
 
         self.include_encryption_configuration(configuration, "query")
 
@@ -319,7 +331,8 @@ class BigQueryCheckOperator(
             )
 
     def execute_complete(self, context: Context, event: dict[str, Any]) -> None:
-        """Act as a callback for when the trigger fires.
+        """
+        Act as a callback for when the trigger fires.
 
         This returns immediately. It relies on trigger to throw an exception,
         otherwise it assumes execution was successful.
@@ -335,7 +348,8 @@ class BigQueryCheckOperator(
 class BigQueryValueCheckOperator(
     _BigQueryDbHookMixin, SQLValueCheckOperator, _BigQueryOperatorsEncryptionConfigurationMixin
 ):
-    """Perform a simple value check using sql code.
+    """
+    Perform a simple value check using sql code.
 
     .. seealso::
         For more information on how to use this operator, take a look at the guide:
@@ -465,7 +479,8 @@ class BigQueryValueCheckOperator(
             raise AirflowException(f"BigQuery job {job.job_id} failed: {job.error_result}")
 
     def execute_complete(self, context: Context, event: dict[str, Any]) -> None:
-        """Act as a callback for when the trigger fires.
+        """
+        Act as a callback for when the trigger fires.
 
         This returns immediately. It relies on trigger to throw an exception,
         otherwise it assumes execution was successful.
@@ -625,7 +640,8 @@ class BigQueryIntervalCheckOperator(
             )
 
     def execute_complete(self, context: Context, event: dict[str, Any]) -> None:
-        """Act as a callback for when the trigger fires.
+        """
+        Act as a callback for when the trigger fires.
 
         This returns immediately. It relies on trigger to throw an exception,
         otherwise it assumes execution was successful.
@@ -1173,7 +1189,8 @@ class BigQueryGetDataOperator(GoogleCloudBaseOperator, _BigQueryOperatorsEncrypt
         )
 
     def execute_complete(self, context: Context, event: dict[str, Any]) -> Any:
-        """Act as a callback for when the trigger fires.
+        """
+        Act as a callback for when the trigger fires.
 
         This returns immediately. It relies on trigger to throw an exception,
         otherwise it assumes execution was successful.
@@ -1190,7 +1207,8 @@ class BigQueryGetDataOperator(GoogleCloudBaseOperator, _BigQueryOperatorsEncrypt
     category=AirflowProviderDeprecationWarning,
 )
 class BigQueryExecuteQueryOperator(GoogleCloudBaseOperator):
-    """Executes BigQuery SQL queries in a specific BigQuery database.
+    """
+    Executes BigQuery SQL queries in a specific BigQuery database.
 
     This operator is deprecated. Please use
     :class:`airflow.providers.google.cloud.operators.bigquery.BigQueryInsertJobOperator`
@@ -1409,7 +1427,8 @@ class BigQueryExecuteQueryOperator(GoogleCloudBaseOperator):
 
 
 class BigQueryCreateEmptyTableOperator(GoogleCloudBaseOperator):
-    """Creates a new table in the specified BigQuery dataset, optionally with schema.
+    """
+    Creates a new table in the specified BigQuery dataset, optionally with schema.
 
     The schema to be used for the BigQuery table may be specified in one of
     two ways. You may either directly pass the schema fields in, or you may
@@ -1658,7 +1677,8 @@ class BigQueryCreateEmptyTableOperator(GoogleCloudBaseOperator):
 
 
 class BigQueryCreateExternalTableOperator(GoogleCloudBaseOperator):
-    """Create a new external table with data from Google Cloud Storage.
+    """
+    Create a new external table with data from Google Cloud Storage.
 
     The schema to be used for the BigQuery table may be specified in one of
     two ways. You may either directly pass the schema fields in, or you may
@@ -1953,7 +1973,8 @@ class BigQueryCreateExternalTableOperator(GoogleCloudBaseOperator):
 
 
 class BigQueryDeleteDatasetOperator(GoogleCloudBaseOperator):
-    """Delete an existing dataset from your Project in BigQuery.
+    """
+    Delete an existing dataset from your Project in BigQuery.
 
     https://cloud.google.com/bigquery/docs/reference/rest/v2/datasets/delete
 
@@ -2028,7 +2049,8 @@ class BigQueryDeleteDatasetOperator(GoogleCloudBaseOperator):
 
 
 class BigQueryCreateEmptyDatasetOperator(GoogleCloudBaseOperator):
-    """Create a new dataset for your Project in BigQuery.
+    """
+    Create a new dataset for your Project in BigQuery.
 
     https://cloud.google.com/bigquery/docs/reference/rest/v2/datasets#resource
 
@@ -2151,7 +2173,8 @@ class BigQueryCreateEmptyDatasetOperator(GoogleCloudBaseOperator):
 
 
 class BigQueryGetDatasetOperator(GoogleCloudBaseOperator):
-    """Get the dataset specified by ID.
+    """
+    Get the dataset specified by ID.
 
     .. seealso::
         For more information on how to use this operator, take a look at the guide:
@@ -2215,7 +2238,8 @@ class BigQueryGetDatasetOperator(GoogleCloudBaseOperator):
 
 
 class BigQueryGetDatasetTablesOperator(GoogleCloudBaseOperator):
-    """Retrieve the list of tables in the specified dataset.
+    """
+    Retrieve the list of tables in the specified dataset.
 
     .. seealso::
         For more information on how to use this operator, take a look at the guide:
@@ -2278,7 +2302,8 @@ class BigQueryGetDatasetTablesOperator(GoogleCloudBaseOperator):
     category=AirflowProviderDeprecationWarning,
 )
 class BigQueryPatchDatasetOperator(GoogleCloudBaseOperator):
-    """Patch a dataset for your Project in BigQuery.
+    """
+    Patch a dataset for your Project in BigQuery.
 
     This operator is deprecated. Please use
     :class:`airflow.providers.google.cloud.operators.bigquery.BigQueryUpdateTableOperator`
@@ -2342,7 +2367,8 @@ class BigQueryPatchDatasetOperator(GoogleCloudBaseOperator):
 
 
 class BigQueryUpdateTableOperator(GoogleCloudBaseOperator):
-    """Update a table for your Project in BigQuery.
+    """
+    Update a table for your Project in BigQuery.
 
     Use ``fields`` to specify which fields of table to update. If a field
     is listed in ``fields`` and is ``None`` in table, it will be deleted.
@@ -2429,7 +2455,8 @@ class BigQueryUpdateTableOperator(GoogleCloudBaseOperator):
 
 
 class BigQueryUpdateDatasetOperator(GoogleCloudBaseOperator):
-    """Update a dataset for your Project in BigQuery.
+    """
+    Update a dataset for your Project in BigQuery.
 
     Use ``fields`` to specify which fields of dataset to update. If a field
     is listed in ``fields`` and is ``None`` in dataset, it will be deleted.
@@ -2511,7 +2538,8 @@ class BigQueryUpdateDatasetOperator(GoogleCloudBaseOperator):
 
 
 class BigQueryDeleteTableOperator(GoogleCloudBaseOperator):
-    """Delete a BigQuery table.
+    """
+    Delete a BigQuery table.
 
     .. seealso::
         For more information on how to use this operator, take a look at the guide:
@@ -2569,7 +2597,8 @@ class BigQueryDeleteTableOperator(GoogleCloudBaseOperator):
 
 
 class BigQueryUpsertTableOperator(GoogleCloudBaseOperator):
-    """Upsert to a BigQuery table.
+    """
+    Upsert to a BigQuery table.
 
     .. seealso::
         For more information on how to use this operator, take a look at the guide:
@@ -2646,7 +2675,8 @@ class BigQueryUpsertTableOperator(GoogleCloudBaseOperator):
 
 
 class BigQueryUpdateTableSchemaOperator(GoogleCloudBaseOperator):
-    """Update BigQuery Table Schema.
+    """
+    Update BigQuery Table Schema.
 
     Updates fields on a table schema based on contents of the supplied schema_fields_updates
     parameter. The supplied schema does not need to be complete, if the field
@@ -2756,7 +2786,8 @@ class BigQueryUpdateTableSchemaOperator(GoogleCloudBaseOperator):
 
 
 class BigQueryInsertJobOperator(GoogleCloudBaseOperator, _BigQueryOpenLineageMixin):
-    """Execute a BigQuery job.
+    """
+    Execute a BigQuery job.
 
     Waits for the job to complete and returns job id.
     This operator work in the following way:
@@ -3005,7 +3036,8 @@ class BigQueryInsertJobOperator(GoogleCloudBaseOperator, _BigQueryOpenLineageMix
             self._handle_job_error(job)
 
     def execute_complete(self, context: Context, event: dict[str, Any]) -> str | None:
-        """Act as a callback for when the trigger fires.
+        """
+        Act as a callback for when the trigger fires.
 
         This returns immediately. It relies on trigger to throw an exception,
         otherwise it assumes execution was successful.
@@ -3017,6 +3049,8 @@ class BigQueryInsertJobOperator(GoogleCloudBaseOperator, _BigQueryOpenLineageMix
             self.task_id,
             event["message"],
         )
+        # Save job_id as an attribute to be later used by listeners
+        self.job_id = event.get("job_id")
         return self.job_id
 
     def on_kill(self) -> None:

airflow/providers/google/cloud/operators/cloud_batch.py

@@ -86,7 +86,9 @@ class CloudBatchSubmitJobOperator(GoogleCloudBaseOperator):
 
     def execute(self, context: Context):
         hook: CloudBatchHook = CloudBatchHook(self.gcp_conn_id, self.impersonation_chain)
-        job = hook.submit_batch_job(job_name=self.job_name, job=self.job, region=self.region)
+        job = hook.submit_batch_job(
+            job_name=self.job_name, job=self.job, region=self.region, project_id=self.project_id
+        )
 
         if not self.deferrable:
             completed_job = hook.wait_for_job(

airflow/providers/google/cloud/operators/cloud_sql.py

@@ -224,7 +224,8 @@ CLOUD_SQL_DATABASE_PATCH_VALIDATION = [
 
 
 class CloudSQLBaseOperator(GoogleCloudBaseOperator):
-    """Abstract base operator for Google Cloud SQL operators.
+    """
+    Abstract base operator for Google Cloud SQL operators.
 
     :param instance: Cloud SQL instance ID. This does not include the project ID.
     :param project_id: Optional, Google Cloud Project ID.  f set to None or missing,
@@ -292,7 +293,8 @@ class CloudSQLBaseOperator(GoogleCloudBaseOperator):
 
 
 class CloudSQLCreateInstanceOperator(CloudSQLBaseOperator):
-    """Create a new Cloud SQL instance.
+    """
+    Create a new Cloud SQL instance.
 
     If an instance with the same name exists, no action will be taken and
     the operator will succeed.
@@ -393,7 +395,8 @@ class CloudSQLCreateInstanceOperator(CloudSQLBaseOperator):
 
 
 class CloudSQLInstancePatchOperator(CloudSQLBaseOperator):
-    """Update settings of a Cloud SQL instance.
+    """
+    Update settings of a Cloud SQL instance.
 
     Caution: This is a partial update, so only included values for the settings will be
     updated.
@@ -485,7 +488,8 @@ class CloudSQLInstancePatchOperator(CloudSQLBaseOperator):
 
 
 class CloudSQLDeleteInstanceOperator(CloudSQLBaseOperator):
-    """Delete a Cloud SQL instance.
+    """
+    Delete a Cloud SQL instance.
 
     .. seealso::
         For more information on how to use this operator, take a look at the guide:
@@ -531,7 +535,8 @@ class CloudSQLDeleteInstanceOperator(CloudSQLBaseOperator):
 
 
 class CloudSQLCloneInstanceOperator(CloudSQLBaseOperator):
-    """Clone an instance to a target instance.
+    """
+    Clone an instance to a target instance.
 
     .. seealso::
         For more information on how to use this operator, take a look at the guide:
@@ -622,7 +627,8 @@ class CloudSQLCloneInstanceOperator(CloudSQLBaseOperator):
 
 
 class CloudSQLCreateInstanceDatabaseOperator(CloudSQLBaseOperator):
-    """Create a new database inside a Cloud SQL instance.
+    """
+    Create a new database inside a Cloud SQL instance.
 
     .. seealso::
         For more information on how to use this operator, take a look at the guide:
@@ -726,7 +732,8 @@ class CloudSQLCreateInstanceDatabaseOperator(CloudSQLBaseOperator):
 
 
 class CloudSQLPatchInstanceDatabaseOperator(CloudSQLBaseOperator):
-    """Update resource containing information about a database using patch semantics.
+    """
+    Update resource containing information about a database using patch semantics.
 
     See: https://cloud.google.com/sql/docs/mysql/admin-api/how-tos/performance#patch
 
@@ -829,7 +836,8 @@ class CloudSQLPatchInstanceDatabaseOperator(CloudSQLBaseOperator):
 
 
 class CloudSQLDeleteInstanceDatabaseOperator(CloudSQLBaseOperator):
-    """Delete a database from a Cloud SQL instance.
+    """
+    Delete a database from a Cloud SQL instance.
 
     .. seealso::
         For more information on how to use this operator, take a look at the guide:
@@ -908,7 +916,8 @@ class CloudSQLDeleteInstanceDatabaseOperator(CloudSQLBaseOperator):
 
 
 class CloudSQLExportInstanceOperator(CloudSQLBaseOperator):
-    """Export data from a Cloud SQL instance to a Cloud Storage bucket.
+    """
+    Export data from a Cloud SQL instance to a Cloud Storage bucket.
 
     The exported format can be a SQL dump or CSV file.
 
@@ -1045,7 +1054,8 @@ class CloudSQLExportInstanceOperator(CloudSQLBaseOperator):
 
 
 class CloudSQLImportInstanceOperator(CloudSQLBaseOperator):
-    """Import data into a Cloud SQL instance from Cloud Storage.
+    """
+    Import data into a Cloud SQL instance from Cloud Storage.
 
     CSV IMPORT
     ``````````
@@ -1157,7 +1167,8 @@ class CloudSQLImportInstanceOperator(CloudSQLBaseOperator):
 
 
 class CloudSQLExecuteQueryOperator(GoogleCloudBaseOperator):
-    """Perform DML or DDL query on an existing Cloud Sql instance.
+    """
+    Perform DML or DDL query on an existing Cloud Sql instance.
 
     It optionally uses cloud-sql-proxy to establish secure connection with the
     database.

airflow/providers/google/cloud/operators/dataform.py

@@ -1143,7 +1143,8 @@ class DataformRemoveDirectoryOperator(GoogleCloudBaseOperator):
 
 
 class DataformInstallNpmPackagesOperator(GoogleCloudBaseOperator):
-    """Install NPM dependencies in the provided workspace.
+    """
+    Install NPM dependencies in the provided workspace.
 
     Requires "package.json" to be created in the workspace.