diff --git a/airflow/providers/microsoft/azure/hooks/cosmos.py b/airflow/providers/microsoft/azure/hooks/cosmos.py index 3a73eeada7d7d..4b42217e68ff9 100644 --- a/airflow/providers/microsoft/azure/hooks/cosmos.py +++ b/airflow/providers/microsoft/azure/hooks/cosmos.py @@ -25,7 +25,6 @@ """ from __future__ import annotations -import json import uuid from typing import Any @@ -145,12 +144,14 @@ def does_collection_exist(self, collection_name: str, database_name: str) -> boo if collection_name is None: raise AirflowBadRequest("Collection name cannot be None.") + # The ignores below is due to typing bug in azure-cosmos 9.2.0 + # https://github.com/Azure/azure-sdk-for-python/issues/31811 existing_container = list( self.get_conn() .get_database_client(self.__get_database_name(database_name)) .query_containers( "SELECT * FROM r WHERE r.id=@id", - parameters=[json.dumps({"name": "@id", "value": collection_name})], + parameters=[{"name": "@id", "value": collection_name}], # type: ignore[list-item] ) ) if not existing_container: @@ -170,12 +171,14 @@ def create_collection( # We need to check to see if this container already exists so we don't try # to create it twice + # The ignores below is due to typing bug in azure-cosmos 9.2.0 + # https://github.com/Azure/azure-sdk-for-python/issues/31811 existing_container = list( self.get_conn() .get_database_client(self.__get_database_name(database_name)) .query_containers( "SELECT * FROM r WHERE r.id=@id", - parameters=[json.dumps({"name": "@id", "value": collection_name})], + parameters=[{"name": "@id", "value": collection_name}], # type: ignore[list-item] ) ) @@ -190,10 +193,12 @@ def does_database_exist(self, database_name: str) -> bool: if database_name is None: raise AirflowBadRequest("Database name cannot be None.") + # The ignores below is due to typing bug in azure-cosmos 9.2.0 + # https://github.com/Azure/azure-sdk-for-python/issues/31811 existing_database = list( self.get_conn().query_databases( "SELECT * FROM r WHERE r.id=@id", - parameters=[json.dumps({"name": "@id", "value": database_name})], + parameters=[{"name": "@id", "value": database_name}], # type: ignore[list-item] ) ) if not existing_database: @@ -208,10 +213,12 @@ def create_database(self, database_name: str) -> None: # We need to check to see if this database already exists so we don't try # to create it twice + # The ignores below is due to typing bug in azure-cosmos 9.2.0 + # https://github.com/Azure/azure-sdk-for-python/issues/31811 existing_database = list( self.get_conn().query_databases( "SELECT * FROM r WHERE r.id=@id", - parameters=[json.dumps({"name": "@id", "value": database_name})], + parameters=[{"name": "@id", "value": database_name}], # type: ignore[list-item] ) )