Skip to content

Commit 6f6ba06

Browse files
Copilotbambriz
Copilot
and
bambriz
committed
Replace hardcoded strings with InternalOptions constants in async container
Co-authored-by: bambriz <[email protected]>
1 parent 3469978 commit 6f6ba06

File tree

3 files changed

+63
-49
lines changed

3 files changed

+63
-49
lines changed

‎sdk/cosmos/azure-cosmos/azure/cosmos/_constants.py‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -257,6 +257,18 @@ class InternalOptions:
257257
PREFIX_PARTITION_KEY_VALUE: Literal["prefixPartitionKeyValue"] = "prefixPartitionKeyValue"
258258
"""Prefix partition key value for queries."""
259259

260+
DISABLE_AUTOMATIC_ID_GENERATION: Literal["disableAutomaticIdGeneration"] = "disableAutomaticIdGeneration"
261+
"""Whether to disable automatic ID generation for documents."""
262+
263+
FILTER_PREDICATE: Literal["filterPredicate"] = "filterPredicate"
264+
"""Filter predicate for query operations."""
265+
266+
CHANGE_FEED_STATE_CONTEXT: Literal["changeFeedStateContext"] = "changeFeedStateContext"
267+
"""Change feed state context for change feed operations."""
268+
269+
CONTAINER_PROPERTIES: Literal["containerProperties"] = "containerProperties"
270+
"""Container properties for requests."""
271+
260272
class Kwargs:
261273
"""Public-facing keyword argument names used in the azure-cosmos package
262274

‎sdk/cosmos/azure-cosmos/azure/cosmos/aio/_container.py‎

Lines changed: 50 additions & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,7 @@
4747
from .._change_feed.feed_range_internal import FeedRangeInternalEpk
4848
from .._cosmos_responses import CosmosDict, CosmosList
4949
from .._constants import _Constants as Constants
50+
InternalOptions = Constants.InternalOptions
5051
from .._routing.routing_range import Range
5152
from .._session_token_helpers import get_latest_session_token
5253
from ..offer import ThroughputProperties
@@ -107,8 +108,8 @@ def __repr__(self) -> str:
107108

108109
async def _get_properties_with_options(self, options: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
109110
kwargs = {}
110-
if options and "excludedLocations" in options:
111-
kwargs[Constants.Kwargs.EXCLUDED_LOCATIONS] = options['excludedLocations']
111+
if options and InternalOptions.EXCLUDED_LOCATIONS in options:
112+
kwargs[Constants.Kwargs.EXCLUDED_LOCATIONS] = options[InternalOptions.EXCLUDED_LOCATIONS]
112113
return await self._get_properties(**kwargs)
113114

114115
async def _get_properties(self, **kwargs: Any) -> Dict[str, Any]:
@@ -203,9 +204,9 @@ async def read(
203204
kwargs[Constants.Kwargs.INITIAL_HEADERS] = initial_headers
204205
request_options = _build_options(kwargs)
205206
if populate_partition_key_range_statistics is not None:
206-
request_options["populatePartitionKeyRangeStatistics"] = populate_partition_key_range_statistics
207+
request_options[InternalOptions.POPULATE_PARTITION_KEY_RANGE_STATISTICS] = populate_partition_key_range_statistics
207208
if populate_quota_info is not None:
208-
request_options["populateQuotaInfo"] = populate_quota_info
209+
request_options[InternalOptions.POPULATE_QUOTA_INFO] = populate_quota_info
209210
container = await self.client_connection.ReadContainer(self.container_link, options=request_options, **kwargs)
210211
# Only cache Container Properties that will not change in the lifetime of the container
211212
self.client_connection._set_container_properties_cache(self.container_link, # pylint: disable=protected-access
@@ -293,11 +294,11 @@ async def create_item(
293294
if throughput_bucket is not None:
294295
kwargs[Constants.Kwargs.THROUGHPUT_BUCKET] = throughput_bucket
295296
request_options = _build_options(kwargs)
296-
request_options["disableAutomaticIdGeneration"] = not enable_automatic_id_generation
297+
request_options[InternalOptions.DISABLE_AUTOMATIC_ID_GENERATION] = not enable_automatic_id_generation
297298
if indexing_directive is not None:
298-
request_options["indexingDirective"] = indexing_directive
299+
request_options[InternalOptions.INDEXING_DIRECTIVE] = indexing_directive
299300
await self._get_properties_with_options(request_options)
300-
request_options["containerRID"] = self.__get_client_container_caches()[self.container_link]["_rid"]
301+
request_options[InternalOptions.CONTAINER_RID] = self.__get_client_container_caches()[self.container_link]["_rid"]
301302

302303
result = await self.client_connection.CreateItem(
303304
database_or_container_link=self.container_link, document=body, options=request_options, **kwargs
@@ -367,12 +368,12 @@ async def read_item(
367368
kwargs[Constants.Kwargs.THROUGHPUT_BUCKET] = throughput_bucket
368369
request_options = _build_options(kwargs)
369370

370-
request_options["partitionKey"] = await self._set_partition_key(partition_key)
371+
request_options[InternalOptions.PARTITION_KEY] = await self._set_partition_key(partition_key)
371372
if max_integrated_cache_staleness_in_ms is not None:
372373
validate_cache_staleness_value(max_integrated_cache_staleness_in_ms)
373-
request_options["maxIntegratedCacheStaleness"] = max_integrated_cache_staleness_in_ms
374+
request_options[InternalOptions.MAX_INTEGRATED_CACHE_STALENESS] = max_integrated_cache_staleness_in_ms
374375
await self._get_properties_with_options(request_options)
375-
request_options["containerRID"] = self.__get_client_container_caches()[self.container_link]["_rid"]
376+
request_options[InternalOptions.CONTAINER_RID] = self.__get_client_container_caches()[self.container_link]["_rid"]
376377

377378
return await self.client_connection.ReadItem(document_link=doc_link, options=request_options, **kwargs)
378379

@@ -419,15 +420,15 @@ def read_all_items(
419420
kwargs[Constants.Kwargs.THROUGHPUT_BUCKET] = throughput_bucket
420421
feed_options = _build_options(kwargs)
421422
if max_item_count is not None:
422-
feed_options["maxItemCount"] = max_item_count
423+
feed_options[InternalOptions.MAX_ITEM_COUNT] = max_item_count
423424
if max_integrated_cache_staleness_in_ms:
424425
validate_cache_staleness_value(max_integrated_cache_staleness_in_ms)
425-
feed_options["maxIntegratedCacheStaleness"] = max_integrated_cache_staleness_in_ms
426+
feed_options[InternalOptions.MAX_INTEGRATED_CACHE_STALENESS] = max_integrated_cache_staleness_in_ms
426427
response_hook = kwargs.pop("response_hook", None)
427428
if response_hook and hasattr(response_hook, "clear"):
428429
response_hook.clear()
429430
if self.container_link in self.__get_client_container_caches():
430-
feed_options["containerRID"] = self.__get_client_container_caches()[self.container_link]["_rid"]
431+
feed_options[InternalOptions.CONTAINER_RID] = self.__get_client_container_caches()[self.container_link]["_rid"]
431432
kwargs["containerProperties"] = self._get_properties_with_options
432433

433434
items = self.client_connection.ReadItems(
@@ -669,20 +670,20 @@ def query_items(
669670

670671
# Update 'feed_options' from 'kwargs'
671672
if utils.valid_key_value_exist(kwargs, "max_item_count"):
672-
feed_options["maxItemCount"] = kwargs.pop(Constants.Kwargs.MAX_ITEM_COUNT)
673+
feed_options[InternalOptions.MAX_ITEM_COUNT] = kwargs.pop(Constants.Kwargs.MAX_ITEM_COUNT)
673674
if utils.valid_key_value_exist(kwargs, "populate_query_metrics"):
674-
feed_options["populateQueryMetrics"] = kwargs.pop(Constants.Kwargs.POPULATE_QUERY_METRICS)
675+
feed_options[InternalOptions.POPULATE_QUERY_METRICS] = kwargs.pop(Constants.Kwargs.POPULATE_QUERY_METRICS)
675676
if utils.valid_key_value_exist(kwargs, "populate_index_metrics"):
676-
feed_options["populateIndexMetrics"] = kwargs.pop(Constants.Kwargs.POPULATE_INDEX_METRICS)
677+
feed_options[InternalOptions.POPULATE_INDEX_METRICS] = kwargs.pop(Constants.Kwargs.POPULATE_INDEX_METRICS)
677678
if utils.valid_key_value_exist(kwargs, "enable_scan_in_query"):
678-
feed_options["enableScanInQuery"] = kwargs.pop(Constants.Kwargs.ENABLE_SCAN_IN_QUERY)
679+
feed_options[InternalOptions.ENABLE_SCAN_IN_QUERY] = kwargs.pop(Constants.Kwargs.ENABLE_SCAN_IN_QUERY)
679680
if utils.valid_key_value_exist(kwargs, "max_integrated_cache_staleness_in_ms"):
680681
max_integrated_cache_staleness_in_ms = kwargs.pop(Constants.Kwargs.MAX_INTEGRATED_CACHE_STALENESS)
681682
validate_cache_staleness_value(max_integrated_cache_staleness_in_ms)
682-
feed_options["maxIntegratedCacheStaleness"] = max_integrated_cache_staleness_in_ms
683+
feed_options[InternalOptions.MAX_INTEGRATED_CACHE_STALENESS] = max_integrated_cache_staleness_in_ms
683684
if utils.valid_key_value_exist(kwargs, "continuation_token_limit"):
684-
feed_options["responseContinuationTokenLimitInKb"] = kwargs.pop(Constants.Kwargs.RESPONSE_CONTINUATION_TOKEN_LIMIT_IN_KB)
685-
feed_options["correlatedActivityId"] = GenerateGuidId()
685+
feed_options[InternalOptions.RESPONSE_CONTINUATION_TOKEN_LIMIT_IN_KB] = kwargs.pop(Constants.Kwargs.RESPONSE_CONTINUATION_TOKEN_LIMIT_IN_KB)
686+
feed_options[InternalOptions.CORRELATED_ACTIVITY_ID] = GenerateGuidId()
686687

687688
# Set query with 'query' and 'parameters' from kwargs
688689
if utils.valid_key_value_exist(kwargs, "parameters"):
@@ -698,10 +699,10 @@ def query_items(
698699
# If 'partition_key' is provided, set 'partitionKey' in 'feed_options'
699700
if utils.valid_key_value_exist(kwargs, "partition_key"):
700701
partition_key = kwargs.pop("partition_key")
701-
feed_options["partitionKey"] = self._set_partition_key(partition_key)
702+
feed_options[InternalOptions.PARTITION_KEY] = self._set_partition_key(partition_key)
702703
# If 'partition_key' or 'feed_range' is not provided, set 'enableCrossPartitionQuery' to True
703704
elif not utils.valid_key_value_exist(kwargs, "feed_range"):
704-
feed_options["enableCrossPartitionQuery"] = True
705+
feed_options[InternalOptions.ENABLE_CROSS_PARTITION_QUERY] = True
705706

706707
# Set 'response_hook'
707708
response_hook = kwargs.pop("response_hook", None)
@@ -938,15 +939,15 @@ def query_items_change_feed( # pylint: disable=unused-argument
938939
if "continuation" in feed_options:
939940
change_feed_state_context["continuation"] = feed_options.pop("continuation")
940941

941-
feed_options["changeFeedStateContext"] = change_feed_state_context
942-
feed_options["containerProperties"] = self._get_properties_with_options(feed_options)
942+
feed_options[InternalOptions.CHANGE_FEED_STATE_CONTEXT] = change_feed_state_context
943+
feed_options[InternalOptions.CONTAINER_PROPERTIES] = self._get_properties_with_options(feed_options)
943944

944945
response_hook = kwargs.pop("response_hook", None)
945946
if hasattr(response_hook, "clear"):
946947
response_hook.clear()
947948

948949
if self.container_link in self.__get_client_container_caches():
949-
feed_options["containerRID"] = self.__get_client_container_caches()[self.container_link]["_rid"]
950+
feed_options[InternalOptions.CONTAINER_RID] = self.__get_client_container_caches()[self.container_link]["_rid"]
950951

951952
result = self.client_connection.QueryItemsChangeFeed(
952953
self.container_link, options=feed_options, response_hook=response_hook, **kwargs
@@ -1026,9 +1027,9 @@ async def upsert_item(
10261027
if throughput_bucket is not None:
10271028
kwargs[Constants.Kwargs.THROUGHPUT_BUCKET] = throughput_bucket
10281029
request_options = _build_options(kwargs)
1029-
request_options["disableAutomaticIdGeneration"] = True
1030+
request_options[InternalOptions.DISABLE_AUTOMATIC_ID_GENERATION] = True
10301031
await self._get_properties_with_options(request_options)
1031-
request_options["containerRID"] = self.__get_client_container_caches()[self.container_link]["_rid"]
1032+
request_options[InternalOptions.CONTAINER_RID] = self.__get_client_container_caches()[self.container_link]["_rid"]
10321033

10331034
result = await self.client_connection.UpsertItem(
10341035
database_or_container_link=self.container_link,
@@ -1115,9 +1116,9 @@ async def replace_item(
11151116
if throughput_bucket is not None:
11161117
kwargs[Constants.Kwargs.THROUGHPUT_BUCKET] = throughput_bucket
11171118
request_options = _build_options(kwargs)
1118-
request_options["disableAutomaticIdGeneration"] = True
1119+
request_options[InternalOptions.DISABLE_AUTOMATIC_ID_GENERATION] = True
11191120
await self._get_properties_with_options(request_options)
1120-
request_options["containerRID"] = self.__get_client_container_caches()[self.container_link]["_rid"]
1121+
request_options[InternalOptions.CONTAINER_RID] = self.__get_client_container_caches()[self.container_link]["_rid"]
11211122

11221123
result = await self.client_connection.ReplaceItem(
11231124
document_link=item_link, new_document=body, options=request_options, **kwargs
@@ -1202,12 +1203,12 @@ async def patch_item(
12021203
if throughput_bucket is not None:
12031204
kwargs[Constants.Kwargs.THROUGHPUT_BUCKET] = throughput_bucket
12041205
request_options = _build_options(kwargs)
1205-
request_options["disableAutomaticIdGeneration"] = True
1206-
request_options["partitionKey"] = await self._set_partition_key(partition_key)
1206+
request_options[InternalOptions.DISABLE_AUTOMATIC_ID_GENERATION] = True
1207+
request_options[InternalOptions.PARTITION_KEY] = await self._set_partition_key(partition_key)
12071208
if filter_predicate is not None:
1208-
request_options["filterPredicate"] = filter_predicate
1209+
request_options[InternalOptions.FILTER_PREDICATE] = filter_predicate
12091210
await self._get_properties_with_options(request_options)
1210-
request_options["containerRID"] = self.__get_client_container_caches()[self.container_link]["_rid"]
1211+
request_options[InternalOptions.CONTAINER_RID] = self.__get_client_container_caches()[self.container_link]["_rid"]
12111212

12121213
item_link = self._get_document_link(item)
12131214
result = await self.client_connection.PatchItem(
@@ -1283,9 +1284,9 @@ async def delete_item(
12831284
if throughput_bucket is not None:
12841285
kwargs[Constants.Kwargs.THROUGHPUT_BUCKET] = throughput_bucket
12851286
request_options = _build_options(kwargs)
1286-
request_options["partitionKey"] = await self._set_partition_key(partition_key)
1287+
request_options[InternalOptions.PARTITION_KEY] = await self._set_partition_key(partition_key)
12871288
await self._get_properties_with_options(request_options)
1288-
request_options["containerRID"] = self.__get_client_container_caches()[self.container_link]["_rid"]
1289+
request_options[InternalOptions.CONTAINER_RID] = self.__get_client_container_caches()[self.container_link]["_rid"]
12891290

12901291
document_link = self._get_document_link(item)
12911292
await self.client_connection.DeleteItem(document_link=document_link, options=request_options, **kwargs)
@@ -1379,9 +1380,9 @@ def list_conflicts(
13791380
"""
13801381
feed_options = _build_options(kwargs)
13811382
if max_item_count is not None:
1382-
feed_options["maxItemCount"] = max_item_count
1383+
feed_options[InternalOptions.MAX_ITEM_COUNT] = max_item_count
13831384
if self.container_link in self.__get_client_container_caches():
1384-
feed_options["containerRID"] = self.__get_client_container_caches()[self.container_link]["_rid"]
1385+
feed_options[InternalOptions.CONTAINER_RID] = self.__get_client_container_caches()[self.container_link]["_rid"]
13851386

