apache-airflow-providers-google 10.17.0rc1__py3-none-any.whl → 10.18.0__py3-none-any.whl

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

@@ -28,6 +28,7 @@ from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
 from airflow.providers.google.cloud.hooks.cloud_storage_transfer_service import (
     ACCESS_KEY_ID,
     AWS_ACCESS_KEY,
+    AWS_ROLE_ARN,
     AWS_S3_DATA_SOURCE,
     BUCKET_NAME,
     DAY,
@@ -62,6 +63,7 @@ from airflow.providers.google.cloud.links.cloud_storage_transfer import (
 )
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
 from airflow.providers.google.cloud.utils.helpers import normalize_directory_path
+from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context
@@ -78,15 +80,23 @@ class TransferJobPreprocessor:
         self.default_schedule = default_schedule
 
     def _inject_aws_credentials(self) -> None:
-        if TRANSFER_SPEC in self.body and AWS_S3_DATA_SOURCE in self.body[TRANSFER_SPEC]:
-            aws_hook = AwsBaseHook(self.aws_conn_id, resource_type="s3")
-            aws_credentials = aws_hook.get_credentials()
-            aws_access_key_id = aws_credentials.access_key  # type: ignore[attr-defined]
-            aws_secret_access_key = aws_credentials.secret_key  # type: ignore[attr-defined]
-            self.body[TRANSFER_SPEC][AWS_S3_DATA_SOURCE][AWS_ACCESS_KEY] = {
-                ACCESS_KEY_ID: aws_access_key_id,
-                SECRET_ACCESS_KEY: aws_secret_access_key,
-            }
+        if TRANSFER_SPEC not in self.body:
+            return
+
+        if AWS_S3_DATA_SOURCE not in self.body[TRANSFER_SPEC]:
+            return
+
+        if AWS_ROLE_ARN in self.body[TRANSFER_SPEC][AWS_S3_DATA_SOURCE]:
+            return
+
+        aws_hook = AwsBaseHook(self.aws_conn_id, resource_type="s3")
+        aws_credentials = aws_hook.get_credentials()
+        aws_access_key_id = aws_credentials.access_key  # type: ignore[attr-defined]
+        aws_secret_access_key = aws_credentials.secret_key  # type: ignore[attr-defined]
+        self.body[TRANSFER_SPEC][AWS_S3_DATA_SOURCE][AWS_ACCESS_KEY] = {
+            ACCESS_KEY_ID: aws_access_key_id,
+            SECRET_ACCESS_KEY: aws_secret_access_key,
+        }
 
     def _reformat_date(self, field_key: str) -> None:
         schedule = self.body[SCHEDULE]
@@ -234,7 +244,7 @@ class CloudDataTransferServiceCreateJobOperator(GoogleCloudBaseOperator):
         aws_conn_id: str | None = "aws_default",
         gcp_conn_id: str = "google_cloud_default",
         api_version: str = "v1",
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         google_impersonation_chain: str | Sequence[str] | None = None,
         **kwargs,
     ) -> None:
@@ -324,7 +334,7 @@ class CloudDataTransferServiceUpdateJobOperator(GoogleCloudBaseOperator):
         aws_conn_id: str | None = "aws_default",
         gcp_conn_id: str = "google_cloud_default",
         api_version: str = "v1",
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         google_impersonation_chain: str | Sequence[str] | None = None,
         **kwargs,
     ) -> None:
@@ -407,7 +417,7 @@ class CloudDataTransferServiceDeleteJobOperator(GoogleCloudBaseOperator):
         job_name: str,
         gcp_conn_id: str = "google_cloud_default",
         api_version: str = "v1",
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         google_impersonation_chain: str | Sequence[str] | None = None,
         **kwargs,
     ) -> None:
