apache-airflow-providers-google 10.7.0__py3-none-any.whl → 10.8.0__py3-none-any.whl

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

@@ -25,7 +25,6 @@ from googleapiclient.errors import HttpError
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException
 from airflow.hooks.base import BaseHook
-from airflow.models import Connection
 from airflow.providers.google.cloud.hooks.cloud_sql import CloudSQLDatabaseHook, CloudSQLHook
 from airflow.providers.google.cloud.links.cloud_sql import CloudSQLInstanceDatabaseLink, CloudSQLInstanceLink
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
@@ -33,10 +32,11 @@ from airflow.providers.google.cloud.triggers.cloud_sql import CloudSQLExportTrig
 from airflow.providers.google.cloud.utils.field_validator import GcpBodyFieldValidator
 from airflow.providers.google.common.hooks.base_google import get_field
 from airflow.providers.google.common.links.storage import FileDetailsLink
-from airflow.providers.mysql.hooks.mysql import MySqlHook
-from airflow.providers.postgres.hooks.postgres import PostgresHook
 
 if TYPE_CHECKING:
+    from airflow.models import Connection
+    from airflow.providers.mysql.hooks.mysql import MySqlHook
+    from airflow.providers.postgres.hooks.postgres import PostgresHook
     from airflow.utils.context import Context
 
 