13861387
result = self.client_connection.ReadConflicts(
13871388
collection_link=self.container_link, feed_options=feed_options, **kwargs
@@ -1417,13 +1418,13 @@ def query_conflicts(
14171418
"""
14181419
feed_options = _build_options(kwargs)
14191420
if max_item_count is not None:
1420-
feed_options["maxItemCount"] = max_item_count
1421+
feed_options[InternalOptions.MAX_ITEM_COUNT] = max_item_count
14211422
if partition_key is not None:
1422-
feed_options["partitionKey"] = self._set_partition_key(partition_key)
1423+
feed_options[InternalOptions.PARTITION_KEY] = self._set_partition_key(partition_key)
14231424
else:
1424-
feed_options["enableCrossPartitionQuery"] = True
1425+
feed_options[InternalOptions.ENABLE_CROSS_PARTITION_QUERY] = True
14251426
if self.container_link in self.__get_client_container_caches():
1426-
feed_options["containerRID"] = self.__get_client_container_caches()[self.container_link]["_rid"]
1427+
feed_options[InternalOptions.CONTAINER_RID] = self.__get_client_container_caches()[self.container_link]["_rid"]
14271428

14281429
result = self.client_connection.QueryConflicts(
14291430
collection_link=self.container_link,
@@ -1455,9 +1456,9 @@ async def get_conflict(
14551456
:rtype: Dict[str, Any]
14561457
"""
14571458
request_options = _build_options(kwargs)
1458-
request_options["partitionKey"] = await self._set_partition_key(partition_key)
1459+
request_options[InternalOptions.PARTITION_KEY] = await self._set_partition_key(partition_key)
14591460
if self.container_link in self.__get_client_container_caches():
1460-
request_options["containerRID"] = self.__get_client_container_caches()[self.container_link]["_rid"]
1461+
request_options[InternalOptions.CONTAINER_RID] = self.__get_client_container_caches()[self.container_link]["_rid"]
14611462
result = await self.client_connection.ReadConflict(
14621463
conflict_link=self._get_conflict_link(conflict), options=request_options, **kwargs
14631464
)
@@ -1485,7 +1486,7 @@ async def delete_conflict(
14851486
:rtype: None
14861487
"""
14871488
request_options = _build_options(kwargs)
1488-
request_options["partitionKey"] = await self._set_partition_key(partition_key)
1489+
request_options[InternalOptions.PARTITION_KEY] = await self._set_partition_key(partition_key)
14891490

14901491
await self.client_connection.DeleteConflict(
14911492
conflict_link=self._get_conflict_link(conflict), options=request_options, **kwargs
@@ -1544,9 +1545,9 @@ async def delete_all_items_by_partition_key(
15441545
kwargs[Constants.Kwargs.THROUGHPUT_BUCKET] = throughput_bucket
15451546
request_options = _build_options(kwargs)
15461547
# regardless if partition key is valid we set it as invalid partition keys are set to a default empty value
1547-
request_options["partitionKey"] = await self._set_partition_key(partition_key)
1548+
request_options[InternalOptions.PARTITION_KEY] = await self._set_partition_key(partition_key)
15481549
await self._get_properties_with_options(request_options)
1549-
request_options["containerRID"] = self.__get_client_container_caches()[self.container_link]["_rid"]
1550+
request_options[InternalOptions.CONTAINER_RID] = self.__get_client_container_caches()[self.container_link]["_rid"]
15501551

15511552
await self.client_connection.DeleteAllItemsByPartitionKey(collection_link=self.container_link,
15521553
options=request_options, **kwargs)
@@ -1611,10 +1612,10 @@ async def execute_item_batch(
16111612
if throughput_bucket is not None:
16121613
kwargs[Constants.Kwargs.THROUGHPUT_BUCKET] = throughput_bucket
16131614
request_options = _build_options(kwargs)
1614-
request_options["partitionKey"] = await self._set_partition_key(partition_key)
1615-
request_options["disableAutomaticIdGeneration"] = True
1615+
request_options[InternalOptions.PARTITION_KEY] = await self._set_partition_key(partition_key)
1616+
request_options[InternalOptions.DISABLE_AUTOMATIC_ID_GENERATION] = True
16161617
await self._get_properties_with_options(request_options)
1617-
request_options["containerRID"] = self.__get_client_container_caches()[self.container_link]["_rid"]
1618+
request_options[InternalOptions.CONTAINER_RID] = self.__get_client_container_caches()[self.container_link]["_rid"]
16181619

16191620
return await self.client_connection.Batch(
16201621
collection_link=self.container_link, batch_operations=batch_operations, options=request_options, **kwargs)

‎sdk/cosmos/azure-cosmos/azure/cosmos/container.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,7 @@
4343
)
4444
from ._change_feed.feed_range_internal import FeedRangeInternalEpk
4545
from ._constants import _Constants as Constants
46+
InternalOptions = Constants.InternalOptions
4647
from ._cosmos_client_connection import CosmosClientConnection
4748
from ._cosmos_responses import CosmosDict, CosmosList
4849
from ._routing.routing_range import Range

0 commit comments

Comments
 (0)