44
55import azure .cosmos .cosmos_client as cosmos_client
66import azure .cosmos .exceptions as exceptions
7+ import flask
78from azure .cosmos import ContainerProxy , PartitionKey
89from flask import Flask
910from werkzeug .exceptions import HTTPException
@@ -101,8 +102,8 @@ def __init__(
101102 raise ValueError ("The cosmos_db module has not been initialized!" )
102103 self .mapper = mapper
103104 self .order_fields = order_fields if order_fields else []
104- self .container : ContainerProxy = self . cosmos_helper . db . get_container_client (
105- container_id
105+ self .container : ContainerProxy = (
106+ self . cosmos_helper . db . get_container_client ( container_id )
106107 )
107108 self .partition_key_attribute = partition_key_attribute
108109
@@ -266,7 +267,6 @@ def find_all(
266267 ),
267268 order_clause = self .create_sql_order_clause (),
268269 )
269-
270270 result = self .container .query_items (
271271 query = query_str ,
272272 parameters = params ,
@@ -277,6 +277,50 @@ def find_all(
277277 function_mapper = self .get_mapper_or_dict (mapper )
278278 return list (map (function_mapper , result ))
279279
280+ def count (
281+ self ,
282+ event_context : EventContext ,
283+ conditions : dict = None ,
284+ custom_sql_conditions : List [str ] = None ,
285+ custom_params : dict = None ,
286+ visible_only = True ,
287+ ):
288+ conditions = conditions if conditions else {}
289+ custom_sql_conditions = (
290+ custom_sql_conditions if custom_sql_conditions else []
291+ )
292+ custom_params = custom_params if custom_params else {}
293+ partition_key_value = self .find_partition_key_value (event_context )
294+ params = [
295+ {"name" : "@partition_key_value" , "value" : partition_key_value },
296+ ]
297+ params .extend (self .generate_params (conditions ))
298+ params .extend (custom_params )
299+ query_str = """
300+ SELECT VALUE COUNT(1) FROM c
301+ WHERE c.{partition_key_attribute}=@partition_key_value
302+ {conditions_clause}
303+ {visibility_condition}
304+ {custom_sql_conditions_clause}
305+ """ .format (
306+ partition_key_attribute = self .partition_key_attribute ,
307+ visibility_condition = self .create_sql_condition_for_visibility (
308+ visible_only
309+ ),
310+ conditions_clause = self .create_sql_where_conditions (conditions ),
311+ custom_sql_conditions_clause = self .create_custom_sql_conditions (
312+ custom_sql_conditions
313+ ),
314+ )
315+
316+ flask .current_app .logger .debug (query_str )
317+ result = self .container .query_items (
318+ query = query_str ,
319+ parameters = params ,
320+ partition_key = partition_key_value ,
321+ )
322+ return result .next ()
323+
280324 def partial_update (
281325 self ,
282326 id : str ,
@@ -286,7 +330,10 @@ def partial_update(
286330 mapper : Callable = None ,
287331 ):
288332 item_data = self .find (
289- id , event_context , visible_only = visible_only , mapper = dict ,
333+ id ,
334+ event_context ,
335+ visible_only = visible_only ,
336+ mapper = dict ,
290337 )
291338 item_data .update (changes )
292339 return self .update (id , item_data , event_context , mapper = mapper )
@@ -304,7 +351,10 @@ def update(
304351 return function_mapper (self .container .replace_item (id , body = item_data ))
305352
306353 def delete (
307- self , id : str , event_context : EventContext , mapper : Callable = None ,
354+ self ,
355+ id : str ,
356+ event_context : EventContext ,
357+ mapper : Callable = None ,
308358 ):
309359 return self .partial_update (
310360 id ,
@@ -327,7 +377,7 @@ def get_mapper_or_dict(self, alternative_mapper: Callable) -> Callable:
327377 def get_page_size_or (self , custom_page_size : int ) -> int :
328378 # TODO The default value should be taken from the Azure Feature Manager
329379 # or any other repository for the settings
330- return custom_page_size or 100
380+ return custom_page_size or 9999
331381
332382 def on_update (self , update_item_data : dict , event_context : EventContext ):
333383 pass
0 commit comments