#
# 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 Google Storage Transfer Service Hook."""
import json
import logging
import time
import warnings
from copy import deepcopy
from datetime import timedelta
from typing import List, Optional, Sequence, Set, Union
from googleapiclient.discovery import Resource, build
from googleapiclient.errors import HttpError
from airflow.exceptions import AirflowException
from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
[docs]log = logging.getLogger(__name__)
# Time to sleep between active checks of the operation results
[docs]TIME_TO_SLEEP_IN_SECONDS = 10
[docs]class GcpTransferJobsStatus:
"""Class with Google Cloud Transfer jobs statuses."""
[docs]class GcpTransferOperationStatus:
"""Class with Google Cloud Transfer operations statuses."""
[docs] IN_PROGRESS = "IN_PROGRESS"
# A list of keywords used to build a request or response
[docs]ACCESS_KEY_ID = "accessKeyId"
[docs]ALREADY_EXISTING_IN_SINK = "overwriteObjectsAlreadyExistingInSink"
[docs]AWS_ACCESS_KEY = "awsAccessKey"
[docs]AWS_S3_DATA_SOURCE = 'awsS3DataSource'
[docs]BUCKET_NAME = 'bucketName'
[docs]DESCRIPTION = "description"
[docs]FILTER_JOB_NAMES = 'job_names'
[docs]FILTER_PROJECT_ID = 'project_id'
[docs]GCS_DATA_SINK = 'gcsDataSink'
[docs]GCS_DATA_SOURCE = 'gcsDataSource'
[docs]HTTP_DATA_SOURCE = 'httpDataSource'
[docs]OBJECT_CONDITIONS = 'object_conditions'
[docs]OPERATIONS = 'operations'
[docs]PROJECT_ID = 'projectId'
[docs]SCHEDULE_END_DATE = 'scheduleEndDate'
[docs]SCHEDULE_START_DATE = 'scheduleStartDate'
[docs]SECRET_ACCESS_KEY = "secretAccessKey"
[docs]START_TIME_OF_DAY = 'startTimeOfDay'
[docs]TRANSFER_JOB = 'transfer_job'
[docs]TRANSFER_JOBS = 'transferJobs'
[docs]TRANSFER_JOB_FIELD_MASK = 'update_transfer_job_field_mask'
[docs]TRANSFER_OPERATIONS = 'transferOperations'
[docs]TRANSFER_OPTIONS = 'transfer_options'
[docs]TRANSFER_SPEC = 'transferSpec'
[docs]ALREADY_EXIST_CODE = 409
[docs]NEGATIVE_STATUSES = {GcpTransferOperationStatus.FAILED, GcpTransferOperationStatus.ABORTED}
[docs]def gen_job_name(job_name: str) -> str:
"""
Adds unique suffix to job name. If suffix already exists, updates it.
Suffix — current timestamp
:param job_name:
:return: job_name with suffix
"""
uniq = int(time.time())
return f"{job_name}_{uniq}"
[docs]class CloudDataTransferServiceHook(GoogleBaseHook):
"""
Hook for Google Storage Transfer Service.
All the methods in the hook where project_id is used must be called with
keyword arguments rather than positional.
"""
def __init__(
self,
api_version: str = 'v1',
gcp_conn_id: str = "google_cloud_default",
delegate_to: Optional[str] = None,
impersonation_chain: Optional[Union[str, Sequence[str]]] = None,
) -> None:
super().__init__(
gcp_conn_id=gcp_conn_id,
delegate_to=delegate_to,
impersonation_chain=impersonation_chain,
)
self.api_version = api_version
self._conn = None
[docs] def get_conn(self) -> Resource:
"""
Retrieves connection to Google Storage Transfer service.
:return: Google Storage Transfer service object
:rtype: dict
"""
if not self._conn:
http_authorized = self._authorize()
self._conn = build(
'storagetransfer', self.api_version, http=http_authorized, cache_discovery=False
)
return self._conn
[docs] def create_transfer_job(self, body: dict) -> dict:
"""
Creates a transfer job that runs periodically.
:param body: (Required) A request body, as described in
https://cloud.google.com/storage-transfer/docs/reference/rest/v1/transferJobs/patch#request-body
:return: transfer job.
See:
https://cloud.google.com/storage-transfer/docs/reference/rest/v1/transferJobs#TransferJob
:rtype: dict
"""
body = self._inject_project_id(body, BODY, PROJECT_ID)
try:
transfer_job = (
self.get_conn().transferJobs().create(body=body).execute(num_retries=self.num_retries)
)
except HttpError as e:
# If status code "Conflict"
# https://cloud.google.com/storage-transfer/docs/reference/rest/v1/transferOperations#Code.ENUM_VALUES.ALREADY_EXISTS
# we should try to find this job
job_name = body.get(JOB_NAME, "")
if int(e.resp.status) == ALREADY_EXIST_CODE and job_name:
transfer_job = self.get_transfer_job(job_name=job_name, project_id=body.get(PROJECT_ID))
# Generate new job_name, if jobs status is deleted
# and try to create this job again
if transfer_job.get(STATUS) == GcpTransferJobsStatus.DELETED:
body[JOB_NAME] = gen_job_name(job_name)
self.log.info(
"Job `%s` has been soft deleted. Creating job with new name `%s`",
job_name,
{body[JOB_NAME]},
)
return (
self.get_conn().transferJobs().create(body=body).execute(num_retries=self.num_retries)
)
elif transfer_job.get(STATUS) == GcpTransferJobsStatus.DISABLED:
return self.enable_transfer_job(job_name=job_name, project_id=body.get(PROJECT_ID))
else:
raise e
self.log.info("Created job %s", transfer_job[NAME])
return transfer_job
@GoogleBaseHook.fallback_to_default_project_id
[docs] def get_transfer_job(self, job_name: str, project_id: str) -> dict:
"""
Gets the latest state of a long-running operation in Google Storage
Transfer Service.
:param job_name: (Required) Name of the job to be fetched
:param project_id: (Optional) the ID of the project that owns the Transfer
Job. If set to None or missing, the default project_id from the Google Cloud
connection is used.
:return: Transfer Job
:rtype: dict
"""
return (
self.get_conn()
.transferJobs()
.get(jobName=job_name, projectId=project_id)
.execute(num_retries=self.num_retries)
)
[docs] def list_transfer_job(self, request_filter: Optional[dict] = None, **kwargs) -> List[dict]:
"""
Lists long-running operations in Google Storage Transfer
Service that match the specified filter.
:param request_filter: (Required) A request filter, as described in
https://cloud.google.com/storage-transfer/docs/reference/rest/v1/transferJobs/list#body.QUERY_PARAMETERS.filter
:return: List of Transfer Jobs
:rtype: list[dict]
"""
# To preserve backward compatibility
# TODO: remove one day
if request_filter is None:
if 'filter' in kwargs:
request_filter = kwargs['filter']
if not isinstance(request_filter, dict):
raise ValueError(f"The request_filter should be dict and is {type(request_filter)}")
warnings.warn("Use 'request_filter' instead of 'filter'", DeprecationWarning)
else:
raise TypeError("list_transfer_job missing 1 required positional argument: 'request_filter'")
conn = self.get_conn()
request_filter = self._inject_project_id(request_filter, FILTER, FILTER_PROJECT_ID)
request = conn.transferJobs().list(filter=json.dumps(request_filter))
jobs: List[dict] = []
while request is not None:
response = request.execute(num_retries=self.num_retries)
jobs.extend(response[TRANSFER_JOBS])
request = conn.transferJobs().list_next(previous_request=request, previous_response=response)
return jobs
@GoogleBaseHook.fallback_to_default_project_id
[docs] def enable_transfer_job(self, job_name: str, project_id: str) -> dict:
"""
New transfers will be performed based on the schedule.
:param job_name: (Required) Name of the job to be updated
:param project_id: (Optional) the ID of the project that owns the Transfer
Job. If set to None or missing, the default project_id from the Google Cloud
connection is used.
:return: If successful, TransferJob.
:rtype: dict
"""
return (
self.get_conn()
.transferJobs()
.patch(
jobName=job_name,
body={
PROJECT_ID: project_id,
TRANSFER_JOB: {STATUS1: GcpTransferJobsStatus.ENABLED},
TRANSFER_JOB_FIELD_MASK: STATUS1,
},
)
.execute(num_retries=self.num_retries)
)
[docs] def update_transfer_job(self, job_name: str, body: dict) -> dict:
"""
Updates a transfer job that runs periodically.
:param job_name: (Required) Name of the job to be updated
:param body: A request body, as described in
https://cloud.google.com/storage-transfer/docs/reference/rest/v1/transferJobs/patch#request-body
:return: If successful, TransferJob.
:rtype: dict
"""
body = self._inject_project_id(body, BODY, PROJECT_ID)
return (
self.get_conn()
.transferJobs()
.patch(jobName=job_name, body=body)
.execute(num_retries=self.num_retries)
)
@GoogleBaseHook.fallback_to_default_project_id
[docs] def delete_transfer_job(self, job_name: str, project_id: str) -> None:
"""
Deletes a transfer job. This is a soft delete. After a transfer job is
deleted, the job and all the transfer executions are subject to garbage
collection. Transfer jobs become eligible for garbage collection
30 days after soft delete.
:param job_name: (Required) Name of the job to be deleted
:param project_id: (Optional) the ID of the project that owns the Transfer
Job. If set to None or missing, the default project_id from the Google Cloud
connection is used.
:rtype: None
"""
(
self.get_conn()
.transferJobs()
.patch(
jobName=job_name,
body={
PROJECT_ID: project_id,
TRANSFER_JOB: {STATUS1: GcpTransferJobsStatus.DELETED},
TRANSFER_JOB_FIELD_MASK: STATUS1,
},
)
.execute(num_retries=self.num_retries)
)
[docs] def cancel_transfer_operation(self, operation_name: str) -> None:
"""
Cancels an transfer operation in Google Storage Transfer Service.
:param operation_name: Name of the transfer operation.
:rtype: None
"""
self.get_conn().transferOperations().cancel(name=operation_name).execute(num_retries=self.num_retries)
[docs] def get_transfer_operation(self, operation_name: str) -> dict:
"""
Gets an transfer operation in Google Storage Transfer Service.
:param operation_name: (Required) Name of the transfer operation.
:return: transfer operation
See:
https://cloud.google.com/storage-transfer/docs/reference/rest/v1/Operation
:rtype: dict
"""
return (
self.get_conn()
.transferOperations()
.get(name=operation_name)
.execute(num_retries=self.num_retries)
)
[docs] def list_transfer_operations(self, request_filter: Optional[dict] = None, **kwargs) -> List[dict]:
"""
Gets an transfer operation in Google Storage Transfer Service.
:param request_filter: (Required) A request filter, as described in
https://cloud.google.com/storage-transfer/docs/reference/rest/v1/transferJobs/list#body.QUERY_PARAMETERS.filter
With one additional improvement:
* project_id is optional if you have a project id defined
in the connection
See: :doc:`/connections/gcp`
:return: transfer operation
:rtype: list[dict]
"""
# To preserve backward compatibility
# TODO: remove one day
if request_filter is None:
if 'filter' in kwargs:
request_filter = kwargs['filter']
if not isinstance(request_filter, dict):
raise ValueError(f"The request_filter should be dict and is {type(request_filter)}")
warnings.warn("Use 'request_filter' instead of 'filter'", DeprecationWarning)
else:
raise TypeError(
"list_transfer_operations missing 1 required positional argument: 'request_filter'"
)
conn = self.get_conn()
request_filter = self._inject_project_id(request_filter, FILTER, FILTER_PROJECT_ID)
operations: List[dict] = []
request = conn.transferOperations().list(name=TRANSFER_OPERATIONS, filter=json.dumps(request_filter))
while request is not None:
response = request.execute(num_retries=self.num_retries)
if OPERATIONS in response:
operations.extend(response[OPERATIONS])
request = conn.transferOperations().list_next(
previous_request=request, previous_response=response
)
return operations
[docs] def pause_transfer_operation(self, operation_name: str) -> None:
"""
Pauses an transfer operation in Google Storage Transfer Service.
:param operation_name: (Required) Name of the transfer operation.
:rtype: None
"""
self.get_conn().transferOperations().pause(name=operation_name).execute(num_retries=self.num_retries)
[docs] def resume_transfer_operation(self, operation_name: str) -> None:
"""
Resumes an transfer operation in Google Storage Transfer Service.
:param operation_name: (Required) Name of the transfer operation.
:rtype: None
"""
self.get_conn().transferOperations().resume(name=operation_name).execute(num_retries=self.num_retries)
[docs] def wait_for_transfer_job(
self,
job: dict,
expected_statuses: Optional[Set[str]] = None,
timeout: Optional[Union[float, timedelta]] = None,
) -> None:
"""
Waits until the job reaches the expected state.
:param job: Transfer job
See:
https://cloud.google.com/storage-transfer/docs/reference/rest/v1/transferJobs#TransferJob
:param expected_statuses: State that is expected
See:
https://cloud.google.com/storage-transfer/docs/reference/rest/v1/transferOperations#Status
:param timeout: Time in which the operation must end in seconds. If not specified, defaults to 60
seconds.
:rtype: None
"""
expected_statuses = (
{GcpTransferOperationStatus.SUCCESS} if not expected_statuses else expected_statuses
)
if timeout is None:
timeout = 60
elif isinstance(timeout, timedelta):
timeout = timeout.total_seconds()
start_time = time.monotonic()
while time.monotonic() - start_time < timeout:
request_filter = {FILTER_PROJECT_ID: job[PROJECT_ID], FILTER_JOB_NAMES: [job[NAME]]}
operations = self.list_transfer_operations(request_filter=request_filter)
for operation in operations:
self.log.info("Progress for operation %s: %s", operation[NAME], operation[METADATA][COUNTERS])
if self.operations_contain_expected_statuses(operations, expected_statuses):
return
time.sleep(TIME_TO_SLEEP_IN_SECONDS)
raise AirflowException("Timeout. The operation could not be completed within the allotted time.")
def _inject_project_id(self, body: dict, param_name: str, target_key: str) -> dict:
body = deepcopy(body)
body[target_key] = body.get(target_key, self.project_id)
if not body.get(target_key):
raise AirflowException(
f"The project id must be passed either as `{target_key}` key in `{param_name}` "
f"parameter or as project_id extra in Google Cloud connection definition. Both are not set!"
)
return body
@staticmethod
[docs] def operations_contain_expected_statuses(
operations: List[dict], expected_statuses: Union[Set[str], str]
) -> bool:
"""
Checks whether the operation list has an operation with the
expected status, then returns true
If it encounters operations in FAILED or ABORTED state
throw :class:`airflow.exceptions.AirflowException`.
:param operations: (Required) List of transfer operations to check.
:param expected_statuses: (Required) status that is expected
See:
https://cloud.google.com/storage-transfer/docs/reference/rest/v1/transferOperations#Status
:return: If there is an operation with the expected state
in the operation list, returns true,
:raises: airflow.exceptions.AirflowException If it encounters operations
with a state in the list,
:rtype: bool
"""
expected_statuses_set = (
{expected_statuses} if isinstance(expected_statuses, str) else set(expected_statuses)
)
if not operations:
return False
current_statuses = {operation[METADATA][STATUS] for operation in operations}
if len(current_statuses - expected_statuses_set) != len(current_statuses):
return True
if len(NEGATIVE_STATUSES - current_statuses) != len(NEGATIVE_STATUSES):
raise AirflowException(
f"An unexpected operation status was encountered. "
f"Expected: {', '.join(expected_statuses_set)}"
)
return False