You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by po...@apache.org on 2023/09/11 20:42:30 UTC
[airflow] branch main updated: Add explicit support of stream (realtime) pipelines for CloudDataFusionStartPipelineOperator (#34271)
This is an automated email from the ASF dual-hosted git repository.
potiuk pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 4dcdc34996 Add explicit support of stream (realtime) pipelines for CloudDataFusionStartPipelineOperator (#34271)
4dcdc34996 is described below
commit 4dcdc349964647ade80702e3d34bbf819ddf7661
Author: max <42...@users.noreply.github.com>
AuthorDate: Mon Sep 11 22:42:22 2023 +0200
Add explicit support of stream (realtime) pipelines for CloudDataFusionStartPipelineOperator (#34271)
---
airflow/providers/google/cloud/hooks/datafusion.py | 48 ++++++++++--
.../providers/google/cloud/operators/datafusion.py | 7 ++
.../providers/google/cloud/triggers/datafusion.py | 6 ++
airflow/providers/google/cloud/utils/datafusion.py | 33 ++++++++
.../google/cloud/hooks/test_datafusion.py | 90 +++++++++++++++++++++-
.../google/cloud/operators/test_datafusion.py | 5 ++
.../google/cloud/triggers/test_datafusion.py | 3 +
.../google/cloud/utils/test_datafusion.py | 37 +++++++++
8 files changed, 219 insertions(+), 10 deletions(-)
diff --git a/airflow/providers/google/cloud/hooks/datafusion.py b/airflow/providers/google/cloud/hooks/datafusion.py
index 54ae768f15..b3cfd6ee9b 100644
--- a/airflow/providers/google/cloud/hooks/datafusion.py
+++ b/airflow/providers/google/cloud/hooks/datafusion.py
@@ -31,6 +31,7 @@ from google.api_core.retry import exponential_sleep_generator
from googleapiclient.discovery import Resource, build
from airflow.exceptions import AirflowException, AirflowNotFoundException
+from airflow.providers.google.cloud.utils.datafusion import DataFusionPipelineType
from airflow.providers.google.common.hooks.base_google import (
PROVIDE_PROJECT_ID,
GoogleBaseAsyncHook,
@@ -105,6 +106,7 @@ class DataFusionHook(GoogleBaseHook):
pipeline_name: str,
pipeline_id: str,
instance_url: str,
+ pipeline_type: DataFusionPipelineType = DataFusionPipelineType.BATCH,
namespace: str = "default",
success_states: list[str] | None = None,
failure_states: list[str] | None = None,
@@ -120,6 +122,7 @@ class DataFusionHook(GoogleBaseHook):
workflow = self.get_pipeline_workflow(
pipeline_name=pipeline_name,
pipeline_id=pipeline_id,
+ pipeline_type=pipeline_type,
instance_url=instance_url,
namespace=namespace,
)
@@ -432,13 +435,14 @@ class DataFusionHook(GoogleBaseHook):
pipeline_name: str,
instance_url: str,
pipeline_id: str,
+ pipeline_type: DataFusionPipelineType = DataFusionPipelineType.BATCH,
namespace: str = "default",
) -> Any:
url = os.path.join(
self._base_url(instance_url, namespace),
quote(pipeline_name),
- "workflows",
- "DataPipelineWorkflow",
+ f"{self.cdap_program_type(pipeline_type=pipeline_type)}s",
+ self.cdap_program_id(pipeline_type=pipeline_type),
"runs",
quote(pipeline_id),
)
@@ -453,6 +457,7 @@ class DataFusionHook(GoogleBaseHook):
self,
pipeline_name: str,
instance_url: str,
+ pipeline_type: DataFusionPipelineType = DataFusionPipelineType.BATCH,
namespace: str = "default",
runtime_args: dict[str, Any] | None = None,
) -> str:
@@ -460,6 +465,7 @@ class DataFusionHook(GoogleBaseHook):
Starts a Cloud Data Fusion pipeline. Works for both batch and stream pipelines.
:param pipeline_name: Your pipeline name.
+ :param pipeline_type: Optional pipeline type (BATCH by default).
:param instance_url: Endpoint on which the REST APIs is accessible for the instance.
:param runtime_args: Optional runtime JSON args to be passed to the pipeline
:param namespace: if your pipeline belongs to a Basic edition instance, the namespace ID
@@ -480,9 +486,9 @@ class DataFusionHook(GoogleBaseHook):
body = [
{
"appId": pipeline_name,
- "programType": "workflow",
- "programId": "DataPipelineWorkflow",
"runtimeargs": runtime_args,
+ "programType": self.cdap_program_type(pipeline_type=pipeline_type),
+ "programId": self.cdap_program_id(pipeline_type=pipeline_type),
}
]
response = self._cdap_request(url=url, method="POST", body=body)
@@ -514,6 +520,30 @@ class DataFusionHook(GoogleBaseHook):
response, f"Stopping a pipeline failed with code {response.status}"
)
+ @staticmethod
+ def cdap_program_type(pipeline_type: DataFusionPipelineType) -> str:
+ """Retrieves CDAP Program type depending on the pipeline type.
+
+ :param pipeline_type: Pipeline type.
+ """
+ program_types = {
+ DataFusionPipelineType.BATCH: "workflow",
+ DataFusionPipelineType.STREAM: "spark",
+ }
+ return program_types.get(pipeline_type, "")
+
+ @staticmethod
+ def cdap_program_id(pipeline_type: DataFusionPipelineType) -> str:
+ """Retrieves CDAP Program id depending on the pipeline type.
+
+ :param pipeline_type: Pipeline type.
+ """
+ program_ids = {
+ DataFusionPipelineType.BATCH: "DataPipelineWorkflow",
+ DataFusionPipelineType.STREAM: "DataStreamsSparkStreaming",
+ }
+ return program_ids.get(pipeline_type, "")
+
class DataFusionAsyncHook(GoogleBaseAsyncHook):
"""Class to get asynchronous hook for DataFusion."""
@@ -561,10 +591,13 @@ class DataFusionAsyncHook(GoogleBaseAsyncHook):
pipeline_name: str,
pipeline_id: str,
session,
+ pipeline_type: DataFusionPipelineType = DataFusionPipelineType.BATCH,
):
+ program_type = self.sync_hook_class.cdap_program_type(pipeline_type=pipeline_type)
+ program_id = self.sync_hook_class.cdap_program_id(pipeline_type=pipeline_type)
base_url_link = self._base_url(instance_url, namespace)
url = urljoin(
- base_url_link, f"{quote(pipeline_name)}/workflows/DataPipelineWorkflow/runs/{quote(pipeline_id)}"
+ base_url_link, f"{quote(pipeline_name)}/{program_type}s/{program_id}/runs/{quote(pipeline_id)}"
)
return await self._get_link(url=url, session=session)
@@ -573,6 +606,7 @@ class DataFusionAsyncHook(GoogleBaseAsyncHook):
pipeline_name: str,
instance_url: str,
pipeline_id: str,
+ pipeline_type: DataFusionPipelineType = DataFusionPipelineType.BATCH,
namespace: str = "default",
success_states: list[str] | None = None,
) -> str:
@@ -581,7 +615,8 @@ class DataFusionAsyncHook(GoogleBaseAsyncHook):
:param pipeline_name: Your pipeline name.
:param instance_url: Endpoint on which the REST APIs is accessible for the instance.
- :param pipeline_id: Unique pipeline ID associated with specific pipeline
+ :param pipeline_id: Unique pipeline ID associated with specific pipeline.
+ :param pipeline_type: Optional pipeline type (by default batch).
:param namespace: if your pipeline belongs to a Basic edition instance, the namespace ID
is always default. If your pipeline belongs to an Enterprise edition instance, you
can create a namespace.
@@ -596,6 +631,7 @@ class DataFusionAsyncHook(GoogleBaseAsyncHook):
namespace=namespace,
pipeline_name=pipeline_name,
pipeline_id=pipeline_id,
+ pipeline_type=pipeline_type,
session=session,
)
pipeline = await pipeline.json(content_type=None)
diff --git a/airflow/providers/google/cloud/operators/datafusion.py b/airflow/providers/google/cloud/operators/datafusion.py
index d82a270871..d63de8442d 100644
--- a/airflow/providers/google/cloud/operators/datafusion.py
+++ b/airflow/providers/google/cloud/operators/datafusion.py
@@ -33,6 +33,7 @@ from airflow.providers.google.cloud.links.datafusion import (
)
from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
from airflow.providers.google.cloud.triggers.datafusion import DataFusionStartPipelineTrigger
+from airflow.providers.google.cloud.utils.datafusion import DataFusionPipelineType
if TYPE_CHECKING:
from airflow.utils.context import Context
@@ -708,6 +709,7 @@ class CloudDataFusionStartPipelineOperator(GoogleCloudBaseOperator):
:ref:`howto/operator:CloudDataFusionStartPipelineOperator`
:param pipeline_name: Your pipeline name.
+ :param pipeline_type: Optional pipeline type (BATCH by default).
:param instance_name: The name of the instance.
:param success_states: If provided the operator will wait for pipeline to be in one of
the provided states.
@@ -752,6 +754,7 @@ class CloudDataFusionStartPipelineOperator(GoogleCloudBaseOperator):
pipeline_name: str,
instance_name: str,
location: str,
+ pipeline_type: DataFusionPipelineType = DataFusionPipelineType.BATCH,
runtime_args: dict[str, Any] | None = None,
success_states: list[str] | None = None,
namespace: str = "default",
@@ -767,6 +770,7 @@ class CloudDataFusionStartPipelineOperator(GoogleCloudBaseOperator):
) -> None:
super().__init__(**kwargs)
self.pipeline_name = pipeline_name
+ self.pipeline_type = pipeline_type
self.runtime_args = runtime_args
self.namespace = namespace
self.instance_name = instance_name
@@ -800,6 +804,7 @@ class CloudDataFusionStartPipelineOperator(GoogleCloudBaseOperator):
api_url = instance["apiEndpoint"]
pipeline_id = hook.start_pipeline(
pipeline_name=self.pipeline_name,
+ pipeline_type=self.pipeline_type,
instance_url=api_url,
namespace=self.namespace,
runtime_args=self.runtime_args,
@@ -824,6 +829,7 @@ class CloudDataFusionStartPipelineOperator(GoogleCloudBaseOperator):
instance_url=api_url,
namespace=self.namespace,
pipeline_name=self.pipeline_name,
+ pipeline_type=self.pipeline_type.value,
pipeline_id=pipeline_id,
poll_interval=self.poll_interval,
gcp_conn_id=self.gcp_conn_id,
@@ -839,6 +845,7 @@ class CloudDataFusionStartPipelineOperator(GoogleCloudBaseOperator):
success_states=self.success_states,
pipeline_id=pipeline_id,
pipeline_name=self.pipeline_name,
+ pipeline_type=self.pipeline_type,
namespace=self.namespace,
instance_url=api_url,
timeout=self.pipeline_timeout,
diff --git a/airflow/providers/google/cloud/triggers/datafusion.py b/airflow/providers/google/cloud/triggers/datafusion.py
index 06bf5e053e..d1419bfbc5 100644
--- a/airflow/providers/google/cloud/triggers/datafusion.py
+++ b/airflow/providers/google/cloud/triggers/datafusion.py
@@ -20,6 +20,7 @@ import asyncio
from typing import Any, AsyncIterator, Sequence
from airflow.providers.google.cloud.hooks.datafusion import DataFusionAsyncHook
+from airflow.providers.google.cloud.utils.datafusion import DataFusionPipelineType
from airflow.triggers.base import BaseTrigger, TriggerEvent
@@ -30,6 +31,7 @@ class DataFusionStartPipelineTrigger(BaseTrigger):
:param pipeline_name: Your pipeline name.
:param instance_url: Endpoint on which the REST APIs is accessible for the instance.
:param pipeline_id: Unique pipeline ID associated with specific pipeline
+ :param pipeline_type: Your pipeline type.
:param namespace: if your pipeline belongs to a Basic edition instance, the namespace ID
is always default. If your pipeline belongs to an Enterprise edition instance, you
can create a namespace.
@@ -51,6 +53,7 @@ class DataFusionStartPipelineTrigger(BaseTrigger):
namespace: str,
pipeline_name: str,
pipeline_id: str,
+ pipeline_type: str,
poll_interval: float = 3.0,
gcp_conn_id: str = "google_cloud_default",
impersonation_chain: str | Sequence[str] | None = None,
@@ -61,6 +64,7 @@ class DataFusionStartPipelineTrigger(BaseTrigger):
self.namespace = namespace
self.pipeline_name = pipeline_name
self.pipeline_id = pipeline_id
+ self.pipeline_type = pipeline_type
self.poll_interval = poll_interval
self.gcp_conn_id = gcp_conn_id
self.impersonation_chain = impersonation_chain
@@ -76,6 +80,7 @@ class DataFusionStartPipelineTrigger(BaseTrigger):
"namespace": self.namespace,
"pipeline_name": self.pipeline_name,
"pipeline_id": self.pipeline_id,
+ "pipeline_type": self.pipeline_type,
"success_states": self.success_states,
},
)
@@ -92,6 +97,7 @@ class DataFusionStartPipelineTrigger(BaseTrigger):
namespace=self.namespace,
pipeline_name=self.pipeline_name,
pipeline_id=self.pipeline_id,
+ pipeline_type=DataFusionPipelineType.from_str(self.pipeline_type),
)
if response_from_hook == "success":
yield TriggerEvent(
diff --git a/airflow/providers/google/cloud/utils/datafusion.py b/airflow/providers/google/cloud/utils/datafusion.py
new file mode 100644
index 0000000000..20101392a9
--- /dev/null
+++ b/airflow/providers/google/cloud/utils/datafusion.py
@@ -0,0 +1,33 @@
+# 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.
+from __future__ import annotations
+
+from enum import Enum
+
+
+class DataFusionPipelineType(Enum):
+ """Enum for Data Fusion pipeline types."""
+
+ BATCH = "batch"
+ STREAM = "stream"
+
+ @staticmethod
+ def from_str(value: str) -> DataFusionPipelineType:
+ value_to_item = {item.value: item for item in DataFusionPipelineType}
+ if value in value_to_item:
+ return value_to_item[value]
+ raise ValueError(f"Invalid value '{value}'. Valid values are: {[i for i in value_to_item.keys()]}")
diff --git a/tests/providers/google/cloud/hooks/test_datafusion.py b/tests/providers/google/cloud/hooks/test_datafusion.py
index 9ee4b9d39e..d6815da5b8 100644
--- a/tests/providers/google/cloud/hooks/test_datafusion.py
+++ b/tests/providers/google/cloud/hooks/test_datafusion.py
@@ -26,6 +26,7 @@ from yarl import URL
from airflow import AirflowException
from airflow.providers.google.cloud.hooks.datafusion import DataFusionAsyncHook, DataFusionHook
+from airflow.providers.google.cloud.utils.datafusion import DataFusionPipelineType
from tests.providers.google.cloud.utils.base_gcp_mock import mock_base_gcp_hook_default_project_id
API_VERSION = "v1beta1"
@@ -50,6 +51,11 @@ CONSTRUCTED_PIPELINE_URL_GET = (
f"googleusercontent.com/api/v3/namespaces/{NAMESPACE}/apps/{PIPELINE_NAME}"
f"/workflows/DataPipelineWorkflow/runs/{PIPELINE_ID}"
)
+CONSTRUCTED_PIPELINE_STREAM_URL_GET = (
+ f"https://{INSTANCE_NAME}-{PROJECT_ID}-dot-eun1.datafusion."
+ f"googleusercontent.com/api/v3/namespaces/{NAMESPACE}/apps/{PIPELINE_NAME}"
+ f"/apsrkss/DataStreamsSparkStreaming/runs/{PIPELINE_ID}"
+)
class MockResponse:
@@ -350,6 +356,29 @@ class TestDataFusionHook:
url=f"{INSTANCE_URL}/v3/namespaces/default/start", method="POST", body=body
)
+ @mock.patch(HOOK_STR.format("DataFusionHook._cdap_request"))
+ def test_start_pipeline_stream(self, mock_request, hook):
+ run_id = 1234
+ mock_request.return_value = mock.MagicMock(status=200, data=f'[{{"runId":{run_id}}}]')
+
+ hook.start_pipeline(
+ pipeline_name=PIPELINE_NAME,
+ instance_url=INSTANCE_URL,
+ runtime_args=RUNTIME_ARGS,
+ pipeline_type=DataFusionPipelineType.STREAM,
+ )
+ body = [
+ {
+ "appId": PIPELINE_NAME,
+ "programType": "spark",
+ "programId": "DataStreamsSparkStreaming",
+ "runtimeargs": RUNTIME_ARGS,
+ }
+ ]
+ mock_request.assert_called_once_with(
+ url=f"{INSTANCE_URL}/v3/namespaces/default/start", method="POST", body=body
+ )
+
@mock.patch(HOOK_STR.format("DataFusionHook._cdap_request"))
def test_start_pipeline_should_fail_if_empty_data_response(self, mock_request, hook):
mock_request.return_value.status = 200
@@ -443,6 +472,22 @@ class TestDataFusionHook:
method="GET",
)
+ @mock.patch(HOOK_STR.format("DataFusionHook._cdap_request"))
+ def test_get_pipeline_workflow_stream(self, mock_request, hook):
+ run_id = 1234
+ mock_request.return_value = mock.MagicMock(status=200, data=f'[{{"runId":{run_id}}}]')
+ hook.get_pipeline_workflow(
+ pipeline_name=PIPELINE_NAME,
+ instance_url=INSTANCE_URL,
+ pipeline_id=PIPELINE_ID,
+ pipeline_type=DataFusionPipelineType.STREAM,
+ )
+ mock_request.assert_called_once_with(
+ url=f"{INSTANCE_URL}/v3/namespaces/default/apps/{PIPELINE_NAME}/"
+ f"sparks/DataStreamsSparkStreaming/runs/{PIPELINE_ID}",
+ method="GET",
+ )
+
@mock.patch(HOOK_STR.format("DataFusionHook._cdap_request"))
def test_get_pipeline_workflow_should_fail_if_empty_data_response(self, mock_request, hook):
mock_request.return_value.status = 200
@@ -474,6 +519,28 @@ class TestDataFusionHook:
method="GET",
)
+ @pytest.mark.parametrize(
+ "pipeline_type, expected_program_type",
+ [
+ (DataFusionPipelineType.BATCH, "workflow"),
+ (DataFusionPipelineType.STREAM, "spark"),
+ ("non existing value", ""),
+ ],
+ )
+ def test_cdap_program_type(self, pipeline_type, expected_program_type):
+ assert DataFusionHook.cdap_program_type(pipeline_type) == expected_program_type
+
+ @pytest.mark.parametrize(
+ "pipeline_type, expected_program_id",
+ [
+ (DataFusionPipelineType.BATCH, "DataPipelineWorkflow"),
+ (DataFusionPipelineType.STREAM, "DataStreamsSparkStreaming"),
+ ("non existing value", ""),
+ ],
+ )
+ def test_cdap_program_id(self, pipeline_type, expected_program_id):
+ assert DataFusionHook.cdap_program_id(pipeline_type) == expected_program_id
+
class TestDataFusionHookAsynch:
def test_delegate_to_runtime_error(self):
@@ -493,13 +560,20 @@ class TestDataFusionHookAsynch:
mocked_link.assert_awaited_once_with(url=CONSTRUCTED_PIPELINE_URL, session=session)
@pytest.mark.asyncio
+ @pytest.mark.parametrize(
+ "pipeline_type, constructed_url",
+ [
+ (DataFusionPipelineType.BATCH, CONSTRUCTED_PIPELINE_URL_GET),
+ (DataFusionPipelineType.STREAM, CONSTRUCTED_PIPELINE_STREAM_URL_GET),
+ ],
+ )
@mock.patch(HOOK_STR.format("DataFusionAsyncHook.get_pipeline"))
async def test_async_get_pipeline_status_completed_should_execute_successfully(
- self, mocked_get, hook_async
+ self, mocked_get, hook_async, pipeline_type, constructed_url
):
response = aiohttp.ClientResponse(
"get",
- URL(CONSTRUCTED_PIPELINE_URL_GET),
+ URL(constructed_url),
request_info=mock.Mock(),
writer=mock.Mock(),
continue100=None,
@@ -523,14 +597,21 @@ class TestDataFusionHookAsynch:
assert pipeline_status == "success"
@pytest.mark.asyncio
+ @pytest.mark.parametrize(
+ "pipeline_type, constructed_url",
+ [
+ (DataFusionPipelineType.BATCH, CONSTRUCTED_PIPELINE_URL_GET),
+ (DataFusionPipelineType.STREAM, CONSTRUCTED_PIPELINE_STREAM_URL_GET),
+ ],
+ )
@mock.patch(HOOK_STR.format("DataFusionAsyncHook.get_pipeline"))
async def test_async_get_pipeline_status_running_should_execute_successfully(
- self, mocked_get, hook_async
+ self, mocked_get, hook_async, pipeline_type, constructed_url
):
"""Assets that the DataFusionAsyncHook returns pending response when job is still in running state"""
response = aiohttp.ClientResponse(
"get",
- URL(CONSTRUCTED_PIPELINE_URL_GET),
+ URL(constructed_url),
request_info=mock.Mock(),
writer=mock.Mock(),
continue100=None,
@@ -548,6 +629,7 @@ class TestDataFusionHookAsynch:
pipeline_name=PIPELINE_NAME,
instance_url=INSTANCE_URL,
pipeline_id=PIPELINE_ID,
+ pipeline_type=pipeline_type,
namespace=NAMESPACE,
)
mocked_get.assert_awaited_once()
diff --git a/tests/providers/google/cloud/operators/test_datafusion.py b/tests/providers/google/cloud/operators/test_datafusion.py
index fd5ed2d648..988a64e9ab 100644
--- a/tests/providers/google/cloud/operators/test_datafusion.py
+++ b/tests/providers/google/cloud/operators/test_datafusion.py
@@ -36,6 +36,7 @@ from airflow.providers.google.cloud.operators.datafusion import (
CloudDataFusionUpdateInstanceOperator,
)
from airflow.providers.google.cloud.triggers.datafusion import DataFusionStartPipelineTrigger
+from airflow.providers.google.cloud.utils.datafusion import DataFusionPipelineType
HOOK_STR = "airflow.providers.google.cloud.operators.datafusion.DataFusionHook"
@@ -235,12 +236,14 @@ class TestCloudDataFusionStartPipelineOperator:
pipeline_name=PIPELINE_NAME,
namespace=NAMESPACE,
runtime_args=RUNTIME_ARGS,
+ pipeline_type=DataFusionPipelineType.BATCH,
)
mock_hook.return_value.wait_for_pipeline_state.assert_called_once_with(
success_states=[*SUCCESS_STATES, PipelineStates.RUNNING],
pipeline_id=PIPELINE_ID,
pipeline_name=PIPELINE_NAME,
+ pipeline_type=DataFusionPipelineType.BATCH,
namespace=NAMESPACE,
instance_url=INSTANCE_URL,
timeout=300,
@@ -275,6 +278,7 @@ class TestCloudDataFusionStartPipelineOperator:
pipeline_name=PIPELINE_NAME,
namespace=NAMESPACE,
runtime_args=RUNTIME_ARGS,
+ pipeline_type=DataFusionPipelineType.BATCH,
)
mock_hook.return_value.wait_for_pipeline_state.assert_not_called()
@@ -374,6 +378,7 @@ class TestCloudDataFusionStartPipelineOperatorAsynch:
pipeline_name=PIPELINE_NAME,
namespace=NAMESPACE,
runtime_args=RUNTIME_ARGS,
+ pipeline_type=DataFusionPipelineType.BATCH,
)
@mock.patch(HOOK_STR)
diff --git a/tests/providers/google/cloud/triggers/test_datafusion.py b/tests/providers/google/cloud/triggers/test_datafusion.py
index 8e85bd0504..6cc2397baf 100644
--- a/tests/providers/google/cloud/triggers/test_datafusion.py
+++ b/tests/providers/google/cloud/triggers/test_datafusion.py
@@ -36,6 +36,7 @@ PROJECT_ID = "test_project_id"
PIPELINE_NAME = "shrubberyPipeline"
PIPELINE = {"test": "pipeline"}
PIPELINE_ID = "test_pipeline_id"
+PIPELINE_TYPE = "batch"
INSTANCE_URL = "http://datafusion.instance.com"
NAMESPACE = "TEST_NAMESPACE"
RUNTIME_ARGS = {"arg1": "a", "arg2": "b"}
@@ -50,6 +51,7 @@ def trigger():
namespace=NAMESPACE,
pipeline_name=PIPELINE_NAME,
pipeline_id=PIPELINE_ID,
+ pipeline_type=PIPELINE_TYPE,
poll_interval=TEST_POLL_INTERVAL,
gcp_conn_id=TEST_GCP_PROJECT_ID,
)
@@ -68,6 +70,7 @@ class TestDataFusionStartPipelineTrigger:
"namespace": NAMESPACE,
"pipeline_name": PIPELINE_NAME,
"pipeline_id": PIPELINE_ID,
+ "pipeline_type": PIPELINE_TYPE,
"gcp_conn_id": TEST_GCP_PROJECT_ID,
"success_states": None,
}
diff --git a/tests/providers/google/cloud/utils/test_datafusion.py b/tests/providers/google/cloud/utils/test_datafusion.py
new file mode 100644
index 0000000000..ab1f3a0927
--- /dev/null
+++ b/tests/providers/google/cloud/utils/test_datafusion.py
@@ -0,0 +1,37 @@
+# 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.
+from __future__ import annotations
+
+import pytest
+
+from airflow.providers.google.cloud.utils.datafusion import DataFusionPipelineType
+
+
+class TestDataFusionPipelineType:
+ @pytest.mark.parametrize(
+ "str_value, expected_item",
+ [
+ ("batch", DataFusionPipelineType.BATCH),
+ ("stream", DataFusionPipelineType.STREAM),
+ ],
+ )
+ def test_from_str(self, str_value, expected_item):
+ assert DataFusionPipelineType.from_str(str_value) == expected_item
+
+ def test_from_str_error(self):
+ with pytest.raises(ValueError):
+ DataFusionPipelineType.from_str("non-existing value")