@@ -467,7 +477,7 @@ class CloudDataTransferServiceGetOperationOperator(GoogleCloudBaseOperator):
     def __init__(
         self,
         *,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         operation_name: str,
         gcp_conn_id: str = "google_cloud_default",
         api_version: str = "v1",
@@ -541,7 +551,7 @@ class CloudDataTransferServiceListOperationsOperator(GoogleCloudBaseOperator):
     def __init__(
         self,
         request_filter: dict,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         gcp_conn_id: str = "google_cloud_default",
         api_version: str = "v1",
         google_impersonation_chain: str | Sequence[str] | None = None,
@@ -818,6 +828,9 @@ class CloudDataTransferServiceS3ToGCSOperator(GoogleCloudBaseOperator):
         account from the list granting this role to the originating account (templated).
     :param delete_job_after_completion: If True, delete the job after complete.
         If set to True, 'wait' must be set to True.
+    :param aws_role_arn: Optional AWS role ARN for workload identity federation. This will
+        override the `aws_conn_id` for authentication between GCP and AWS; see
+        https://cloud.google.com/storage-transfer/docs/reference/rest/v1/TransferSpec#AwsS3Data
     """
 
     template_fields: Sequence[str] = (
@@ -829,6 +842,7 @@ class CloudDataTransferServiceS3ToGCSOperator(GoogleCloudBaseOperator):
         "description",
         "object_conditions",
         "google_impersonation_chain",
+        "aws_role_arn",
     )
     ui_color = "#e09411"
 
@@ -839,7 +853,7 @@ class CloudDataTransferServiceS3ToGCSOperator(GoogleCloudBaseOperator):
         gcs_bucket: str,
         s3_path: str | None = None,
         gcs_path: str | None = None,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         aws_conn_id: str | None = "aws_default",
         gcp_conn_id: str = "google_cloud_default",
         description: str | None = None,
@@ -850,6 +864,7 @@ class CloudDataTransferServiceS3ToGCSOperator(GoogleCloudBaseOperator):
         timeout: float | None = None,
         google_impersonation_chain: str | Sequence[str] | None = None,
         delete_job_after_completion: bool = False,
+        aws_role_arn: str | None = None,
         **kwargs,
     ) -> None:
         super().__init__(**kwargs)
@@ -868,6 +883,7 @@ class CloudDataTransferServiceS3ToGCSOperator(GoogleCloudBaseOperator):
         self.timeout = timeout
         self.google_impersonation_chain = google_impersonation_chain
         self.delete_job_after_completion = delete_job_after_completion
+        self.aws_role_arn = aws_role_arn
         self._validate_inputs()
 
     def _validate_inputs(self) -> None:
@@ -918,6 +934,9 @@ class CloudDataTransferServiceS3ToGCSOperator(GoogleCloudBaseOperator):
         if self.transfer_options is not None:
             body[TRANSFER_SPEC][TRANSFER_OPTIONS] = self.transfer_options  # type: ignore[index]
 
+        if self.aws_role_arn is not None:
+            body[TRANSFER_SPEC][AWS_S3_DATA_SOURCE][AWS_ROLE_ARN] = self.aws_role_arn  # type: ignore[index]
+
         return body
 
 
@@ -1007,7 +1026,7 @@ class CloudDataTransferServiceGCSToGCSOperator(GoogleCloudBaseOperator):
         destination_bucket: str,
         source_path: str | None = None,
         destination_path: str | None = None,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         gcp_conn_id: str = "google_cloud_default",
         description: str | None = None,
         schedule: dict | None = None,

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

@@ -36,6 +36,7 @@ from airflow.providers.google.cloud.links.compute import (
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
 from airflow.providers.google.cloud.utils.field_sanitizer import GcpBodyFieldSanitizer
 from airflow.providers.google.cloud.utils.field_validator import GcpBodyFieldValidator
+from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
 
 if TYPE_CHECKING:
     from google.api_core.retry import Retry
@@ -51,7 +52,7 @@ class ComputeEngineBaseOperator(GoogleCloudBaseOperator):
         *,
         zone: str,
         resource_id: str | None = None,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         gcp_conn_id: str = "google_cloud_default",
         api_version: str = "v1",
         impersonation_chain: str | Sequence[str] | None = None,
@@ -138,7 +139,7 @@ class ComputeEngineInsertInstanceOperator(ComputeEngineBaseOperator):
         body: dict,
         zone: str,
         resource_id: str | None = None,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         request_id: str | None = None,
         retry: Retry | None = None,
         timeout: float | None = None,
@@ -320,7 +321,7 @@ class ComputeEngineInsertInstanceFromTemplateOperator(ComputeEngineBaseOperator)
         body: dict,
         zone: str,
         resource_id: str | None = None,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         request_id: str | None = None,
         retry: Retry | None = None,
         timeout: float | None = None,
@@ -477,7 +478,7 @@ class ComputeEngineDeleteInstanceOperator(ComputeEngineBaseOperator):
         resource_id: str,
         zone: str,
         request_id: str | None = None,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         retry: Retry | None = None,
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
@@ -721,7 +722,7 @@ class ComputeEngineSetMachineTypeOperator(ComputeEngineBaseOperator):
         zone: str,
         resource_id: str,
         body: dict,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         gcp_conn_id: str = "google_cloud_default",
         api_version: str = "v1",
         validate_body: bool = True,
@@ -882,7 +883,7 @@ class ComputeEngineInsertInstanceTemplateOperator(ComputeEngineBaseOperator):
         self,
         *,
         body: dict,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         resource_id: str | None = None,
         request_id: str | None = None,
         retry: Retry | None = None,
@@ -1044,7 +1045,7 @@ class ComputeEngineDeleteInstanceTemplateOperator(ComputeEngineBaseOperator):
         *,
         resource_id: str,
         request_id: str | None = None,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         retry: Retry | None = None,
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
@@ -1164,7 +1165,7 @@ class ComputeEngineCopyInstanceTemplateOperator(ComputeEngineBaseOperator):
         *,
         resource_id: str,
         body_patch: dict,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         request_id: str | None = None,
         gcp_conn_id: str = "google_cloud_default",
         api_version: str = "v1",
@@ -1323,7 +1324,7 @@ class ComputeEngineInstanceGroupUpdateManagerTemplateOperator(ComputeEngineBaseO
         zone: str,
         source_template: str,
         destination_template: str,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         update_policy: dict[str, Any] | None = None,
         request_id: str | None = None,
         gcp_conn_id: str = "google_cloud_default",
@@ -1468,7 +1469,7 @@ class ComputeEngineInsertInstanceGroupManagerOperator(ComputeEngineBaseOperator)
         *,
         body: dict,
         zone: str,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         resource_id: str | None = None,
         request_id: str | None = None,
         gcp_conn_id: str = "google_cloud_default",
@@ -1630,7 +1631,7 @@ class ComputeEngineDeleteInstanceGroupManagerOperator(ComputeEngineBaseOperator)
         *,
         resource_id: str,
         zone: str,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         request_id: str | None = None,
         gcp_conn_id: str = "google_cloud_default",
         api_version="v1",

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

@@ -38,6 +38,7 @@ from airflow.providers.google.cloud.links.datacatalog import (
     DataCatalogTagTemplateLink,
 )
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
+from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
 
 if TYPE_CHECKING:
     from google.api_core.retry import Retry
@@ -104,7 +105,7 @@ class CloudDataCatalogCreateEntryOperator(GoogleCloudBaseOperator):
         entry_group: str,
         entry_id: str,
         entry: dict | Entry,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         retry: Retry | _MethodDefault = DEFAULT,
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
@@ -220,7 +221,7 @@ class CloudDataCatalogCreateEntryGroupOperator(GoogleCloudBaseOperator):
         location: str,
         entry_group_id: str,
         entry_group: dict | EntryGroup,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         retry: Retry | _MethodDefault = DEFAULT,
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
@@ -336,7 +337,7 @@ class CloudDataCatalogCreateTagOperator(GoogleCloudBaseOperator):
         entry: str,
         tag: dict | Tag,
         template_id: str | None = None,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         retry: Retry | _MethodDefault = DEFAULT,
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
@@ -467,7 +468,7 @@ class CloudDataCatalogCreateTagTemplateOperator(GoogleCloudBaseOperator):
         location: str,
         tag_template_id: str,
         tag_template: dict | TagTemplate,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         retry: Retry | _MethodDefault = DEFAULT,
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
@@ -583,7 +584,7 @@ class CloudDataCatalogCreateTagTemplateFieldOperator(GoogleCloudBaseOperator):
         tag_template: str,
         tag_template_field_id: str,
         tag_template_field: dict | TagTemplateField,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         retry: Retry | _MethodDefault = DEFAULT,
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
@@ -690,7 +691,7 @@ class CloudDataCatalogDeleteEntryOperator(GoogleCloudBaseOperator):
         location: str,
         entry_group: str,
         entry: str,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         retry: Retry | _MethodDefault = DEFAULT,
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
@@ -774,7 +775,7 @@ class CloudDataCatalogDeleteEntryGroupOperator(GoogleCloudBaseOperator):
         *,
         location: str,
         entry_group: str,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         retry: Retry | _MethodDefault = DEFAULT,
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
@@ -860,7 +861,7 @@ class CloudDataCatalogDeleteTagOperator(GoogleCloudBaseOperator):
         entry_group: str,
         entry: str,
         tag: str,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         retry: Retry | _MethodDefault = DEFAULT,
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
@@ -949,7 +950,7 @@ class CloudDataCatalogDeleteTagTemplateOperator(GoogleCloudBaseOperator):
         location: str,
         tag_template: str,
         force: bool,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         retry: Retry | _MethodDefault = DEFAULT,
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
@@ -1037,7 +1038,7 @@ class CloudDataCatalogDeleteTagTemplateFieldOperator(GoogleCloudBaseOperator):
         tag_template: str,
         field: str,
         force: bool,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         retry: Retry | _MethodDefault = DEFAULT,
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
@@ -1125,7 +1126,7 @@ class CloudDataCatalogGetEntryOperator(GoogleCloudBaseOperator):
         location: str,
         entry_group: str,
         entry: str,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         retry: Retry | _MethodDefault = DEFAULT,
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
@@ -1220,7 +1221,7 @@ class CloudDataCatalogGetEntryGroupOperator(GoogleCloudBaseOperator):
         location: str,
         entry_group: str,
         read_mask: FieldMask,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         retry: Retry | _MethodDefault = DEFAULT,
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
@@ -1308,7 +1309,7 @@ class CloudDataCatalogGetTagTemplateOperator(GoogleCloudBaseOperator):
         *,
         location: str,
         tag_template: str,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         retry: Retry | _MethodDefault = DEFAULT,
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
@@ -1403,7 +1404,7 @@ class CloudDataCatalogListTagsOperator(GoogleCloudBaseOperator):
         entry_group: str,
         entry: str,
         page_size: int = 100,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         retry: Retry | _MethodDefault = DEFAULT,
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
@@ -1497,7 +1498,7 @@ class CloudDataCatalogLookupEntryOperator(GoogleCloudBaseOperator):
         *,
         linked_resource: str | None = None,
         sql_resource: str | None = None,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         retry: Retry | _MethodDefault = DEFAULT,
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
@@ -1593,7 +1594,7 @@ class CloudDataCatalogRenameTagTemplateFieldOperator(GoogleCloudBaseOperator):
         tag_template: str,
         field: str,
         new_tag_template_field_id: str,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         retry: Retry | _MethodDefault = DEFAULT,
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
@@ -1808,7 +1809,7 @@ class CloudDataCatalogUpdateEntryOperator(GoogleCloudBaseOperator):
         location: str | None = None,
         entry_group: str | None = None,
         entry_id: str | None = None,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         retry: Retry | _MethodDefault = DEFAULT,
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
@@ -1920,7 +1921,7 @@ class CloudDataCatalogUpdateTagOperator(GoogleCloudBaseOperator):
         entry_group: str | None = None,
         entry: str | None = None,
         tag_id: str | None = None,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         retry: Retry | _MethodDefault = DEFAULT,
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
@@ -2034,7 +2035,7 @@ class CloudDataCatalogUpdateTagTemplateOperator(GoogleCloudBaseOperator):
         update_mask: dict | FieldMask,
         location: str | None = None,
         tag_template_id: str | None = None,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         retry: Retry | _MethodDefault = DEFAULT,
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),
@@ -2151,7 +2152,7 @@ class CloudDataCatalogUpdateTagTemplateFieldOperator(GoogleCloudBaseOperator):
         location: str | None = None,
         tag_template: str | None = None,
         tag_template_field_id: str | None = None,
-        project_id: str | None = None,
+        project_id: str = PROVIDE_PROJECT_ID,
         retry: Retry | _MethodDefault = DEFAULT,
         timeout: float | None = None,
         metadata: Sequence[tuple[str, str]] = (),