diff --git a/airflow/providers/microsoft/azure/CHANGELOG.rst b/airflow/providers/microsoft/azure/CHANGELOG.rst index 4030c3c7ed92e..9fe19f899ec61 100644 --- a/airflow/providers/microsoft/azure/CHANGELOG.rst +++ b/airflow/providers/microsoft/azure/CHANGELOG.rst @@ -36,6 +36,11 @@ Breaking changes * ``azure_synapse_pipeline`` connection type has been changed to ``azure_synapse``. * The usage of ``default_conn_name=azure_synapse_connection`` is deprecated and will be removed in future. The new default connection name for ``AzureSynapsePipelineHook`` is: ``default_conn_name=azure_synapse_default``. +Significant changes +~~~~~~~~~~~~~~~~~~~ +.. warning:: + * We bumped the minimum version of azure-cosmos to 4.6.0, and providing a partition key is now required to create, get or delete a container and to get a document. + 9.0.1 ..... diff --git a/airflow/providers/microsoft/azure/hooks/cosmos.py b/airflow/providers/microsoft/azure/hooks/cosmos.py index fe4aa88b5f7a5..9d8b360a75522 100644 --- a/airflow/providers/microsoft/azure/hooks/cosmos.py +++ b/airflow/providers/microsoft/azure/hooks/cosmos.py @@ -27,9 +27,10 @@ from __future__ import annotations import uuid -from typing import Any +from typing import TYPE_CHECKING, Any, List, Union from urllib.parse import urlparse +from azure.cosmos import PartitionKey from azure.cosmos.cosmos_client import CosmosClient from azure.cosmos.exceptions import CosmosHttpResponseError from azure.mgmt.cosmosdb import CosmosDBManagementClient @@ -42,6 +43,9 @@ get_sync_default_azure_credential, ) +if TYPE_CHECKING: + PartitionKeyType = Union[str, List[str]] + class AzureCosmosDBHook(BaseHook): """ @@ -111,6 +115,7 @@ def __init__(self, azure_cosmos_conn_id: str = default_conn_name) -> None: self.default_database_name = None self.default_collection_name = None + self.default_partition_key = None def _get_field(self, extras, name): return get_field( @@ -153,6 +158,7 @@ def get_conn(self) -> CosmosClient: self.default_database_name = self._get_field(extras, "database_name") self.default_collection_name = self._get_field(extras, "collection_name") + self.default_partition_key = self._get_field(extras, "partition_key") # Initialize the Python Azure Cosmos DB client self._conn = CosmosClient(endpoint_uri, {"masterKey": master_key}) @@ -180,6 +186,18 @@ def __get_collection_name(self, collection_name: str | None = None) -> str: return coll_name + def __get_partition_key(self, partition_key: PartitionKeyType | None = None) -> PartitionKeyType: + self.get_conn() + if partition_key is None: + part_key = self.default_partition_key + else: + part_key = partition_key + + if part_key is None: + raise AirflowBadRequest("Partition key must be specified") + + return part_key + def does_collection_exist(self, collection_name: str, database_name: str) -> bool: """Check if a collection exists in CosmosDB.""" if collection_name is None: @@ -204,7 +222,7 @@ def create_collection( self, collection_name: str, database_name: str | None = None, - partition_key: str | None = None, + partition_key: PartitionKeyType | None = None, ) -> None: """Create a new collection in the CosmosDB database.""" if collection_name is None: @@ -226,7 +244,8 @@ def create_collection( # Only create if we did not find it already existing if not existing_container: self.get_conn().get_database_client(self.__get_database_name(database_name)).create_container( - collection_name, partition_key=partition_key + collection_name, + partition_key=PartitionKey(path=self.__get_partition_key(partition_key)), ) def does_database_exist(self, database_name: str) -> bool: @@ -328,7 +347,7 @@ def delete_document( document_id: str, database_name: str | None = None, collection_name: str | None = None, - partition_key: str | None = None, + partition_key: PartitionKeyType | None = None, ) -> None: """Delete an existing document out of a collection in the CosmosDB database.""" if document_id is None: @@ -337,7 +356,7 @@ def delete_document( self.get_conn() .get_database_client(self.__get_database_name(database_name)) .get_container_client(self.__get_collection_name(collection_name)) - .delete_item(document_id, partition_key=partition_key) + .delete_item(document_id, partition_key=self.__get_partition_key(partition_key)) ) def get_document( @@ -345,7 +364,7 @@ def get_document( document_id: str, database_name: str | None = None, collection_name: str | None = None, - partition_key: str | None = None, + partition_key: PartitionKeyType | None = None, ): """Get a document from an existing collection in the CosmosDB database.""" if document_id is None: @@ -356,7 +375,7 @@ def get_document( self.get_conn() .get_database_client(self.__get_database_name(database_name)) .get_container_client(self.__get_collection_name(collection_name)) - .read_item(document_id, partition_key=partition_key) + .read_item(document_id, partition_key=self.__get_partition_key(partition_key)) ) except CosmosHttpResponseError: return None @@ -366,7 +385,7 @@ def get_documents( sql_string: str, database_name: str | None = None, collection_name: str | None = None, - partition_key: str | None = None, + partition_key: PartitionKeyType | None = None, ) -> list | None: """Get a list of documents from an existing collection in the CosmosDB database via SQL query.""" if sql_string is None: @@ -377,7 +396,7 @@ def get_documents( self.get_conn() .get_database_client(self.__get_database_name(database_name)) .get_container_client(self.__get_collection_name(collection_name)) - .query_items(sql_string, partition_key=partition_key) + .query_items(sql_string, partition_key=self.__get_partition_key(partition_key)) ) return list(result_iterable) except CosmosHttpResponseError: diff --git a/airflow/providers/microsoft/azure/provider.yaml b/airflow/providers/microsoft/azure/provider.yaml index fe3e45a92f6dd..2ddc479b63c0a 100644 --- a/airflow/providers/microsoft/azure/provider.yaml +++ b/airflow/providers/microsoft/azure/provider.yaml @@ -79,8 +79,7 @@ dependencies: - apache-airflow>=2.6.0 - adlfs>=2023.10.0 - azure-batch>=8.0.0 - # azure-cosmos 4.6.0 fail on mypy, limit version till we fix the issue - - azure-cosmos>=4.0.0,<4.6.0 + - azure-cosmos>=4.6.0 - azure-mgmt-cosmosdb - azure-datalake-store>=0.0.45 - azure-identity>=1.3.1 diff --git a/docs/apache-airflow-providers-microsoft-azure/connections/azure_cosmos.rst b/docs/apache-airflow-providers-microsoft-azure/connections/azure_cosmos.rst index e7933fc89797f..2d7754f2cce3b 100644 --- a/docs/apache-airflow-providers-microsoft-azure/connections/azure_cosmos.rst +++ b/docs/apache-airflow-providers-microsoft-azure/connections/azure_cosmos.rst @@ -56,6 +56,7 @@ Extra (optional) * ``database_name``: Specify the azure cosmos database to use. * ``collection_name``: Specify the azure cosmos collection to use. + * ``partition_key``: Specify the partition key for the collection. * ``subscription_id``: Specify the ID of the subscription used for the initial connection. Required for falling back to DefaultAzureCredential_ * ``resource_group_name``: Specify the Azure Resource Group Name under which the desired azure cosmos resides. Required for falling back to DefaultAzureCredential_ * ``managed_identity_client_id``: The client ID of a user-assigned managed identity. If provided with `workload_identity_tenant_id`, they'll pass to ``DefaultAzureCredential``. diff --git a/generated/provider_dependencies.json b/generated/provider_dependencies.json index 284ff97cc7616..9315766f81926 100644 --- a/generated/provider_dependencies.json +++ b/generated/provider_dependencies.json @@ -682,7 +682,7 @@ "adlfs>=2023.10.0", "apache-airflow>=2.6.0", "azure-batch>=8.0.0", - "azure-cosmos>=4.0.0,<4.6.0", + "azure-cosmos>=4.6.0", "azure-datalake-store>=0.0.45", "azure-identity>=1.3.1", "azure-keyvault-secrets>=4.1.0", diff --git a/tests/providers/microsoft/azure/hooks/test_cosmos.py b/tests/providers/microsoft/azure/hooks/test_cosmos.py index b63b46410e3f4..bc6b4f4277a98 100644 --- a/tests/providers/microsoft/azure/hooks/test_cosmos.py +++ b/tests/providers/microsoft/azure/hooks/test_cosmos.py @@ -23,6 +23,7 @@ from unittest.mock import PropertyMock import pytest +from azure.cosmos import PartitionKey from azure.cosmos.cosmos_client import CosmosClient from airflow.exceptions import AirflowException @@ -43,6 +44,7 @@ def setup_test_cases(self, create_mock_connection): self.test_collection_name = "test_collection_name" self.test_database_default = "test_database_default" self.test_collection_default = "test_collection_default" + self.test_partition_key = "/test_partition_key" create_mock_connection( Connection( conn_id="azure_cosmos_test_key_id", @@ -52,6 +54,7 @@ def setup_test_cases(self, create_mock_connection): extra={ "database_name": self.test_database_default, "collection_name": self.test_collection_default, + "partition_key": self.test_partition_key, }, ) ) @@ -115,11 +118,11 @@ def test_create_container_exception(self, mock_cosmos): @mock.patch(f"{MODULE}.CosmosClient") def test_create_container(self, mock_cosmos): hook = AzureCosmosDBHook(azure_cosmos_conn_id="azure_cosmos_test_key_id") - hook.create_collection(self.test_collection_name, self.test_database_name) + hook.create_collection(self.test_collection_name, self.test_database_name, partition_key="/id") expected_calls = [ mock.call() .get_database_client("test_database_name") - .create_container("test_collection_name", partition_key=None) + .create_container("test_collection_name", partition_key=PartitionKey(path="/id")) ] mock_cosmos.assert_any_call(self.test_end_point, {"masterKey": self.test_master_key}) mock_cosmos.assert_has_calls(expected_calls) @@ -131,7 +134,9 @@ def test_create_container_default(self, mock_cosmos): expected_calls = [ mock.call() .get_database_client("test_database_name") - .create_container("test_collection_name", partition_key=None) + .create_container( + "test_collection_name", partition_key=PartitionKey(path=self.test_partition_key) + ) ] mock_cosmos.assert_any_call(self.test_end_point, {"masterKey": self.test_master_key}) mock_cosmos.assert_has_calls(expected_calls) diff --git a/tests/providers/microsoft/azure/operators/test_cosmos.py b/tests/providers/microsoft/azure/operators/test_cosmos.py index 231a56c2b1b9d..d3627482bbb3c 100644 --- a/tests/providers/microsoft/azure/operators/test_cosmos.py +++ b/tests/providers/microsoft/azure/operators/test_cosmos.py @@ -35,6 +35,7 @@ def setup_test_cases(self, create_mock_connection): self.test_master_key = "magic_test_key" self.test_database_name = "test_database_name" self.test_collection_name = "test_collection_name" + self.test_partition_key = "test_partition_key" create_mock_connection( Connection( conn_id="azure_cosmos_test_key_id", @@ -44,6 +45,7 @@ def setup_test_cases(self, create_mock_connection): extra={ "database_name": self.test_database_name, "collection_name": self.test_collection_name, + "partition_key": self.test_partition_key, }, ) )