Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -603,15 +603,17 @@ def __init__(
super().__init__(**kwargs)
self.source_bucket = source_bucket
self.source_object = source_object
self.destination_bucket = destination_bucket or self.source_bucket
self.destination_object = destination_object or self.source_object
self.destination_bucket = destination_bucket
self.destination_object = destination_object

self.gcp_conn_id = gcp_conn_id
self.transform_script = transform_script
self.output_encoding = sys.getdefaultencoding()
self.impersonation_chain = impersonation_chain

def execute(self, context: Context) -> None:
destination_bucket = self.destination_bucket or self.source_bucket
destination_object = self.destination_object or self.source_object
hook = GCSHook(gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain)

with NamedTemporaryFile() as source_file, NamedTemporaryFile() as destination_file:
Expand All @@ -637,29 +639,31 @@ def execute(self, context: Context) -> None:

self.log.info("Transformation succeeded. Output temporarily located at %s", destination_file.name)

self.log.info("Uploading file to %s as %s", self.destination_bucket, self.destination_object)
self.log.info("Uploading file to %s as %s", destination_bucket, destination_object)
FileDetailsLink.persist(
context=context,
uri=f"{self.destination_bucket}/{self.destination_object}",
uri=f"{destination_bucket}/{destination_object}",
project_id=hook.project_id,
)
hook.upload(
bucket_name=self.destination_bucket,
object_name=self.destination_object,
bucket_name=destination_bucket,
object_name=destination_object,
filename=destination_file.name,
)

def get_openlineage_facets_on_start(self):
from airflow.providers.common.compat.openlineage.facet import Dataset
from airflow.providers.openlineage.extractors import OperatorLineage

destination_bucket = self.destination_bucket or self.source_bucket
destination_object = self.destination_object or self.source_object
input_dataset = Dataset(
namespace=f"gs://{self.source_bucket}",
name=self.source_object,
)
output_dataset = Dataset(
namespace=f"gs://{self.destination_bucket}",
name=self.destination_object,
namespace=f"gs://{destination_bucket}",
name=destination_object,
)

return OperatorLineage(inputs=[input_dataset], outputs=[output_dataset])
Expand Down
43 changes: 43 additions & 0 deletions providers/google/tests/unit/google/cloud/operators/test_gcs.py
Original file line number Diff line number Diff line change
Expand Up @@ -336,6 +336,49 @@ def test_execute(self, mock_hook, mock_subprocess, mock_tempfile):
filename=destination,
)

@mock.patch("airflow.providers.google.cloud.operators.gcs.NamedTemporaryFile")
@mock.patch("airflow.providers.google.cloud.operators.gcs.subprocess")
@mock.patch("airflow.providers.google.cloud.operators.gcs.GCSHook")
def test_execute_defaults_destination_to_source(self, mock_hook, mock_subprocess, mock_tempfile):
source_bucket = TEST_BUCKET
source_object = "test.txt"
transform_script = "script.py"

source = "source"
destination = "destination"

mock1 = mock.Mock()
mock2 = mock.Mock()
mock1.name = source
mock2.name = destination

mock_tempfile.return_value.__enter__.side_effect = [mock1, mock2]

mock_proc = mock.MagicMock()
mock_proc.returncode = 0
mock_proc.stdout.readline = lambda: b""
mock_proc.wait.return_value = None
mock_popen = mock.MagicMock()
mock_popen.return_value.__enter__.return_value = mock_proc

mock_subprocess.Popen = mock_popen
mock_subprocess.PIPE = "pipe"
mock_subprocess.STDOUT = "stdout"

op = GCSFileTransformOperator(
task_id=TASK_ID,
source_bucket=source_bucket,
source_object=source_object,
transform_script=transform_script,
)
op.execute(context=mock.MagicMock())

mock_hook.return_value.upload.assert_called_with(
bucket_name=source_bucket,
object_name=source_object,
filename=destination,
)

def test_get_openlineage_facets_on_start(self):
expected_input = Dataset(
namespace=f"gs://{TEST_BUCKET}",
Expand Down
1 change: 0 additions & 1 deletion scripts/ci/prek/validate_operators_init_exemptions.txt
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,6 @@ providers/google/src/airflow/providers/google/cloud/operators/cloud_storage_tran
providers/google/src/airflow/providers/google/cloud/operators/dataproc.py::DataprocCreateClusterOperator
providers/google/src/airflow/providers/google/cloud/operators/dataproc.py::DataprocSubmitJobOperator
providers/google/src/airflow/providers/google/cloud/operators/functions.py::CloudFunctionDeployFunctionOperator
providers/google/src/airflow/providers/google/cloud/operators/gcs.py::GCSFileTransformOperator
providers/google/src/airflow/providers/google/cloud/sensors/bigquery_dts.py::BigQueryDataTransferServiceTransferRunSensor
providers/google/src/airflow/providers/google/cloud/sensors/cloud_composer.py::CloudComposerExternalTaskSensor
providers/google/src/airflow/providers/google/cloud/transfers/azure_fileshare_to_gcs.py::AzureFileShareToGCSOperator
Expand Down
Loading