fix cosmos hook static checks by making providing partition_key mandatory

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
 .....
 

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,15 +356,15 @@ 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(
         self,
         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:

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

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``.

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",

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)

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,
                 },
             )
         )