@@ -44,175 +44,180 @@ SETTINGS = "settings"
 SETTINGS_VERSION = "settingsVersion"
 
 CLOUD_SQL_CREATE_VALIDATION: Sequence[dict] = [
-    dict(name="name", allow_empty=False),
-    dict(
-        name="settings",
-        type="dict",
-        fields=[
-            dict(name="tier", allow_empty=False),
-            dict(
-                name="backupConfiguration",
-                type="dict",
-                fields=[
-                    dict(name="binaryLogEnabled", optional=True),
-                    dict(name="enabled", optional=True),
-                    dict(name="replicationLogArchivingEnabled", optional=True),
-                    dict(name="startTime", allow_empty=False, optional=True),
+    {"name": "name", "allow_empty": False},
+    {
+        "name": "settings",
+        "type": "dict",
+        "fields": [
+            {"name": "tier", "allow_empty": False},
+            {
+                "name": "backupConfiguration",
+                "type": "dict",
+                "fields": [
+                    {"name": "binaryLogEnabled", "optional": True},
+                    {"name": "enabled", "optional": True},
+                    {"name": "replicationLogArchivingEnabled", "optional": True},
+                    {"name": "startTime", "allow_empty": False, "optional": True},
                 ],
-                optional=True,
-            ),
-            dict(name="activationPolicy", allow_empty=False, optional=True),
-            dict(name="authorizedGaeApplications", type="list", optional=True),
-            dict(name="crashSafeReplicationEnabled", optional=True),
-            dict(name="dataDiskSizeGb", optional=True),
-            dict(name="dataDiskType", allow_empty=False, optional=True),
-            dict(name="databaseFlags", type="list", optional=True),
-            dict(
-                name="ipConfiguration",
-                type="dict",
-                fields=[
-                    dict(
-                        name="authorizedNetworks",
-                        type="list",
-                        fields=[
-                            dict(name="expirationTime", optional=True),
-                            dict(name="name", allow_empty=False, optional=True),
-                            dict(name="value", allow_empty=False, optional=True),
+                "optional": True,
+            },
+            {"name": "activationPolicy", "allow_empty": False, "optional": True},
+            {"name": "authorizedGaeApplications", "type": "list", "optional": True},
+            {"name": "crashSafeReplicationEnabled", "optional": True},
+            {"name": "dataDiskSizeGb", "optional": True},
+            {"name": "dataDiskType", "allow_empty": False, "optional": True},
+            {"name": "databaseFlags", "type": "list", "optional": True},
+            {
+                "name": "ipConfiguration",
+                "type": "dict",
+                "fields": [
+                    {
+                        "name": "authorizedNetworks",
+                        "type": "list",
+                        "fields": [
+                            {"name": "expirationTime", "optional": True},
+                            {"name": "name", "allow_empty": False, "optional": True},
+                            {"name": "value", "allow_empty": False, "optional": True},
                         ],
-                        optional=True,
-                    ),
-                    dict(name="ipv4Enabled", optional=True),
-                    dict(name="privateNetwork", allow_empty=False, optional=True),
-                    dict(name="requireSsl", optional=True),
+                        "optional": True,
+                    },
+                    {"name": "ipv4Enabled", "optional": True},
+                    {"name": "privateNetwork", "allow_empty": False, "optional": True},
+                    {"name": "requireSsl", "optional": True},
                 ],
-                optional=True,
-            ),
-            dict(
-                name="locationPreference",
-                type="dict",
-                fields=[
-                    dict(name="followGaeApplication", allow_empty=False, optional=True),
-                    dict(name="zone", allow_empty=False, optional=True),
+                "optional": True,
+            },
+            {
+                "name": "locationPreference",
+                "type": "dict",
+                "fields": [
+                    {"name": "followGaeApplication", "allow_empty": False, "optional": True},
+                    {"name": "zone", "allow_empty": False, "optional": True},
                 ],
-                optional=True,
-            ),
-            dict(
-                name="maintenanceWindow",
-                type="dict",
-                fields=[
-                    dict(name="hour", optional=True),
-                    dict(name="day", optional=True),
-                    dict(name="updateTrack", allow_empty=False, optional=True),
+                "optional": True,
+            },
+            {
+                "name": "maintenanceWindow",
+                "type": "dict",
+                "fields": [
+                    {"name": "hour", "optional": True},
+                    {"name": "day", "optional": True},
+                    {"name": "updateTrack", "allow_empty": False, "optional": True},
                 ],
-                optional=True,
-            ),
-            dict(name="pricingPlan", allow_empty=False, optional=True),
-            dict(name="replicationType", allow_empty=False, optional=True),
-            dict(name="storageAutoResize", optional=True),
-            dict(name="storageAutoResizeLimit", optional=True),
-            dict(name="userLabels", type="dict", optional=True),
+                "optional": True,
+            },
+            {"name": "pricingPlan", "allow_empty": False, "optional": True},
+            {"name": "replicationType", "allow_empty": False, "optional": True},
+            {"name": "storageAutoResize", "optional": True},
+            {"name": "storageAutoResizeLimit", "optional": True},
+            {"name": "userLabels", "type": "dict", "optional": True},
         ],
-    ),
-    dict(name="databaseVersion", allow_empty=False, optional=True),
-    dict(name="failoverReplica", type="dict", fields=[dict(name="name", allow_empty=False)], optional=True),
-    dict(name="masterInstanceName", allow_empty=False, optional=True),
-    dict(name="onPremisesConfiguration", type="dict", optional=True),
-    dict(name="region", allow_empty=False, optional=True),
-    dict(
-        name="replicaConfiguration",
-        type="dict",
-        fields=[
-            dict(name="failoverTarget", optional=True),
-            dict(
-                name="mysqlReplicaConfiguration",
-                type="dict",
-                fields=[
-                    dict(name="caCertificate", allow_empty=False, optional=True),
-                    dict(name="clientCertificate", allow_empty=False, optional=True),
-                    dict(name="clientKey", allow_empty=False, optional=True),
-                    dict(name="connectRetryInterval", optional=True),
-                    dict(name="dumpFilePath", allow_empty=False, optional=True),
-                    dict(name="masterHeartbeatPeriod", optional=True),
-                    dict(name="password", allow_empty=False, optional=True),
-                    dict(name="sslCipher", allow_empty=False, optional=True),
-                    dict(name="username", allow_empty=False, optional=True),
-                    dict(name="verifyServerCertificate", optional=True),
+    },
+    {"name": "databaseVersion", "allow_empty": False, "optional": True},
+    {
+        "name": "failoverReplica",
+        "type": "dict",
+        "fields": [{"name": "name", "allow_empty": False}],
+        "optional": True,
+    },
+    {"name": "masterInstanceName", "allow_empty": False, "optional": True},
+    {"name": "onPremisesConfiguration", "type": "dict", "optional": True},
+    {"name": "region", "allow_empty": False, "optional": True},
+    {
+        "name": "replicaConfiguration",
+        "type": "dict",
+        "fields": [
+            {"name": "failoverTarget", "optional": True},
+            {
+                "name": "mysqlReplicaConfiguration",
+                "type": "dict",
+                "fields": [
+                    {"name": "caCertificate", "allow_empty": False, "optional": True},
+                    {"name": "clientCertificate", "allow_empty": False, "optional": True},
+                    {"name": "clientKey", "allow_empty": False, "optional": True},
+                    {"name": "connectRetryInterval", "optional": True},
+                    {"name": "dumpFilePath", "allow_empty": False, "optional": True},
+                    {"name": "masterHeartbeatPeriod", "optional": True},
+                    {"name": "password", "allow_empty": False, "optional": True},
+                    {"name": "sslCipher", "allow_empty": False, "optional": True},
+                    {"name": "username", "allow_empty": False, "optional": True},
+                    {"name": "verifyServerCertificate", "optional": True},
                 ],
-                optional=True,
-            ),
+                "optional": True,
+            },
         ],
-        optional=True,
-    ),
+        "optional": True,
+    },
 ]
 CLOUD_SQL_EXPORT_VALIDATION = [
-    dict(
-        name="exportContext",
-        type="dict",
-        fields=[
-            dict(name="fileType", allow_empty=False),
-            dict(name="uri", allow_empty=False),
-            dict(name="databases", optional=True, type="list"),
-            dict(
-                name="sqlExportOptions",
-                type="dict",
-                optional=True,
-                fields=[
-                    dict(name="tables", optional=True, type="list"),
-                    dict(name="schemaOnly", optional=True),
-                    dict(
-                        name="mysqlExportOptions",
-                        type="dict",
-                        optional=True,
-                        fields=[dict(name="masterData")],
-                    ),
+    {
+        "name": "exportContext",
+        "type": "dict",
+        "fields": [
+            {"name": "fileType", "allow_empty": False},
+            {"name": "uri", "allow_empty": False},
+            {"name": "databases", "optional": True, "type": "list"},
+            {
+                "name": "sqlExportOptions",
+                "type": "dict",
+                "optional": True,
+                "fields": [
+                    {"name": "tables", "optional": True, "type": "list"},
+                    {"name": "schemaOnly", "optional": True},
+                    {
+                        "name": "mysqlExportOptions",
+                        "type": "dict",
+                        "optional": True,
+                        "fields": [{"name": "masterData"}],
+                    },
                 ],
-            ),
-            dict(
-                name="csvExportOptions",
-                type="dict",
-                optional=True,
-                fields=[
-                    dict(name="selectQuery"),
-                    dict(name="escapeCharacter", optional=True),
-                    dict(name="quoteCharacter", optional=True),
-                    dict(name="fieldsTerminatedBy", optional=True),
-                    dict(name="linesTerminatedBy", optional=True),
+            },
+            {
+                "name": "csvExportOptions",
+                "type": "dict",
+                "optional": True,
+                "fields": [
+                    {"name": "selectQuery"},
+                    {"name": "escapeCharacter", "optional": True},
+                    {"name": "quoteCharacter", "optional": True},
+                    {"name": "fieldsTerminatedBy", "optional": True},
+                    {"name": "linesTerminatedBy", "optional": True},
                 ],
-            ),
-            dict(name="offload", optional=True),
+            },
+            {"name": "offload", "optional": True},
         ],
-    )
+    }
 ]
 CLOUD_SQL_IMPORT_VALIDATION = [
-    dict(
-        name="importContext",
-        type="dict",
-        fields=[
-            dict(name="fileType", allow_empty=False),
-            dict(name="uri", allow_empty=False),
-            dict(name="database", optional=True, allow_empty=False),
-            dict(name="importUser", optional=True),
-            dict(
-                name="csvImportOptions",
-                type="dict",
-                optional=True,
-                fields=[dict(name="table"), dict(name="columns", type="list", optional=True)],
-            ),
+    {
+        "name": "importContext",
+        "type": "dict",
+        "fields": [
+            {"name": "fileType", "allow_empty": False},
+            {"name": "uri", "allow_empty": False},
+            {"name": "database", "optional": True, "allow_empty": False},
+            {"name": "importUser", "optional": True},
+            {
+                "name": "csvImportOptions",
+                "type": "dict",
+                "optional": True,
+                "fields": [{"name": "table"}, {"name": "columns", "type": "list", "optional": True}],
+            },
         ],
-    )
+    }
 ]
 CLOUD_SQL_DATABASE_CREATE_VALIDATION = [
-    dict(name="instance", allow_empty=False),
-    dict(name="name", allow_empty=False),
-    dict(name="project", allow_empty=False),
+    {"name": "instance", "allow_empty": False},
+    {"name": "name", "allow_empty": False},
+    {"name": "project", "allow_empty": False},
 ]
 CLOUD_SQL_DATABASE_PATCH_VALIDATION = [
-    dict(name="instance", optional=True),
-    dict(name="name", optional=True),
-    dict(name="project", optional=True),
-    dict(name="etag", optional=True),
-    dict(name="charset", optional=True),
-    dict(name="collation", optional=True),
+    {"name": "instance", "optional": True},
+    {"name": "name", "optional": True},
+    {"name": "project", "optional": True},
+    {"name": "etag", "optional": True},
+    {"name": "charset", "optional": True},
+    {"name": "collation", "optional": True},
 ]
 
 

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

@@ -22,7 +22,6 @@ from copy import deepcopy
 from typing import TYPE_CHECKING, Any, Sequence
 
 from google.api_core import exceptions
-from google.api_core.retry import Retry
 from google.cloud.compute_v1.types import Instance, InstanceGroupManager, InstanceTemplate
 from json_merge_patch import merge
 
@@ -38,6 +37,8 @@ from airflow.providers.google.cloud.utils.field_sanitizer import GcpBodyFieldSan
 from airflow.providers.google.cloud.utils.field_validator import GcpBodyFieldValidator
 
 if TYPE_CHECKING:
+    from google.api_core.retry import Retry
+
     from airflow.utils.context import Context
 
 
@@ -173,14 +174,13 @@ class ComputeEngineInsertInstanceOperator(ComputeEngineBaseOperator):
     def check_body_fields(self) -> None:
         required_params = ["machine_type", "disks", "network_interfaces"]
         for param in required_params:
-            if param in self.body:
-                continue
-            readable_param = param.replace("_", " ")
-            raise AirflowException(
-                f"The body '{self.body}' should contain at least {readable_param} for the new operator "
-                f"in the '{param}' field. Check (google.cloud.compute_v1.types.Instance) "
-                f"for more details about body fields description."
-            )
+            if param not in self.body:
+                readable_param = param.replace("_", " ")
+                raise AirflowException(
+                    f"The body '{self.body}' should contain at least {readable_param} for the new operator "
+                    f"in the '{param}' field. Check (google.cloud.compute_v1.types.Instance) "
+                    f"for more details about body fields description."
+                )
 
     def _validate_inputs(self) -> None:
         super()._validate_inputs()
@@ -661,7 +661,7 @@ class ComputeEngineStopInstanceOperator(ComputeEngineBaseOperator):
 
 
 SET_MACHINE_TYPE_VALIDATION_SPECIFICATION = [
-    dict(name="machineType", regexp="^.+$"),
+    {"name": "machineType", "regexp": "^.+$"},
 ]
 
 
@@ -768,44 +768,44 @@ class ComputeEngineSetMachineTypeOperator(ComputeEngineBaseOperator):
 
 
 GCE_INSTANCE_TEMPLATE_VALIDATION_PATCH_SPECIFICATION: list[dict[str, Any]] = [
-    dict(name="name", regexp="^.+$"),
-    dict(name="description", optional=True),
-    dict(
-        name="properties",
-        type="dict",
-        optional=True,
-        fields=[
-            dict(name="description", optional=True),
-            dict(name="tags", optional=True, fields=[dict(name="items", optional=True)]),
-            dict(name="machineType", optional=True),
-            dict(name="canIpForward", optional=True),
-            dict(name="networkInterfaces", optional=True),  # not validating deeper
-            dict(name="disks", optional=True),  # not validating the array deeper
-            dict(
-                name="metadata",
-                optional=True,
-                fields=[
-                    dict(name="fingerprint", optional=True),
-                    dict(name="items", optional=True),
-                    dict(name="kind", optional=True),
+    {"name": "name", "regexp": "^.+$"},
+    {"name": "description", "optional": True},
+    {
+        "name": "properties",
+        "type": "dict",
+        "optional": True,
+        "fields": [
+            {"name": "description", "optional": True},
+            {"name": "tags", "optional": True, "fields": [{"name": "items", "optional": True}]},
+            {"name": "machineType", "optional": True},
+            {"name": "canIpForward", "optional": True},
+            {"name": "networkInterfaces", "optional": True},  # not validating deeper
+            {"name": "disks", "optional": True},  # not validating the array deeper
+            {
+                "name": "metadata",
+                "optional": True,
+                "fields": [
+                    {"name": "fingerprint", "optional": True},
+                    {"name": "items", "optional": True},
+                    {"name": "kind", "optional": True},
                 ],
-            ),
-            dict(name="serviceAccounts", optional=True),  # not validating deeper
-            dict(
-                name="scheduling",
-                optional=True,
-                fields=[
-                    dict(name="onHostMaintenance", optional=True),
-                    dict(name="automaticRestart", optional=True),
-                    dict(name="preemptible", optional=True),
-                    dict(name="nodeAffinities", optional=True),  # not validating deeper
+            },
+            {"name": "serviceAccounts", "optional": True},  # not validating deeper
+            {
+                "name": "scheduling",
+                "optional": True,
+                "fields": [
+                    {"name": "onHostMaintenance", "optional": True},
+                    {"name": "automaticRestart", "optional": True},
+                    {"name": "preemptible", "optional": True},
+                    {"name": "nodeAffinities", "optional": True},  # not validating deeper
                 ],
-            ),
-            dict(name="labels", optional=True),
-            dict(name="guestAccelerators", optional=True),  # not validating deeper
-            dict(name="minCpuPlatform", optional=True),
+            },
+            {"name": "labels", "optional": True},
+            {"name": "guestAccelerators", "optional": True},  # not validating deeper
+            {"name": "minCpuPlatform", "optional": True},
         ],
-    ),
+    },
 ]
 
 GCE_INSTANCE_FIELDS_TO_SANITIZE = [
@@ -914,14 +914,13 @@ class ComputeEngineInsertInstanceTemplateOperator(ComputeEngineBaseOperator):
     def check_body_fields(self) -> None:
         required_params = ["machine_type", "disks", "network_interfaces"]
         for param in required_params:
-            if param in self.body["properties"]:
-                continue
-            readable_param = param.replace("_", " ")
-            raise AirflowException(
-                f"The body '{self.body}' should contain at least {readable_param} for the new operator "
-                f"in the '{param}' field. Check (google.cloud.compute_v1.types.Instance) "
-                f"for more details about body fields description."
-            )
+            if param not in self.body["properties"]:
+                readable_param = param.replace("_", " ")
+                raise AirflowException(
+                    f"The body '{self.body}' should contain at least {readable_param} for the new operator "
+                    f"in the '{param}' field. Check (google.cloud.compute_v1.types.Instance) "
+                    f"for more details about body fields description."
+                )
 
     def _validate_all_body_fields(self) -> None:
         if self._field_validator:
@@ -1499,14 +1498,13 @@ class ComputeEngineInsertInstanceGroupManagerOperator(ComputeEngineBaseOperator)
     def check_body_fields(self) -> None:
         required_params = ["base_instance_name", "target_size", "instance_template"]
         for param in required_params:
-            if param in self.body:
-                continue
-            readable_param = param.replace("_", " ")
-            raise AirflowException(
-                f"The body '{self.body}' should contain at least {readable_param} for the new operator "
-                f"in the '{param}' field. Check (google.cloud.compute_v1.types.Instance) "
-                f"for more details about body fields description."
-            )
+            if param not in self.body:
+                readable_param = param.replace("_", " ")
+                raise AirflowException(
+                    f"The body '{self.body}' should contain at least {readable_param} for the new operator "
+                    f"in the '{param}' field. Check (google.cloud.compute_v1.types.Instance) "
+                    f"for more details about body fields description."
+                )
 
     def _validate_all_body_fields(self) -> None:
         if self._field_validator:

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

@@ -20,7 +20,6 @@ from typing import TYPE_CHECKING, Sequence
 
 from google.api_core.exceptions import AlreadyExists, NotFound
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
-from google.api_core.retry import Retry
 from google.cloud.datacatalog import (
     DataCatalogClient,
     Entry,
@@ -31,7 +30,6 @@ from google.cloud.datacatalog import (
     TagTemplate,
     TagTemplateField,
 )
-from google.protobuf.field_mask_pb2 import FieldMask
 
 from airflow.providers.google.cloud.hooks.datacatalog import CloudDataCatalogHook
 from airflow.providers.google.cloud.links.datacatalog import (
@@ -42,6 +40,9 @@ from airflow.providers.google.cloud.links.datacatalog import (
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
 
 if TYPE_CHECKING:
+    from google.api_core.retry import Retry
+    from google.protobuf.field_mask_pb2 import FieldMask
+
     from airflow.utils.context import Context
 
 

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

@@ -1188,7 +1188,9 @@ class DataflowCreatePythonJobOperator(GoogleCloudBaseOperator):
         pipeline_options.update(self.options)
 
         # Convert argument names from lowerCamelCase to snake case.
-        camel_to_snake = lambda name: re.sub(r"[A-Z]", lambda x: "_" + x.group(0).lower(), name)
+        def camel_to_snake(name):
+            return re.sub("[A-Z]", lambda x: "_" + x.group(0).lower(), name)
+
         formatted_pipeline_options = {camel_to_snake(key): pipeline_options[key] for key in pipeline_options}
 
         def set_current_job_id(job_id):

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

@@ -25,10 +25,11 @@ from airflow.providers.google.cloud.links.dataform import (
 )
 
 if TYPE_CHECKING:
+    from google.api_core.retry import Retry
+
     from airflow.utils.context import Context
 
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
-from google.api_core.retry import Retry
 from google.cloud.dataform_v1beta1.types import (
     CompilationResult,
     InstallNpmPackagesResponse,

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

@@ -783,7 +783,7 @@ class CloudDataFusionStartPipelineOperator(GoogleCloudBaseOperator):
         if success_states:
             self.success_states = success_states
         else:
-            self.success_states = SUCCESS_STATES + [PipelineStates.RUNNING]
+            self.success_states = [*SUCCESS_STATES, PipelineStates.RUNNING]
 
     def execute(self, context: Context) -> str:
         hook = DataFusionHook(