dabla commented on code in PR #62772:
URL: https://github.com/apache/airflow/pull/62772#discussion_r2976995590
##########
providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/container_instance.py:
##########
@@ -170,3 +177,97 @@ def test_connection(self):
return False, str(e)
return True, "Successfully connected to Azure Container Instance."
+
+
+class AzureContainerInstanceAsyncHook(AzureContainerInstanceHook):
+ """
+ An async hook for communicating with Azure Container Instances.
+
+ :param azure_conn_id: :ref:`Azure connection id<howto/connection:azure>` of
+ a service principal which will be used to start the container instance.
+ """
+
+ def __init__(self, azure_conn_id: str =
AzureContainerInstanceHook.default_conn_name) -> None:
+ self._async_conn: AsyncContainerInstanceManagementClient | None = None
+ self._async_credential: Any = None
+ super().__init__(azure_conn_id=azure_conn_id)
+
+ async def __aenter__(self) -> AzureContainerInstanceAsyncHook:
+ return self
Review Comment:
Don't really like the contextmanager on hook level, I would prefer to see
this on the get_async_conn method instead, as the the connection we want to
make sure is closed correctly in case of exception, not the hook, as the hook
is just a way to get the connection, but it isn't the resource itself, the
connection is.
So more something like this:
```
@asynccontextmanager
async def get_async_conn(self):
credential = ...
client = AsyncContainerInstanceManagementClient(...)
try:
yield client
finally:
await client.close()
if hasattr(credential, "close"):
await credential.close()
```
So later on it can be used like this in the trigger:
```
async with hook.get_async_conn() as client:
await client.container_groups.get(...)
```
But I really like the fact that you implemented a contextmanager, because
that's a clean way to cope with connections.
##########
providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/container_instances.py:
##########
@@ -303,12 +315,14 @@ def _ensure_identity(identity: ContainerGroupIdentity |
dict | None) -> Containe
)
return identity
+ @cached_property
+ def _ci_hook(self) -> AzureContainerInstanceHook:
Review Comment:
Why not just call this property hook?
##########
providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/container_instance.py:
##########
@@ -170,3 +177,97 @@ def test_connection(self):
return False, str(e)
return True, "Successfully connected to Azure Container Instance."
+
+
+class AzureContainerInstanceAsyncHook(AzureContainerInstanceHook):
+ """
+ An async hook for communicating with Azure Container Instances.
+
+ :param azure_conn_id: :ref:`Azure connection id<howto/connection:azure>` of
+ a service principal which will be used to start the container instance.
+ """
+
+ def __init__(self, azure_conn_id: str =
AzureContainerInstanceHook.default_conn_name) -> None:
+ self._async_conn: AsyncContainerInstanceManagementClient | None = None
Review Comment:
First call super().__init(), then initialise additional fields
##########
providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/container_instance.py:
##########
@@ -170,3 +177,97 @@ def test_connection(self):
return False, str(e)
return True, "Successfully connected to Azure Container Instance."
+
+
+class AzureContainerInstanceAsyncHook(AzureContainerInstanceHook):
+ """
+ An async hook for communicating with Azure Container Instances.
+
+ :param azure_conn_id: :ref:`Azure connection id<howto/connection:azure>` of
+ a service principal which will be used to start the container instance.
+ """
+
+ def __init__(self, azure_conn_id: str =
AzureContainerInstanceHook.default_conn_name) -> None:
+ self._async_conn: AsyncContainerInstanceManagementClient | None = None
+ self._async_credential: Any = None
+ super().__init__(azure_conn_id=azure_conn_id)
+
+ async def __aenter__(self) -> AzureContainerInstanceAsyncHook:
+ return self
+
+ async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) ->
None:
+ await self.close()
+
+ async def close(self) -> None:
+ """Close the async connection and credential."""
+ if self._async_conn is not None:
+ await self._async_conn.close()
+ self._async_conn = None
+ if self._async_credential is not None and
hasattr(self._async_credential, "close"):
+ await self._async_credential.close()
+ self._async_credential = None
+
+ async def get_async_conn(self) -> AsyncContainerInstanceManagementClient:
+ """Return (or create) the async management client."""
+ if self._async_conn is not None:
+ return self._async_conn
+
+ conn = self.get_connection(self.conn_id)
Review Comment:
I think we need to use the get_async_connection from common compat here
instead of self.get_connection(self.conn_id)
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
