You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "phanikumv (via GitHub)" <gi...@apache.org> on 2023/02/28 07:40:50 UTC
[GitHub] [airflow] phanikumv opened a new pull request, #29801: Add deferrable `AzureDataFactoryPipelineRunStatusSensor`
phanikumv opened a new pull request, #29801:
URL: https://github.com/apache/airflow/pull/29801
This PR donates the `AzureDataFactoryPipelineRunStatusAsyncSensor` from [astronomer-providers](https://github.com/astronomer/astronomer-providers) repo to OSS
<!--
Thank you for contributing! Please make sure that your code changes
are covered with tests. And in case of new features or big changes
remember to adjust the documentation.
Feel free to ping committers for the review!
In case of an existing issue, reference it using one of the following:
closes: #ISSUE
related: #ISSUE
How to write a good git commit message:
http://chris.beams.io/posts/git-commit/
-->
---
**^ Add meaningful description above**
Read the **[Pull Request Guidelines](https://github.com/apache/airflow/blob/main/CONTRIBUTING.rst#pull-request-guidelines)** for more information.
In case of fundamental code changes, an Airflow Improvement Proposal ([AIP](https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvement+Proposals)) is needed.
In case of a new dependency, check compliance with the [ASF 3rd Party License Policy](https://www.apache.org/legal/resolved.html#category-x).
In case of backwards incompatible changes please leave a note in a newsfragment file, named `{pr_number}.significant.rst` or `{issue_number}.significant.rst`, in [newsfragments](https://github.com/apache/airflow/tree/main/newsfragments).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] phanikumv commented on a diff in pull request #29801: Add deferrable `AzureDataFactoryPipelineRunStatusSensor`
Posted by "phanikumv (via GitHub)" <gi...@apache.org>.
phanikumv commented on code in PR #29801:
URL: https://github.com/apache/airflow/pull/29801#discussion_r1125631700
##########
airflow/providers/microsoft/azure/hooks/data_factory.py:
##########
@@ -1039,3 +1048,120 @@ def test_connection(self) -> tuple[bool, str]:
return success
except Exception as e:
return False, str(e)
+
+
+def provide_targeted_factory_async(func: T) -> T:
+ """
+ Provide the targeted factory to the async decorated function in case it isn't specified.
+
+ If ``resource_group_name`` or ``factory_name`` is not provided it defaults to the value specified in
+ the connection extras.
+ """
+ signature = inspect.signature(func)
+
+ @wraps(func)
+ async def wrapper(*args: Any, **kwargs: Any) -> Any:
+ bound_args = signature.bind(*args, **kwargs)
+
+ async def bind_argument(arg: Any, default_key: str) -> None:
+ # Check if arg was not included in the function signature or, if it is, the value is not provided.
+ if arg not in bound_args.arguments or bound_args.arguments[arg] is None:
+ self = args[0]
+ conn = await sync_to_async(self.get_connection)(self.conn_id)
+ default_value = conn.extra_dejson.get(default_key)
Review Comment:
done
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] phanikumv commented on a diff in pull request #29801: Add deferrable `AzureDataFactoryPipelineRunStatusSensor`
Posted by "phanikumv (via GitHub)" <gi...@apache.org>.
phanikumv commented on code in PR #29801:
URL: https://github.com/apache/airflow/pull/29801#discussion_r1132622222
##########
airflow/providers/microsoft/azure/hooks/data_factory.py:
##########
@@ -1039,3 +1048,124 @@ def test_connection(self) -> tuple[bool, str]:
return success
except Exception as e:
return False, str(e)
+
+
+def provide_targeted_factory_async(func: T) -> T:
+ """
+ Provide the targeted factory to the async decorated function in case it isn't specified.
+
+ If ``resource_group_name`` or ``factory_name`` is not provided it defaults to the value specified in
+ the connection extras.
+ """
+ signature = inspect.signature(func)
+
+ @wraps(func)
+ async def wrapper(*args: Any, **kwargs: Any) -> Any:
+ bound_args = signature.bind(*args, **kwargs)
+
+ async def bind_argument(arg: Any, default_key: str) -> None:
+ # Check if arg was not included in the function signature or, if it is, the value is not provided.
+ if arg not in bound_args.arguments or bound_args.arguments[arg] is None:
+ self = args[0]
+ conn = await sync_to_async(self.get_connection)(self.conn_id)
+ extras = conn.extra_dejson
+ default_value = extras.get(default_key) or extras.get(
+ f"extra__azure_data_factory__{default_key}"
+ )
+ if not default_value:
+ raise AirflowException("Could not determine the targeted data factory.")
+
+ bound_args.arguments[arg] = default_value
+
+ await bind_argument("resource_group_name", "resource_group_name")
+ await bind_argument("factory_name", "factory_name")
+
+ return await func(*bound_args.args, **bound_args.kwargs)
+
+ return cast(T, wrapper)
+
+
+class AzureDataFactoryAsyncHook(AzureDataFactoryHook):
+ """
+ An Async Hook that connects to Azure DataFactory to perform pipeline operations
+
+ :param azure_data_factory_conn_id: The :ref:`Azure Data Factory connection id<howto/connection:adf>`.
+ """
+
+ def __init__(self, azure_data_factory_conn_id: str):
+ self._async_conn: AsyncDataFactoryManagementClient = None
+ self.conn_id = azure_data_factory_conn_id
+ super().__init__(azure_data_factory_conn_id=azure_data_factory_conn_id)
+
+ async def get_async_conn(self) -> AsyncDataFactoryManagementClient:
+ """Get async connection and connect to azure data factory"""
+ if self._conn is not None:
Review Comment:
correct , fixed it
##########
airflow/providers/microsoft/azure/hooks/data_factory.py:
##########
@@ -1039,3 +1048,124 @@ def test_connection(self) -> tuple[bool, str]:
return success
except Exception as e:
return False, str(e)
+
+
+def provide_targeted_factory_async(func: T) -> T:
+ """
+ Provide the targeted factory to the async decorated function in case it isn't specified.
+
+ If ``resource_group_name`` or ``factory_name`` is not provided it defaults to the value specified in
+ the connection extras.
+ """
+ signature = inspect.signature(func)
+
+ @wraps(func)
+ async def wrapper(*args: Any, **kwargs: Any) -> Any:
+ bound_args = signature.bind(*args, **kwargs)
+
+ async def bind_argument(arg: Any, default_key: str) -> None:
+ # Check if arg was not included in the function signature or, if it is, the value is not provided.
+ if arg not in bound_args.arguments or bound_args.arguments[arg] is None:
+ self = args[0]
+ conn = await sync_to_async(self.get_connection)(self.conn_id)
+ extras = conn.extra_dejson
+ default_value = extras.get(default_key) or extras.get(
+ f"extra__azure_data_factory__{default_key}"
+ )
+ if not default_value:
+ raise AirflowException("Could not determine the targeted data factory.")
+
+ bound_args.arguments[arg] = default_value
+
+ await bind_argument("resource_group_name", "resource_group_name")
+ await bind_argument("factory_name", "factory_name")
+
+ return await func(*bound_args.args, **bound_args.kwargs)
+
+ return cast(T, wrapper)
+
+
+class AzureDataFactoryAsyncHook(AzureDataFactoryHook):
+ """
+ An Async Hook that connects to Azure DataFactory to perform pipeline operations
+
+ :param azure_data_factory_conn_id: The :ref:`Azure Data Factory connection id<howto/connection:adf>`.
+ """
+
+ def __init__(self, azure_data_factory_conn_id: str):
Review Comment:
done
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] josh-fell commented on a diff in pull request #29801: Add deferrable `AzureDataFactoryPipelineRunStatusSensor`
Posted by "josh-fell (via GitHub)" <gi...@apache.org>.
josh-fell commented on code in PR #29801:
URL: https://github.com/apache/airflow/pull/29801#discussion_r1132638461
##########
airflow/providers/microsoft/azure/hooks/data_factory.py:
##########
@@ -1039,3 +1048,124 @@ def test_connection(self) -> tuple[bool, str]:
return success
except Exception as e:
return False, str(e)
+
+
+def provide_targeted_factory_async(func: T) -> T:
+ """
+ Provide the targeted factory to the async decorated function in case it isn't specified.
+
+ If ``resource_group_name`` or ``factory_name`` is not provided it defaults to the value specified in
+ the connection extras.
+ """
+ signature = inspect.signature(func)
+
+ @wraps(func)
+ async def wrapper(*args: Any, **kwargs: Any) -> Any:
+ bound_args = signature.bind(*args, **kwargs)
+
+ async def bind_argument(arg: Any, default_key: str) -> None:
+ # Check if arg was not included in the function signature or, if it is, the value is not provided.
+ if arg not in bound_args.arguments or bound_args.arguments[arg] is None:
+ self = args[0]
+ conn = await sync_to_async(self.get_connection)(self.conn_id)
+ extras = conn.extra_dejson
+ default_value = extras.get(default_key) or extras.get(
+ f"extra__azure_data_factory__{default_key}"
+ )
+ if not default_value:
+ raise AirflowException("Could not determine the targeted data factory.")
+
+ bound_args.arguments[arg] = default_value
+
+ await bind_argument("resource_group_name", "resource_group_name")
+ await bind_argument("factory_name", "factory_name")
+
+ return await func(*bound_args.args, **bound_args.kwargs)
+
+ return cast(T, wrapper)
+
+
+class AzureDataFactoryAsyncHook(AzureDataFactoryHook):
+ """
+ An Async Hook that connects to Azure DataFactory to perform pipeline operations
+
+ :param azure_data_factory_conn_id: The :ref:`Azure Data Factory connection id<howto/connection:adf>`.
+ """
+
+ def __init__(self, azure_data_factory_conn_id: str):
+ self._async_conn: AsyncDataFactoryManagementClient = None
+ self.conn_id = azure_data_factory_conn_id
+ super().__init__(azure_data_factory_conn_id=azure_data_factory_conn_id)
+
+ async def get_async_conn(self) -> AsyncDataFactoryManagementClient:
+ """Get async connection and connect to azure data factory"""
+ if self._conn is not None:
Review Comment:
`self._async_conn` isn't actually being set anywhere. Can you add that so the `is not None` check here will just return an existing connection object?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] phanikumv commented on a diff in pull request #29801: Add deferrable `AzureDataFactoryPipelineRunStatusSensor`
Posted by "phanikumv (via GitHub)" <gi...@apache.org>.
phanikumv commented on code in PR #29801:
URL: https://github.com/apache/airflow/pull/29801#discussion_r1132621625
##########
tests/providers/microsoft/azure/hooks/test_azure_data_factory.py:
##########
@@ -708,3 +724,245 @@ def test_backcompat_prefix_both_prefers_short(mock_connect):
hook = AzureDataFactoryHook("my_conn")
hook.delete_factory(factory_name="n/a")
mock_connect.return_value.factories.delete.assert_called_with("non-prefixed", "n/a")
+
+
+class TestAzureDataFactoryAsyncHook:
+ @pytest.mark.asyncio
+ @async_mock.patch(f"{MODULE}.hooks.data_factory.AzureDataFactoryAsyncHook.get_async_conn")
+ @async_mock.patch(f"{MODULE}.hooks.data_factory.AzureDataFactoryAsyncHook.get_pipeline_run")
+ async def test_get_adf_pipeline_run_status_queued(self, mock_get_pipeline_run, mock_conn):
+ """Test get_adf_pipeline_run_status function with mocked status"""
+ mock_status = "Queued"
+ mock_get_pipeline_run.return_value.status = mock_status
+ hook = AzureDataFactoryAsyncHook(AZURE_DATA_FACTORY_CONN_ID)
+ response = await hook.get_adf_pipeline_run_status(RUN_ID, RESOURCE_GROUP_NAME, DATAFACTORY_NAME)
+ assert response == mock_status
+
+ @pytest.mark.asyncio
+ @async_mock.patch(f"{MODULE}.hooks.data_factory.AzureDataFactoryAsyncHook.get_async_conn")
+ @async_mock.patch(f"{MODULE}.hooks.data_factory.AzureDataFactoryAsyncHook.get_pipeline_run")
+ async def test_get_adf_pipeline_run_status_inprogress(
+ self,
+ mock_get_pipeline_run,
+ mock_conn,
+ ):
+ """Test get_adf_pipeline_run_status function with mocked status"""
+ mock_status = "InProgress"
+ mock_get_pipeline_run.return_value.status = mock_status
+ hook = AzureDataFactoryAsyncHook(AZURE_DATA_FACTORY_CONN_ID)
+ response = await hook.get_adf_pipeline_run_status(RUN_ID, RESOURCE_GROUP_NAME, DATAFACTORY_NAME)
+ assert response == mock_status
+
+ @pytest.mark.asyncio
+ @async_mock.patch(f"{MODULE}.hooks.data_factory.AzureDataFactoryAsyncHook.get_async_conn")
+ @async_mock.patch(f"{MODULE}.hooks.data_factory.AzureDataFactoryAsyncHook.get_pipeline_run")
+ async def test_get_adf_pipeline_run_status_success(self, mock_get_pipeline_run, mock_conn):
+ """Test get_adf_pipeline_run_status function with mocked status"""
+ mock_status = "Succeeded"
+ mock_get_pipeline_run.return_value.status = mock_status
+ hook = AzureDataFactoryAsyncHook(AZURE_DATA_FACTORY_CONN_ID)
+ response = await hook.get_adf_pipeline_run_status(RUN_ID, RESOURCE_GROUP_NAME, DATAFACTORY_NAME)
+ assert response == mock_status
+
+ @pytest.mark.asyncio
+ @async_mock.patch(f"{MODULE}.hooks.data_factory.AzureDataFactoryAsyncHook.get_async_conn")
+ @async_mock.patch(f"{MODULE}.hooks.data_factory.AzureDataFactoryAsyncHook.get_pipeline_run")
+ async def test_get_adf_pipeline_run_status_failed(self, mock_get_pipeline_run, mock_conn):
+ """Test get_adf_pipeline_run_status function with mocked status"""
+ mock_status = "Failed"
+ mock_get_pipeline_run.return_value.status = mock_status
+ hook = AzureDataFactoryAsyncHook(AZURE_DATA_FACTORY_CONN_ID)
+ response = await hook.get_adf_pipeline_run_status(RUN_ID, RESOURCE_GROUP_NAME, DATAFACTORY_NAME)
+ assert response == mock_status
+
+ @pytest.mark.asyncio
+ @async_mock.patch(f"{MODULE}.hooks.data_factory.AzureDataFactoryAsyncHook.get_async_conn")
+ @async_mock.patch(f"{MODULE}.hooks.data_factory.AzureDataFactoryAsyncHook.get_pipeline_run")
+ async def test_get_adf_pipeline_run_status_cancelled(self, mock_get_pipeline_run, mock_conn):
+ """Test get_adf_pipeline_run_status function with mocked status"""
+ mock_status = "Cancelled"
+ mock_get_pipeline_run.return_value.status = mock_status
+ hook = AzureDataFactoryAsyncHook(AZURE_DATA_FACTORY_CONN_ID)
+ response = await hook.get_adf_pipeline_run_status(RUN_ID, RESOURCE_GROUP_NAME, DATAFACTORY_NAME)
+ assert response == mock_status
+
+ @pytest.mark.asyncio
+ @async_mock.patch(f"{MODULE}.hooks.data_factory.AzureDataFactoryAsyncHook.get_async_conn")
+ @async_mock.patch(f"{MODULE}.hooks.data_factory.AzureDataFactoryAsyncHook.get_pipeline_run")
+ async def test_get_adf_pipeline_run_status_exception(self, mock_get_pipeline_run, mock_conn):
+ """Test get_adf_pipeline_run_status function with exception"""
+ mock_get_pipeline_run.side_effect = Exception("Test exception")
+ hook = AzureDataFactoryAsyncHook(AZURE_DATA_FACTORY_CONN_ID)
+ with pytest.raises(AirflowException):
+ await hook.get_adf_pipeline_run_status(RUN_ID, RESOURCE_GROUP_NAME, DATAFACTORY_NAME)
+
+ @pytest.mark.asyncio
+ @async_mock.patch("azure.mgmt.datafactory.models._models_py3.PipelineRun")
+ @async_mock.patch(f"{MODULE}.hooks.data_factory.AzureDataFactoryAsyncHook.get_connection")
+ @async_mock.patch(f"{MODULE}.hooks.data_factory.AzureDataFactoryAsyncHook.get_async_conn")
+ async def test_get_pipeline_run_exception_without_resource(
+ self, mock_conn, mock_get_connection, mock_pipeline_run
+ ):
+ """
+ Test get_pipeline_run function without passing the resource name to check the decorator function and
+ raise exception
+ """
+ mock_connection = Connection(
+ extra=json.dumps({"extra__azure_data_factory__factory_name": DATAFACTORY_NAME})
+ )
+ mock_get_connection.return_value = mock_connection
+ mock_conn.return_value.__aenter__.return_value.pipeline_runs.get.return_value = mock_pipeline_run
+ hook = AzureDataFactoryAsyncHook(AZURE_DATA_FACTORY_CONN_ID)
+ with pytest.raises(AirflowException):
+ await hook.get_pipeline_run(RUN_ID, None, DATAFACTORY_NAME)
+
+ @pytest.mark.asyncio
+ @async_mock.patch(f"{MODULE}.hooks.data_factory.AzureDataFactoryAsyncHook.get_async_conn")
+ async def test_get_pipeline_run_exception(self, mock_conn):
+ """Test get_pipeline_run function with exception"""
+ mock_conn.return_value.__aenter__.return_value.pipeline_runs.get.side_effect = Exception(
+ "Test exception"
+ )
+ hook = AzureDataFactoryAsyncHook(AZURE_DATA_FACTORY_CONN_ID)
+ with pytest.raises(AirflowException):
+ await hook.get_pipeline_run(RUN_ID, RESOURCE_GROUP_NAME, DATAFACTORY_NAME)
+
+ @pytest.mark.asyncio
+ @async_mock.patch(f"{MODULE}.hooks.data_factory.AzureDataFactoryAsyncHook.get_connection")
+ async def test_get_async_conn(self, mock_connection):
+ """"""
+ mock_conn = Connection(
+ conn_id=DEFAULT_CONNECTION_CLIENT_SECRET,
+ conn_type="azure_data_factory",
+ login="clientId",
+ password="clientSecret",
+ extra=json.dumps(
+ {
+ "extra__azure_data_factory__tenantId": "tenantId",
+ "extra__azure_data_factory__subscriptionId": "subscriptionId",
+ "extra__azure_data_factory__resource_group_name": RESOURCE_GROUP_NAME,
+ "extra__azure_data_factory__factory_name": DATAFACTORY_NAME,
+ }
+ ),
+ )
+ mock_connection.return_value = mock_conn
+ hook = AzureDataFactoryAsyncHook(AZURE_DATA_FACTORY_CONN_ID)
+ response = await hook.get_async_conn()
+ assert isinstance(response, DataFactoryManagementClient)
+
+ @pytest.mark.asyncio
+ @async_mock.patch(f"{MODULE}.hooks.data_factory.AzureDataFactoryAsyncHook.get_connection")
+ async def test_get_async_conn_without_login_id(self, mock_connection):
+ """Test get_async_conn function without login id"""
+ mock_conn = Connection(
+ conn_id=DEFAULT_CONNECTION_CLIENT_SECRET,
+ conn_type="azure_data_factory",
+ extra=json.dumps(
+ {
+ "extra__azure_data_factory__tenantId": "tenantId",
+ "extra__azure_data_factory__subscriptionId": "subscriptionId",
+ "extra__azure_data_factory__resource_group_name": RESOURCE_GROUP_NAME,
+ "extra__azure_data_factory__factory_name": DATAFACTORY_NAME,
+ }
+ ),
+ )
+ mock_connection.return_value = mock_conn
+ hook = AzureDataFactoryAsyncHook(AZURE_DATA_FACTORY_CONN_ID)
+ response = await hook.get_async_conn()
+ assert isinstance(response, DataFactoryManagementClient)
+
+ @pytest.mark.asyncio
+ @pytest.mark.parametrize(
+ "mock_connection_params",
+ [
+ {
+ "extra__azure_data_factory__tenantId": "tenantId",
+ "extra__azure_data_factory__resource_group_name": RESOURCE_GROUP_NAME,
+ "extra__azure_data_factory__factory_name": DATAFACTORY_NAME,
+ }
+ ],
+ )
+ @async_mock.patch(f"{MODULE}.hooks.data_factory.AzureDataFactoryAsyncHook.get_connection")
+ async def test_get_async_conn_key_error_tenantId(self, mock_connection, mock_connection_params):
+ """Test get_async_conn function with raising key error"""
Review Comment:
done
##########
airflow/providers/microsoft/azure/sensors/data_factory.py:
##########
@@ -78,3 +83,52 @@ def poke(self, context: Context) -> bool:
raise AzureDataFactoryPipelineRunException(f"Pipeline run {self.run_id} has been cancelled.")
return pipeline_run_status == AzureDataFactoryPipelineRunStatus.SUCCEEDED
+
+
+class AzureDataFactoryPipelineRunStatusAsyncSensor(AzureDataFactoryPipelineRunStatusSensor):
+ """
+ Checks the status of a pipeline run asynchronously.
+
+ :param azure_data_factory_conn_id: The connection identifier for connecting to Azure Data Factory.
+ :param run_id: The pipeline run identifier.
+ :param resource_group_name: The resource group name.
+ :param factory_name: The data factory name.
+ :param poll_interval: polling period in seconds to check for the status
Review Comment:
fixed
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] phanikumv commented on a diff in pull request #29801: Add deferrable `AzureDataFactoryPipelineRunStatusSensor`
Posted by "phanikumv (via GitHub)" <gi...@apache.org>.
phanikumv commented on code in PR #29801:
URL: https://github.com/apache/airflow/pull/29801#discussion_r1125620257
##########
tests/providers/microsoft/azure/hooks/test_azure_data_factory.py:
##########
@@ -708,3 +723,214 @@ def test_backcompat_prefix_both_prefers_short(mock_connect):
hook = AzureDataFactoryHook("my_conn")
hook.delete_factory(factory_name="n/a")
mock_connect.return_value.factories.delete.assert_called_with("non-prefixed", "n/a")
+
+
+class TestAzureDataFactoryAsyncHook:
+ @pytest.mark.asyncio
+ @pytest.mark.parametrize(
Review Comment:
yeah, initially I thought of putting all the status in one test by parameterization..but asynctest tests on Python 3.7 wont run like that, and will give the below error
```FAILED test_azure_data_factory.py::TestAzureDataFactoryAsyncHook::test_get_adf_pipeline_run_status_queued[InProgress] - AttributeError: '_patch' object has no attribute '_exit_stack'```
when used in combination with pytest.mark.parametrize it leads to mock __enter()__ being called only once per test case (regardless of number of "subcases" created with parametrisation), but the patch __exit()__ is called after each of the parametrisation "subtests", so only the first "subtest" passes and all other fail. Here is a link about the issue for reference
https://github.com/Martiusweb/asynctest/issues/152#issue-666326338
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] josh-fell merged pull request #29801: Add deferrable `AzureDataFactoryPipelineRunStatusSensor`
Posted by "josh-fell (via GitHub)" <gi...@apache.org>.
josh-fell merged PR #29801:
URL: https://github.com/apache/airflow/pull/29801
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] kaxil commented on a diff in pull request #29801: Add deferrable `AzureDataFactoryPipelineRunStatusSensor`
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #29801:
URL: https://github.com/apache/airflow/pull/29801#discussion_r1124626407
##########
airflow/providers/microsoft/azure/sensors/data_factory.py:
##########
@@ -78,3 +84,60 @@ def poke(self, context: Context) -> bool:
raise AzureDataFactoryPipelineRunException(f"Pipeline run {self.run_id} has been cancelled.")
return pipeline_run_status == AzureDataFactoryPipelineRunStatus.SUCCEEDED
+
+
+class AzureDataFactoryPipelineRunStatusAsyncSensor(AzureDataFactoryPipelineRunStatusSensor):
+ """
+ Checks the status of a pipeline run asynchronously.
+
+ :param azure_data_factory_conn_id: The connection identifier for connecting to Azure Data Factory.
+ :param run_id: The pipeline run identifier.
+ :param resource_group_name: The resource group name.
+ :param factory_name: The data factory name.
+ :param poll_interval: polling period in seconds to check for the status
+ """
+
+ def __init__(
+ self,
+ *,
+ poll_interval: float = 5,
+ **kwargs: Any,
+ ):
+ # TODO: Remove once deprecated
+ if poll_interval:
+ self.poke_interval = poll_interval
+ warnings.warn(
+ "Argument `poll_interval` is deprecated and will be removed "
+ "in a future release. Please use `poke_interval` instead.",
+ DeprecationWarning,
+ stacklevel=2,
Review Comment:
Why deprecate when we have not even mergred, released? Copy/paste error from `astronomer-providers`?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] phanikumv commented on a diff in pull request #29801: Add deferrable `AzureDataFactoryPipelineRunStatusSensor`
Posted by "phanikumv (via GitHub)" <gi...@apache.org>.
phanikumv commented on code in PR #29801:
URL: https://github.com/apache/airflow/pull/29801#discussion_r1125620257
##########
tests/providers/microsoft/azure/hooks/test_azure_data_factory.py:
##########
@@ -708,3 +723,214 @@ def test_backcompat_prefix_both_prefers_short(mock_connect):
hook = AzureDataFactoryHook("my_conn")
hook.delete_factory(factory_name="n/a")
mock_connect.return_value.factories.delete.assert_called_with("non-prefixed", "n/a")
+
+
+class TestAzureDataFactoryAsyncHook:
+ @pytest.mark.asyncio
+ @pytest.mark.parametrize(
Review Comment:
yeah, initially I thought of putting all the status in one test by parameterization..but asyncio tests on Python 3.7 wont run like that, and will give the below error
```FAILED test_azure_data_factory.py::TestAzureDataFactoryAsyncHook::test_get_adf_pipeline_run_status_queued[InProgress] - AttributeError: '_patch' object has no attribute '_exit_stack'```
when used in combination with pytest.mark.parametrize it leads to mock __enter()__ being called only once per test case (regardless of number of "subcases" created with parametrisation), but the patch __exit()__ is called after each of the parametrisation "subtests", so only the first "subtest" passes and all other fail. Here is a link about the issue for reference
https://github.com/Martiusweb/asynctest/issues/152#issue-666326338
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] phanikumv commented on a diff in pull request #29801: Add deferrable `AzureDataFactoryPipelineRunStatusSensor`
Posted by "phanikumv (via GitHub)" <gi...@apache.org>.
phanikumv commented on code in PR #29801:
URL: https://github.com/apache/airflow/pull/29801#discussion_r1132621981
##########
airflow/providers/microsoft/azure/hooks/data_factory.py:
##########
@@ -1039,3 +1048,124 @@ def test_connection(self) -> tuple[bool, str]:
return success
except Exception as e:
return False, str(e)
+
+
+def provide_targeted_factory_async(func: T) -> T:
+ """
+ Provide the targeted factory to the async decorated function in case it isn't specified.
+
+ If ``resource_group_name`` or ``factory_name`` is not provided it defaults to the value specified in
+ the connection extras.
+ """
+ signature = inspect.signature(func)
+
+ @wraps(func)
+ async def wrapper(*args: Any, **kwargs: Any) -> Any:
+ bound_args = signature.bind(*args, **kwargs)
+
+ async def bind_argument(arg: Any, default_key: str) -> None:
+ # Check if arg was not included in the function signature or, if it is, the value is not provided.
+ if arg not in bound_args.arguments or bound_args.arguments[arg] is None:
+ self = args[0]
+ conn = await sync_to_async(self.get_connection)(self.conn_id)
+ extras = conn.extra_dejson
+ default_value = extras.get(default_key) or extras.get(
+ f"extra__azure_data_factory__{default_key}"
+ )
+ if not default_value:
+ raise AirflowException("Could not determine the targeted data factory.")
+
+ bound_args.arguments[arg] = default_value
+
+ await bind_argument("resource_group_name", "resource_group_name")
+ await bind_argument("factory_name", "factory_name")
+
+ return await func(*bound_args.args, **bound_args.kwargs)
+
+ return cast(T, wrapper)
+
+
+class AzureDataFactoryAsyncHook(AzureDataFactoryHook):
+ """
+ An Async Hook that connects to Azure DataFactory to perform pipeline operations
+
+ :param azure_data_factory_conn_id: The :ref:`Azure Data Factory connection id<howto/connection:adf>`.
+ """
+
+ def __init__(self, azure_data_factory_conn_id: str):
+ self._async_conn: AsyncDataFactoryManagementClient = None
+ self.conn_id = azure_data_factory_conn_id
+ super().__init__(azure_data_factory_conn_id=azure_data_factory_conn_id)
+
+ async def get_async_conn(self) -> AsyncDataFactoryManagementClient:
+ """Get async connection and connect to azure data factory"""
+ if self._conn is not None:
+ return self._conn
+
+ conn = await sync_to_async(self.get_connection)(self.conn_id)
+ extras = conn.extra_dejson
+ tenant = get_field(extras, "tenantId")
+
+ try:
+ subscription_id = get_field(extras, "subscriptionId", strict=True)
+ except KeyError:
+ raise ValueError("A Subscription ID is required to connect to Azure Data Factory.")
+
+ credential: AsyncCredentials
+ if conn.login is not None and conn.password is not None:
+ if not tenant:
+ raise ValueError("A Tenant ID is required when authenticating with Client ID and Secret.")
+
+ credential = AsyncClientSecretCredential(
+ client_id=conn.login, client_secret=conn.password, tenant_id=tenant
+ )
+ else:
+ credential = AsyncDefaultAzureCredential()
+
+ return AsyncDataFactoryManagementClient(
+ credential=credential,
+ subscription_id=subscription_id,
+ )
+
+ @provide_targeted_factory_async
+ async def get_pipeline_run(
+ self,
+ run_id: str,
+ resource_group_name: str | None = None,
+ factory_name: str | None = None,
+ **config: Any,
+ ) -> PipelineRun:
+ """
+ Connect to Azure Data Factory asynchronously to get the pipeline run details by run id
+
+ :param run_id: The pipeline run identifier.
+ :param resource_group_name: The resource group name.
+ :param factory_name: The factory name.
Review Comment:
done
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] phanikumv commented on pull request #29801: Add deferrable `AzureDataFactoryPipelineRunStatusSensor`
Posted by "phanikumv (via GitHub)" <gi...@apache.org>.
phanikumv commented on PR #29801:
URL: https://github.com/apache/airflow/pull/29801#issuecomment-1463322920
> Could you add any docs or example DAG for this?
this is done, added docs and example DAG
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] phanikumv commented on a diff in pull request #29801: Add deferrable `AzureDataFactoryPipelineRunStatusSensor`
Posted by "phanikumv (via GitHub)" <gi...@apache.org>.
phanikumv commented on code in PR #29801:
URL: https://github.com/apache/airflow/pull/29801#discussion_r1125612675
##########
airflow/providers/microsoft/azure/sensors/data_factory.py:
##########
@@ -78,3 +84,60 @@ def poke(self, context: Context) -> bool:
raise AzureDataFactoryPipelineRunException(f"Pipeline run {self.run_id} has been cancelled.")
return pipeline_run_status == AzureDataFactoryPipelineRunStatus.SUCCEEDED
+
+
+class AzureDataFactoryPipelineRunStatusAsyncSensor(AzureDataFactoryPipelineRunStatusSensor):
+ """
+ Checks the status of a pipeline run asynchronously.
+
+ :param azure_data_factory_conn_id: The connection identifier for connecting to Azure Data Factory.
+ :param run_id: The pipeline run identifier.
+ :param resource_group_name: The resource group name.
+ :param factory_name: The data factory name.
+ :param poll_interval: polling period in seconds to check for the status
+ """
+
+ def __init__(
+ self,
+ *,
+ poll_interval: float = 5,
+ **kwargs: Any,
+ ):
+ # TODO: Remove once deprecated
+ if poll_interval:
+ self.poke_interval = poll_interval
+ warnings.warn(
+ "Argument `poll_interval` is deprecated and will be removed "
+ "in a future release. Please use `poke_interval` instead.",
+ DeprecationWarning,
+ stacklevel=2,
Review Comment:
yes it is a typo. Corrected it.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] josh-fell commented on a diff in pull request #29801: Add deferrable `AzureDataFactoryPipelineRunStatusSensor`
Posted by "josh-fell (via GitHub)" <gi...@apache.org>.
josh-fell commented on code in PR #29801:
URL: https://github.com/apache/airflow/pull/29801#discussion_r1132491038
##########
airflow/providers/microsoft/azure/hooks/data_factory.py:
##########
@@ -1039,3 +1048,124 @@ def test_connection(self) -> tuple[bool, str]:
return success
except Exception as e:
return False, str(e)
+
+
+def provide_targeted_factory_async(func: T) -> T:
+ """
+ Provide the targeted factory to the async decorated function in case it isn't specified.
+
+ If ``resource_group_name`` or ``factory_name`` is not provided it defaults to the value specified in
+ the connection extras.
+ """
+ signature = inspect.signature(func)
+
+ @wraps(func)
+ async def wrapper(*args: Any, **kwargs: Any) -> Any:
+ bound_args = signature.bind(*args, **kwargs)
+
+ async def bind_argument(arg: Any, default_key: str) -> None:
+ # Check if arg was not included in the function signature or, if it is, the value is not provided.
+ if arg not in bound_args.arguments or bound_args.arguments[arg] is None:
+ self = args[0]
+ conn = await sync_to_async(self.get_connection)(self.conn_id)
+ extras = conn.extra_dejson
+ default_value = extras.get(default_key) or extras.get(
+ f"extra__azure_data_factory__{default_key}"
+ )
+ if not default_value:
+ raise AirflowException("Could not determine the targeted data factory.")
+
+ bound_args.arguments[arg] = default_value
+
+ await bind_argument("resource_group_name", "resource_group_name")
+ await bind_argument("factory_name", "factory_name")
+
+ return await func(*bound_args.args, **bound_args.kwargs)
+
+ return cast(T, wrapper)
+
+
+class AzureDataFactoryAsyncHook(AzureDataFactoryHook):
+ """
+ An Async Hook that connects to Azure DataFactory to perform pipeline operations
+
+ :param azure_data_factory_conn_id: The :ref:`Azure Data Factory connection id<howto/connection:adf>`.
+ """
+
+ def __init__(self, azure_data_factory_conn_id: str):
Review Comment:
WDYT about making `azure_data_factory_conn_id` optional like it is in AzureDataFactoryHook?
##########
airflow/providers/microsoft/azure/hooks/data_factory.py:
##########
@@ -1039,3 +1048,124 @@ def test_connection(self) -> tuple[bool, str]:
return success
except Exception as e:
return False, str(e)
+
+
+def provide_targeted_factory_async(func: T) -> T:
+ """
+ Provide the targeted factory to the async decorated function in case it isn't specified.
+
+ If ``resource_group_name`` or ``factory_name`` is not provided it defaults to the value specified in
+ the connection extras.
+ """
+ signature = inspect.signature(func)
+
+ @wraps(func)
+ async def wrapper(*args: Any, **kwargs: Any) -> Any:
+ bound_args = signature.bind(*args, **kwargs)
+
+ async def bind_argument(arg: Any, default_key: str) -> None:
+ # Check if arg was not included in the function signature or, if it is, the value is not provided.
+ if arg not in bound_args.arguments or bound_args.arguments[arg] is None:
+ self = args[0]
+ conn = await sync_to_async(self.get_connection)(self.conn_id)
+ extras = conn.extra_dejson
+ default_value = extras.get(default_key) or extras.get(
+ f"extra__azure_data_factory__{default_key}"
+ )
+ if not default_value:
+ raise AirflowException("Could not determine the targeted data factory.")
+
+ bound_args.arguments[arg] = default_value
+
+ await bind_argument("resource_group_name", "resource_group_name")
+ await bind_argument("factory_name", "factory_name")
+
+ return await func(*bound_args.args, **bound_args.kwargs)
+
+ return cast(T, wrapper)
+
+
+class AzureDataFactoryAsyncHook(AzureDataFactoryHook):
+ """
+ An Async Hook that connects to Azure DataFactory to perform pipeline operations
+
+ :param azure_data_factory_conn_id: The :ref:`Azure Data Factory connection id<howto/connection:adf>`.
+ """
+
+ def __init__(self, azure_data_factory_conn_id: str):
+ self._async_conn: AsyncDataFactoryManagementClient = None
+ self.conn_id = azure_data_factory_conn_id
+ super().__init__(azure_data_factory_conn_id=azure_data_factory_conn_id)
+
+ async def get_async_conn(self) -> AsyncDataFactoryManagementClient:
+ """Get async connection and connect to azure data factory"""
+ if self._conn is not None:
Review Comment:
I think this function should use `self._async_conn` instead?
##########
airflow/providers/microsoft/azure/hooks/data_factory.py:
##########
@@ -1039,3 +1048,124 @@ def test_connection(self) -> tuple[bool, str]:
return success
except Exception as e:
return False, str(e)
+
+
+def provide_targeted_factory_async(func: T) -> T:
+ """
+ Provide the targeted factory to the async decorated function in case it isn't specified.
+
+ If ``resource_group_name`` or ``factory_name`` is not provided it defaults to the value specified in
+ the connection extras.
+ """
+ signature = inspect.signature(func)
+
+ @wraps(func)
+ async def wrapper(*args: Any, **kwargs: Any) -> Any:
+ bound_args = signature.bind(*args, **kwargs)
+
+ async def bind_argument(arg: Any, default_key: str) -> None:
+ # Check if arg was not included in the function signature or, if it is, the value is not provided.
+ if arg not in bound_args.arguments or bound_args.arguments[arg] is None:
+ self = args[0]
+ conn = await sync_to_async(self.get_connection)(self.conn_id)
+ extras = conn.extra_dejson
+ default_value = extras.get(default_key) or extras.get(
+ f"extra__azure_data_factory__{default_key}"
+ )
+ if not default_value:
+ raise AirflowException("Could not determine the targeted data factory.")
+
+ bound_args.arguments[arg] = default_value
+
+ await bind_argument("resource_group_name", "resource_group_name")
+ await bind_argument("factory_name", "factory_name")
+
+ return await func(*bound_args.args, **bound_args.kwargs)
+
+ return cast(T, wrapper)
+
+
+class AzureDataFactoryAsyncHook(AzureDataFactoryHook):
+ """
+ An Async Hook that connects to Azure DataFactory to perform pipeline operations
+
+ :param azure_data_factory_conn_id: The :ref:`Azure Data Factory connection id<howto/connection:adf>`.
+ """
+
+ def __init__(self, azure_data_factory_conn_id: str):
+ self._async_conn: AsyncDataFactoryManagementClient = None
+ self.conn_id = azure_data_factory_conn_id
+ super().__init__(azure_data_factory_conn_id=azure_data_factory_conn_id)
+
+ async def get_async_conn(self) -> AsyncDataFactoryManagementClient:
+ """Get async connection and connect to azure data factory"""
+ if self._conn is not None:
+ return self._conn
+
+ conn = await sync_to_async(self.get_connection)(self.conn_id)
+ extras = conn.extra_dejson
+ tenant = get_field(extras, "tenantId")
+
+ try:
+ subscription_id = get_field(extras, "subscriptionId", strict=True)
+ except KeyError:
+ raise ValueError("A Subscription ID is required to connect to Azure Data Factory.")
+
+ credential: AsyncCredentials
+ if conn.login is not None and conn.password is not None:
+ if not tenant:
+ raise ValueError("A Tenant ID is required when authenticating with Client ID and Secret.")
+
+ credential = AsyncClientSecretCredential(
+ client_id=conn.login, client_secret=conn.password, tenant_id=tenant
+ )
+ else:
+ credential = AsyncDefaultAzureCredential()
+
+ return AsyncDataFactoryManagementClient(
+ credential=credential,
+ subscription_id=subscription_id,
+ )
+
+ @provide_targeted_factory_async
+ async def get_pipeline_run(
+ self,
+ run_id: str,
+ resource_group_name: str | None = None,
+ factory_name: str | None = None,
+ **config: Any,
+ ) -> PipelineRun:
+ """
+ Connect to Azure Data Factory asynchronously to get the pipeline run details by run id
+
+ :param run_id: The pipeline run identifier.
+ :param resource_group_name: The resource group name.
+ :param factory_name: The factory name.
Review Comment:
```suggestion
:param factory_name: The factory name.
:param config: Extra parameters for the ADF client.
```
To be consistent with the other sync functions and to let users know they can pass in additional configuration should they choose to use this function directly.
##########
airflow/providers/microsoft/azure/hooks/data_factory.py:
##########
@@ -1039,3 +1048,124 @@ def test_connection(self) -> tuple[bool, str]:
return success
except Exception as e:
return False, str(e)
+
+
+def provide_targeted_factory_async(func: T) -> T:
+ """
+ Provide the targeted factory to the async decorated function in case it isn't specified.
+
+ If ``resource_group_name`` or ``factory_name`` is not provided it defaults to the value specified in
+ the connection extras.
+ """
+ signature = inspect.signature(func)
+
+ @wraps(func)
+ async def wrapper(*args: Any, **kwargs: Any) -> Any:
+ bound_args = signature.bind(*args, **kwargs)
+
+ async def bind_argument(arg: Any, default_key: str) -> None:
+ # Check if arg was not included in the function signature or, if it is, the value is not provided.
+ if arg not in bound_args.arguments or bound_args.arguments[arg] is None:
+ self = args[0]
+ conn = await sync_to_async(self.get_connection)(self.conn_id)
+ extras = conn.extra_dejson
+ default_value = extras.get(default_key) or extras.get(
+ f"extra__azure_data_factory__{default_key}"
+ )
+ if not default_value:
+ raise AirflowException("Could not determine the targeted data factory.")
+
+ bound_args.arguments[arg] = default_value
+
+ await bind_argument("resource_group_name", "resource_group_name")
+ await bind_argument("factory_name", "factory_name")
+
+ return await func(*bound_args.args, **bound_args.kwargs)
+
+ return cast(T, wrapper)
+
+
+class AzureDataFactoryAsyncHook(AzureDataFactoryHook):
+ """
+ An Async Hook that connects to Azure DataFactory to perform pipeline operations
+
+ :param azure_data_factory_conn_id: The :ref:`Azure Data Factory connection id<howto/connection:adf>`.
+ """
+
+ def __init__(self, azure_data_factory_conn_id: str):
+ self._async_conn: AsyncDataFactoryManagementClient = None
+ self.conn_id = azure_data_factory_conn_id
+ super().__init__(azure_data_factory_conn_id=azure_data_factory_conn_id)
+
+ async def get_async_conn(self) -> AsyncDataFactoryManagementClient:
+ """Get async connection and connect to azure data factory"""
+ if self._conn is not None:
Review Comment:
And `self._async_conn` should be returned too so the initial `is not None` check acts as a form of "caching".
##########
airflow/providers/microsoft/azure/sensors/data_factory.py:
##########
@@ -78,3 +83,52 @@ def poke(self, context: Context) -> bool:
raise AzureDataFactoryPipelineRunException(f"Pipeline run {self.run_id} has been cancelled.")
return pipeline_run_status == AzureDataFactoryPipelineRunStatus.SUCCEEDED
+
+
+class AzureDataFactoryPipelineRunStatusAsyncSensor(AzureDataFactoryPipelineRunStatusSensor):
+ """
+ Checks the status of a pipeline run asynchronously.
+
+ :param azure_data_factory_conn_id: The connection identifier for connecting to Azure Data Factory.
+ :param run_id: The pipeline run identifier.
+ :param resource_group_name: The resource group name.
+ :param factory_name: The data factory name.
+ :param poll_interval: polling period in seconds to check for the status
Review Comment:
```suggestion
:param poke_interval: polling period in seconds to check for the status
```
##########
tests/providers/microsoft/azure/hooks/test_azure_data_factory.py:
##########
@@ -708,3 +724,245 @@ def test_backcompat_prefix_both_prefers_short(mock_connect):
hook = AzureDataFactoryHook("my_conn")
hook.delete_factory(factory_name="n/a")
mock_connect.return_value.factories.delete.assert_called_with("non-prefixed", "n/a")
+
+
+class TestAzureDataFactoryAsyncHook:
+ @pytest.mark.asyncio
+ @async_mock.patch(f"{MODULE}.hooks.data_factory.AzureDataFactoryAsyncHook.get_async_conn")
+ @async_mock.patch(f"{MODULE}.hooks.data_factory.AzureDataFactoryAsyncHook.get_pipeline_run")
+ async def test_get_adf_pipeline_run_status_queued(self, mock_get_pipeline_run, mock_conn):
+ """Test get_adf_pipeline_run_status function with mocked status"""
+ mock_status = "Queued"
+ mock_get_pipeline_run.return_value.status = mock_status
+ hook = AzureDataFactoryAsyncHook(AZURE_DATA_FACTORY_CONN_ID)
+ response = await hook.get_adf_pipeline_run_status(RUN_ID, RESOURCE_GROUP_NAME, DATAFACTORY_NAME)
+ assert response == mock_status
+
+ @pytest.mark.asyncio
+ @async_mock.patch(f"{MODULE}.hooks.data_factory.AzureDataFactoryAsyncHook.get_async_conn")
+ @async_mock.patch(f"{MODULE}.hooks.data_factory.AzureDataFactoryAsyncHook.get_pipeline_run")
+ async def test_get_adf_pipeline_run_status_inprogress(
+ self,
+ mock_get_pipeline_run,
+ mock_conn,
+ ):
+ """Test get_adf_pipeline_run_status function with mocked status"""
+ mock_status = "InProgress"
+ mock_get_pipeline_run.return_value.status = mock_status
+ hook = AzureDataFactoryAsyncHook(AZURE_DATA_FACTORY_CONN_ID)
+ response = await hook.get_adf_pipeline_run_status(RUN_ID, RESOURCE_GROUP_NAME, DATAFACTORY_NAME)
+ assert response == mock_status
+
+ @pytest.mark.asyncio
+ @async_mock.patch(f"{MODULE}.hooks.data_factory.AzureDataFactoryAsyncHook.get_async_conn")
+ @async_mock.patch(f"{MODULE}.hooks.data_factory.AzureDataFactoryAsyncHook.get_pipeline_run")
+ async def test_get_adf_pipeline_run_status_success(self, mock_get_pipeline_run, mock_conn):
+ """Test get_adf_pipeline_run_status function with mocked status"""
+ mock_status = "Succeeded"
+ mock_get_pipeline_run.return_value.status = mock_status
+ hook = AzureDataFactoryAsyncHook(AZURE_DATA_FACTORY_CONN_ID)
+ response = await hook.get_adf_pipeline_run_status(RUN_ID, RESOURCE_GROUP_NAME, DATAFACTORY_NAME)
+ assert response == mock_status
+
+ @pytest.mark.asyncio
+ @async_mock.patch(f"{MODULE}.hooks.data_factory.AzureDataFactoryAsyncHook.get_async_conn")
+ @async_mock.patch(f"{MODULE}.hooks.data_factory.AzureDataFactoryAsyncHook.get_pipeline_run")
+ async def test_get_adf_pipeline_run_status_failed(self, mock_get_pipeline_run, mock_conn):
+ """Test get_adf_pipeline_run_status function with mocked status"""
+ mock_status = "Failed"
+ mock_get_pipeline_run.return_value.status = mock_status
+ hook = AzureDataFactoryAsyncHook(AZURE_DATA_FACTORY_CONN_ID)
+ response = await hook.get_adf_pipeline_run_status(RUN_ID, RESOURCE_GROUP_NAME, DATAFACTORY_NAME)
+ assert response == mock_status
+
+ @pytest.mark.asyncio
+ @async_mock.patch(f"{MODULE}.hooks.data_factory.AzureDataFactoryAsyncHook.get_async_conn")
+ @async_mock.patch(f"{MODULE}.hooks.data_factory.AzureDataFactoryAsyncHook.get_pipeline_run")
+ async def test_get_adf_pipeline_run_status_cancelled(self, mock_get_pipeline_run, mock_conn):
+ """Test get_adf_pipeline_run_status function with mocked status"""
+ mock_status = "Cancelled"
+ mock_get_pipeline_run.return_value.status = mock_status
+ hook = AzureDataFactoryAsyncHook(AZURE_DATA_FACTORY_CONN_ID)
+ response = await hook.get_adf_pipeline_run_status(RUN_ID, RESOURCE_GROUP_NAME, DATAFACTORY_NAME)
+ assert response == mock_status
+
+ @pytest.mark.asyncio
+ @async_mock.patch(f"{MODULE}.hooks.data_factory.AzureDataFactoryAsyncHook.get_async_conn")
+ @async_mock.patch(f"{MODULE}.hooks.data_factory.AzureDataFactoryAsyncHook.get_pipeline_run")
+ async def test_get_adf_pipeline_run_status_exception(self, mock_get_pipeline_run, mock_conn):
+ """Test get_adf_pipeline_run_status function with exception"""
+ mock_get_pipeline_run.side_effect = Exception("Test exception")
+ hook = AzureDataFactoryAsyncHook(AZURE_DATA_FACTORY_CONN_ID)
+ with pytest.raises(AirflowException):
+ await hook.get_adf_pipeline_run_status(RUN_ID, RESOURCE_GROUP_NAME, DATAFACTORY_NAME)
+
+ @pytest.mark.asyncio
+ @async_mock.patch("azure.mgmt.datafactory.models._models_py3.PipelineRun")
+ @async_mock.patch(f"{MODULE}.hooks.data_factory.AzureDataFactoryAsyncHook.get_connection")
+ @async_mock.patch(f"{MODULE}.hooks.data_factory.AzureDataFactoryAsyncHook.get_async_conn")
+ async def test_get_pipeline_run_exception_without_resource(
+ self, mock_conn, mock_get_connection, mock_pipeline_run
+ ):
+ """
+ Test get_pipeline_run function without passing the resource name to check the decorator function and
+ raise exception
+ """
+ mock_connection = Connection(
+ extra=json.dumps({"extra__azure_data_factory__factory_name": DATAFACTORY_NAME})
+ )
+ mock_get_connection.return_value = mock_connection
+ mock_conn.return_value.__aenter__.return_value.pipeline_runs.get.return_value = mock_pipeline_run
+ hook = AzureDataFactoryAsyncHook(AZURE_DATA_FACTORY_CONN_ID)
+ with pytest.raises(AirflowException):
+ await hook.get_pipeline_run(RUN_ID, None, DATAFACTORY_NAME)
+
+ @pytest.mark.asyncio
+ @async_mock.patch(f"{MODULE}.hooks.data_factory.AzureDataFactoryAsyncHook.get_async_conn")
+ async def test_get_pipeline_run_exception(self, mock_conn):
+ """Test get_pipeline_run function with exception"""
+ mock_conn.return_value.__aenter__.return_value.pipeline_runs.get.side_effect = Exception(
+ "Test exception"
+ )
+ hook = AzureDataFactoryAsyncHook(AZURE_DATA_FACTORY_CONN_ID)
+ with pytest.raises(AirflowException):
+ await hook.get_pipeline_run(RUN_ID, RESOURCE_GROUP_NAME, DATAFACTORY_NAME)
+
+ @pytest.mark.asyncio
+ @async_mock.patch(f"{MODULE}.hooks.data_factory.AzureDataFactoryAsyncHook.get_connection")
+ async def test_get_async_conn(self, mock_connection):
+ """"""
+ mock_conn = Connection(
+ conn_id=DEFAULT_CONNECTION_CLIENT_SECRET,
+ conn_type="azure_data_factory",
+ login="clientId",
+ password="clientSecret",
+ extra=json.dumps(
+ {
+ "extra__azure_data_factory__tenantId": "tenantId",
+ "extra__azure_data_factory__subscriptionId": "subscriptionId",
+ "extra__azure_data_factory__resource_group_name": RESOURCE_GROUP_NAME,
+ "extra__azure_data_factory__factory_name": DATAFACTORY_NAME,
+ }
+ ),
+ )
+ mock_connection.return_value = mock_conn
+ hook = AzureDataFactoryAsyncHook(AZURE_DATA_FACTORY_CONN_ID)
+ response = await hook.get_async_conn()
+ assert isinstance(response, DataFactoryManagementClient)
+
+ @pytest.mark.asyncio
+ @async_mock.patch(f"{MODULE}.hooks.data_factory.AzureDataFactoryAsyncHook.get_connection")
+ async def test_get_async_conn_without_login_id(self, mock_connection):
+ """Test get_async_conn function without login id"""
+ mock_conn = Connection(
+ conn_id=DEFAULT_CONNECTION_CLIENT_SECRET,
+ conn_type="azure_data_factory",
+ extra=json.dumps(
+ {
+ "extra__azure_data_factory__tenantId": "tenantId",
+ "extra__azure_data_factory__subscriptionId": "subscriptionId",
+ "extra__azure_data_factory__resource_group_name": RESOURCE_GROUP_NAME,
+ "extra__azure_data_factory__factory_name": DATAFACTORY_NAME,
+ }
+ ),
+ )
+ mock_connection.return_value = mock_conn
+ hook = AzureDataFactoryAsyncHook(AZURE_DATA_FACTORY_CONN_ID)
+ response = await hook.get_async_conn()
+ assert isinstance(response, DataFactoryManagementClient)
+
+ @pytest.mark.asyncio
+ @pytest.mark.parametrize(
+ "mock_connection_params",
+ [
+ {
+ "extra__azure_data_factory__tenantId": "tenantId",
+ "extra__azure_data_factory__resource_group_name": RESOURCE_GROUP_NAME,
+ "extra__azure_data_factory__factory_name": DATAFACTORY_NAME,
+ }
+ ],
+ )
+ @async_mock.patch(f"{MODULE}.hooks.data_factory.AzureDataFactoryAsyncHook.get_connection")
+ async def test_get_async_conn_key_error_tenantId(self, mock_connection, mock_connection_params):
+ """Test get_async_conn function with raising key error"""
Review Comment:
Would you mind updating the docstrings for this function and `test_get_async_conn_key_error_subscriptionId`? They are both the same.
Although I think the function names should be switched? `test_get_async_conn_key_error_tenantId` is testing that a subscription ID is not part of the connection while `test_get_async_conn_key_error_subscriptionId` is testing that tenant ID is required when using login/password.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] josh-fell commented on a diff in pull request #29801: Add deferrable `AzureDataFactoryPipelineRunStatusSensor`
Posted by "josh-fell (via GitHub)" <gi...@apache.org>.
josh-fell commented on code in PR #29801:
URL: https://github.com/apache/airflow/pull/29801#discussion_r1124899383
##########
airflow/providers/microsoft/azure/hooks/data_factory.py:
##########
@@ -1039,3 +1048,120 @@ def test_connection(self) -> tuple[bool, str]:
return success
except Exception as e:
return False, str(e)
+
+
+def provide_targeted_factory_async(func: T) -> T:
+ """
+ Provide the targeted factory to the async decorated function in case it isn't specified.
+
+ If ``resource_group_name`` or ``factory_name`` is not provided it defaults to the value specified in
+ the connection extras.
+ """
+ signature = inspect.signature(func)
+
+ @wraps(func)
+ async def wrapper(*args: Any, **kwargs: Any) -> Any:
+ bound_args = signature.bind(*args, **kwargs)
+
+ async def bind_argument(arg: Any, default_key: str) -> None:
+ # Check if arg was not included in the function signature or, if it is, the value is not provided.
+ if arg not in bound_args.arguments or bound_args.arguments[arg] is None:
+ self = args[0]
+ conn = await sync_to_async(self.get_connection)(self.conn_id)
+ default_value = conn.extra_dejson.get(default_key)
Review Comment:
Can you update this to match the sync version? Prefixing connection extras with `extra__...` is no longer needed, but we do keep the check for backwards compat.
##########
airflow/providers/microsoft/azure/hooks/data_factory.py:
##########
@@ -1039,3 +1048,120 @@ def test_connection(self) -> tuple[bool, str]:
return success
except Exception as e:
return False, str(e)
+
+
+def provide_targeted_factory_async(func: T) -> T:
+ """
+ Provide the targeted factory to the async decorated function in case it isn't specified.
+
+ If ``resource_group_name`` or ``factory_name`` is not provided it defaults to the value specified in
+ the connection extras.
+ """
+ signature = inspect.signature(func)
+
+ @wraps(func)
+ async def wrapper(*args: Any, **kwargs: Any) -> Any:
+ bound_args = signature.bind(*args, **kwargs)
+
+ async def bind_argument(arg: Any, default_key: str) -> None:
+ # Check if arg was not included in the function signature or, if it is, the value is not provided.
+ if arg not in bound_args.arguments or bound_args.arguments[arg] is None:
+ self = args[0]
+ conn = await sync_to_async(self.get_connection)(self.conn_id)
+ default_value = conn.extra_dejson.get(default_key)
+ if not default_value:
+ raise AirflowException("Could not determine the targeted data factory.")
+
+ bound_args.arguments[arg] = conn.extra_dejson[default_key]
+
+ await bind_argument("resource_group_name", "extra__azure_data_factory__resource_group_name")
+ await bind_argument("factory_name", "extra__azure_data_factory__factory_name")
+
+ return await func(*bound_args.args, **bound_args.kwargs)
+
+ return cast(T, wrapper)
+
+
+class AzureDataFactoryAsyncHook(AzureDataFactoryHook):
+ """
+ An Async Hook that connects to Azure DataFactory to perform pipeline operations
+
+ :param azure_data_factory_conn_id: The :ref:`Azure Data Factory connection id<howto/connection:adf>`.
+ """
+
+ def __init__(self, azure_data_factory_conn_id: str):
+ self._async_conn: AsyncDataFactoryManagementClient = None
+ self.conn_id = azure_data_factory_conn_id
+ super().__init__(azure_data_factory_conn_id=azure_data_factory_conn_id)
+
+ async def get_async_conn(self) -> AsyncDataFactoryManagementClient:
+ """Get async connection and connect to azure data factory"""
+ if self._conn is not None:
+ return self._conn
+
+ conn = await sync_to_async(self.get_connection)(self.conn_id)
+ tenant = conn.extra_dejson.get("extra__azure_data_factory__tenantId")
Review Comment:
Same idea here. Should use the `get_field()` function to retrieve the extras.
##########
tests/providers/microsoft/azure/hooks/test_azure_data_factory.py:
##########
@@ -708,3 +723,214 @@ def test_backcompat_prefix_both_prefers_short(mock_connect):
hook = AzureDataFactoryHook("my_conn")
hook.delete_factory(factory_name="n/a")
mock_connect.return_value.factories.delete.assert_called_with("non-prefixed", "n/a")
+
+
+class TestAzureDataFactoryAsyncHook:
+ @pytest.mark.asyncio
+ @pytest.mark.parametrize(
Review Comment:
Why parametrize if each status is its own explicit test? You can inject test case `ids` like `pytest.mark.parametrize(..., ids=["test_status1", "test_status2"])`
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] phanikumv commented on a diff in pull request #29801: Add deferrable `AzureDataFactoryPipelineRunStatusSensor`
Posted by "phanikumv (via GitHub)" <gi...@apache.org>.
phanikumv commented on code in PR #29801:
URL: https://github.com/apache/airflow/pull/29801#discussion_r1125620257
##########
tests/providers/microsoft/azure/hooks/test_azure_data_factory.py:
##########
@@ -708,3 +723,214 @@ def test_backcompat_prefix_both_prefers_short(mock_connect):
hook = AzureDataFactoryHook("my_conn")
hook.delete_factory(factory_name="n/a")
mock_connect.return_value.factories.delete.assert_called_with("non-prefixed", "n/a")
+
+
+class TestAzureDataFactoryAsyncHook:
+ @pytest.mark.asyncio
+ @pytest.mark.parametrize(
Review Comment:
yeah, initially I thought of putting all the status in one test by parameterization..but asyncio tests on Python 3.7 wont run like that, and will give the below error
FAILED test_azure_data_factory.py::TestAzureDataFactoryAsyncHook::test_get_adf_pipeline_run_status_queued[InProgress] - AttributeError: '_patch' object has no attribute '_exit_stack'
when used in combination with pytest.mark.parametrize it leads to mock __enter()__ being called only once per test case (regardless of number of "subcases" created with parametrisation), but the patch __exit()__ is called after each of the parametrisation "subtests", so only the first "subtest" passes and all other fail. Here is a link about the issue for reference
https://github.com/Martiusweb/asynctest/issues/152#issue-666326338
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] phanikumv commented on a diff in pull request #29801: Add deferrable `AzureDataFactoryPipelineRunStatusSensor`
Posted by "phanikumv (via GitHub)" <gi...@apache.org>.
phanikumv commented on code in PR #29801:
URL: https://github.com/apache/airflow/pull/29801#discussion_r1125631634
##########
airflow/providers/microsoft/azure/hooks/data_factory.py:
##########
@@ -1039,3 +1048,120 @@ def test_connection(self) -> tuple[bool, str]:
return success
except Exception as e:
return False, str(e)
+
+
+def provide_targeted_factory_async(func: T) -> T:
+ """
+ Provide the targeted factory to the async decorated function in case it isn't specified.
+
+ If ``resource_group_name`` or ``factory_name`` is not provided it defaults to the value specified in
+ the connection extras.
+ """
+ signature = inspect.signature(func)
+
+ @wraps(func)
+ async def wrapper(*args: Any, **kwargs: Any) -> Any:
+ bound_args = signature.bind(*args, **kwargs)
+
+ async def bind_argument(arg: Any, default_key: str) -> None:
+ # Check if arg was not included in the function signature or, if it is, the value is not provided.
+ if arg not in bound_args.arguments or bound_args.arguments[arg] is None:
+ self = args[0]
+ conn = await sync_to_async(self.get_connection)(self.conn_id)
+ default_value = conn.extra_dejson.get(default_key)
+ if not default_value:
+ raise AirflowException("Could not determine the targeted data factory.")
+
+ bound_args.arguments[arg] = conn.extra_dejson[default_key]
+
+ await bind_argument("resource_group_name", "extra__azure_data_factory__resource_group_name")
+ await bind_argument("factory_name", "extra__azure_data_factory__factory_name")
+
+ return await func(*bound_args.args, **bound_args.kwargs)
+
+ return cast(T, wrapper)
+
+
+class AzureDataFactoryAsyncHook(AzureDataFactoryHook):
+ """
+ An Async Hook that connects to Azure DataFactory to perform pipeline operations
+
+ :param azure_data_factory_conn_id: The :ref:`Azure Data Factory connection id<howto/connection:adf>`.
+ """
+
+ def __init__(self, azure_data_factory_conn_id: str):
+ self._async_conn: AsyncDataFactoryManagementClient = None
+ self.conn_id = azure_data_factory_conn_id
+ super().__init__(azure_data_factory_conn_id=azure_data_factory_conn_id)
+
+ async def get_async_conn(self) -> AsyncDataFactoryManagementClient:
+ """Get async connection and connect to azure data factory"""
+ if self._conn is not None:
+ return self._conn
+
+ conn = await sync_to_async(self.get_connection)(self.conn_id)
+ tenant = conn.extra_dejson.get("extra__azure_data_factory__tenantId")
Review Comment:
done
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] phanikumv commented on a diff in pull request #29801: Add deferrable `AzureDataFactoryPipelineRunStatusSensor`
Posted by "phanikumv (via GitHub)" <gi...@apache.org>.
phanikumv commented on code in PR #29801:
URL: https://github.com/apache/airflow/pull/29801#discussion_r1133483001
##########
airflow/providers/microsoft/azure/hooks/data_factory.py:
##########
@@ -1039,3 +1048,124 @@ def test_connection(self) -> tuple[bool, str]:
return success
except Exception as e:
return False, str(e)
+
+
+def provide_targeted_factory_async(func: T) -> T:
+ """
+ Provide the targeted factory to the async decorated function in case it isn't specified.
+
+ If ``resource_group_name`` or ``factory_name`` is not provided it defaults to the value specified in
+ the connection extras.
+ """
+ signature = inspect.signature(func)
+
+ @wraps(func)
+ async def wrapper(*args: Any, **kwargs: Any) -> Any:
+ bound_args = signature.bind(*args, **kwargs)
+
+ async def bind_argument(arg: Any, default_key: str) -> None:
+ # Check if arg was not included in the function signature or, if it is, the value is not provided.
+ if arg not in bound_args.arguments or bound_args.arguments[arg] is None:
+ self = args[0]
+ conn = await sync_to_async(self.get_connection)(self.conn_id)
+ extras = conn.extra_dejson
+ default_value = extras.get(default_key) or extras.get(
+ f"extra__azure_data_factory__{default_key}"
+ )
+ if not default_value:
+ raise AirflowException("Could not determine the targeted data factory.")
+
+ bound_args.arguments[arg] = default_value
+
+ await bind_argument("resource_group_name", "resource_group_name")
+ await bind_argument("factory_name", "factory_name")
+
+ return await func(*bound_args.args, **bound_args.kwargs)
+
+ return cast(T, wrapper)
+
+
+class AzureDataFactoryAsyncHook(AzureDataFactoryHook):
+ """
+ An Async Hook that connects to Azure DataFactory to perform pipeline operations
+
+ :param azure_data_factory_conn_id: The :ref:`Azure Data Factory connection id<howto/connection:adf>`.
+ """
+
+ def __init__(self, azure_data_factory_conn_id: str):
+ self._async_conn: AsyncDataFactoryManagementClient = None
+ self.conn_id = azure_data_factory_conn_id
+ super().__init__(azure_data_factory_conn_id=azure_data_factory_conn_id)
+
+ async def get_async_conn(self) -> AsyncDataFactoryManagementClient:
+ """Get async connection and connect to azure data factory"""
+ if self._conn is not None:
Review Comment:
this is done in [264e90d](https://github.com/apache/airflow/pull/29801/commits/264e90dd62f4f6be0498b7df0fb03594b5936153)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org