Skip to content

[Cosmos] Session container fixes new branch #41678

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
wants to merge 71 commits into
base: main
Choose a base branch
from
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
71 commits
Select commit Hold shift + click to select a range
2550197
session container fixes
simorenoh Apr 3, 2025
bc80922
Merge branch 'main' into session_token_fixes
simorenoh May 4, 2025
6002825
Merge branch 'main' into session_token_fixes
simorenoh May 8, 2025
9cf54e7
async changes
simorenoh May 8, 2025
2320353
sync changes
simorenoh May 9, 2025
dcc11dd
mypy/pylint
simorenoh May 9, 2025
10c2c9a
Update _session.py
simorenoh May 9, 2025
2c7a2db
mark query plan as fetched for query
simorenoh May 9, 2025
1215129
Merge branch 'main' into session_token_fixes
simorenoh May 9, 2025
31a927e
adjust logic after merging
simorenoh May 12, 2025
73fa94c
Update _base.py
simorenoh May 12, 2025
450070a
Update _base.py
simorenoh May 12, 2025
7411502
tests - missing sync mwr
simorenoh May 12, 2025
b2fb943
sync mwr tests, test fixes
simorenoh May 13, 2025
6899cdb
Update test_session_async.py
simorenoh May 13, 2025
a79be96
Merge branch 'main' into session_token_fixes
simorenoh May 19, 2025
209d3d7
Update test_backwards_compatibility.py
simorenoh May 21, 2025
e80d2c7
Update test_backwards_compatibility.py
simorenoh May 22, 2025
24c36b8
Update test_backwards_compatibility_async.py
simorenoh May 23, 2025
d037304
Update test_backwards_compatibility_async.py
simorenoh May 23, 2025
a78d07d
Update test_backwards_compatibility.py
simorenoh May 23, 2025
d815374
Merge branch 'main' into session_token_fixes
simorenoh May 27, 2025
3ab0a06
Merge branch 'main' into session_token_fixes
simorenoh May 27, 2025
9ca6054
Update execution_dispatcher.py
simorenoh May 27, 2025
b806e20
merge leftovers
simorenoh May 27, 2025
34d63a8
slip
simorenoh May 27, 2025
82ff60c
Update test_backwards_compatibility.py
simorenoh May 28, 2025
b06ad7d
Update test_backwards_compatibility.py
simorenoh May 28, 2025
eb885dd
Update test_query_hybrid_search_async.py
simorenoh May 28, 2025
008ac69
further changes, changelog
simorenoh May 28, 2025
a800f28
Merge branch 'main' into session_token_fixes
simorenoh May 28, 2025
7775860
add tests
simorenoh May 28, 2025
da9eaf1
Merge branch 'session_token_fixes' of https://github.com/Azure/azure-…
simorenoh May 28, 2025
d4dd2d4
typehint
simorenoh May 28, 2025
7522407
Merge branch 'main' into session_token_fixes
simorenoh May 29, 2025
6c124bc
address comments
simorenoh May 29, 2025
2854cb5
Update _session.py
simorenoh May 29, 2025
cce381f
Update _base.py
simorenoh May 29, 2025
b20a89f
ci tests
simorenoh May 29, 2025
78f3827
Merge branch 'main' into session_token_fixes
simorenoh May 29, 2025
a5b18fe
merging main
simorenoh May 29, 2025
bc3d853
Update _cosmos_client_connection.py
simorenoh May 29, 2025
575ea0c
tests
simorenoh May 29, 2025
2a19516
change session token logic
simorenoh May 30, 2025
f94098f
Update _session.py
simorenoh May 30, 2025
883c861
small fixes
simorenoh May 30, 2025
b274126
Update _session.py
simorenoh May 30, 2025
3918a5c
Update _session.py
simorenoh May 30, 2025
3327652
updates
simorenoh May 30, 2025
9e08f93
Merge branch 'main' into session_token_fixes
simorenoh May 30, 2025
c7b5ca5
update tests that checked for compound tokens - this will no longer b…
simorenoh May 30, 2025
9417196
Update test_config.py
simorenoh Jun 2, 2025
525ad6c
Update execution_dispatcher.py
simorenoh Jun 2, 2025
d1a014f
remove query logic
simorenoh Jun 2, 2025
5c825e7
remove partition split testing
simorenoh Jun 2, 2025
76f707b
oylint
simorenoh Jun 2, 2025
29ea4b1
Update dev_requirements.txt
simorenoh Jun 4, 2025
8a5def6
Merge branch 'main' into session_token_fixes
simorenoh Jun 9, 2025
9123dfe
delete duplicate
simorenoh Jun 10, 2025
ad35674
Merge branch 'main' into session_token_fixes
simorenoh Jun 10, 2025
d49c983
Merge branch 'main' into session_token_fixes
simorenoh Jun 11, 2025
8b22152
Merge branch 'main' into session_token_fixes
simorenoh Jun 11, 2025
c7fddc4
Merge remote-tracking branch 'upstream/main' into session_token_fixes
bambriz Jun 12, 2025
8be9ffe
Merge remote-tracking branch 'upstream/main' into session_token_fixes
bambriz Jun 17, 2025
69af540
update tests for partition split
bambriz Jun 19, 2025
bca0aa8
Merge remote-tracking branch 'upstream/main' into session_token_fixes…
bambriz Jun 20, 2025
db5233f
change timeout of split partition key tests
bambriz Jun 21, 2025
af0b1e0
update test
bambriz Jun 23, 2025
a644f34
Merge remote-tracking branch 'upstream/main' into session_token_fixes…
bambriz Jun 27, 2025
d7364ac
Test fixes
bambriz Jun 30, 2025
b52c4ac
update tests
bambriz Jul 1, 2025
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions sdk/cosmos/azure-cosmos/CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -28,8 +28,11 @@
#### Bugs Fixed
* Fixed how resource tokens are parsed for metadata calls in the lifecycle of a document operation. See [PR 40302](https://github.com/Azure/azure-sdk-for-python/pull/40302).
* Fixed issue where Query Change Feed did not return items if the container uses legacy Hash V1 Partition Keys. This also fixes issues with not being able to change feed query for Specific Partition Key Values for HPK. See [PR 41270](https://github.com/Azure/azure-sdk-for-python/pull/41270/)
* Fixed session container compound session token logic. The SDK will now only send the relevant partition-local session tokens for each read request, as opposed to the entire compound session token for the container. See [PR 40366](https://github.com/Azure/azure-sdk-for-python/pull/40366).
* Write requests for single-write region accounts will no longer send session tokens when using session consistency. See [PR 40366](https://github.com/Azure/azure-sdk-for-python/pull/40366).

#### Other Changes
* Cross-partition queries will now always send a query plan before attempting to execute. See [PR 40366](https://github.com/Azure/azure-sdk-for-python/pull/40366).
* Added Client Generated Activity IDs to all Requests. Cosmos Diagnostics Logs will more clearly show the Activity ID for each request and response. [PR 41013](https://github.com/Azure/azure-sdk-for-python/pull/41013)

### 4.12.0b1 (2025-05-19)
Expand Down
104 changes: 74 additions & 30 deletions sdk/cosmos/azure-cosmos/azure/cosmos/_base.py
Original file line number Diff line number Diff line change
Expand Up @@ -44,7 +44,9 @@
if TYPE_CHECKING:
from ._cosmos_client_connection import CosmosClientConnection
from .aio._cosmos_client_connection_async import CosmosClientConnection as AsyncClientConnection
from ._request_object import RequestObject

# pylint: disable=protected-access

_COMMON_OPTIONS = {
'initial_headers': 'initialHeaders',
Expand Down Expand Up @@ -172,37 +174,9 @@ def GetHeaders( # pylint: disable=too-many-statements,too-many-branches
if options.get("indexingDirective"):
headers[http_constants.HttpHeaders.IndexingDirective] = options["indexingDirective"]

consistency_level = None

# get default client consistency level
default_client_consistency_level = headers.get(http_constants.HttpHeaders.ConsistencyLevel)

# set consistency level. check if set via options, this will override the default
# set request consistency level - if session consistency, the client should be setting this on its own
if options.get("consistencyLevel"):
consistency_level = options["consistencyLevel"]
# TODO: move this line outside of if-else cause to remove the code duplication
headers[http_constants.HttpHeaders.ConsistencyLevel] = consistency_level
elif default_client_consistency_level is not None:
consistency_level = default_client_consistency_level
headers[http_constants.HttpHeaders.ConsistencyLevel] = consistency_level

# figure out if consistency level for this request is session
is_session_consistency = consistency_level == documents.ConsistencyLevel.Session

# set session token if required
if is_session_consistency is True and not IsMasterResource(resource_type):
# if there is a token set via option, then use it to override default
if options.get("sessionToken"):
headers[http_constants.HttpHeaders.SessionToken] = options["sessionToken"]
else:
# check if the client's default consistency is session (and request consistency level is same),
# then update from session container
if default_client_consistency_level == documents.ConsistencyLevel.Session and \
cosmos_client_connection.session:
# populate session token from the client's session container
headers[http_constants.HttpHeaders.SessionToken] = cosmos_client_connection.session.get_session_token(
path
)
headers[http_constants.HttpHeaders.ConsistencyLevel] = options["consistencyLevel"]

if options.get("enableScanInQuery"):
headers[http_constants.HttpHeaders.EnableScanInQuery] = options["enableScanInQuery"]
Expand Down Expand Up @@ -345,6 +319,76 @@ def GetHeaders( # pylint: disable=too-many-statements,too-many-branches

return headers

def _is_session_token_request(
cosmos_client_connection: Union["CosmosClientConnection", "AsyncClientConnection"],
headers: dict,
request_object: "RequestObject") -> bool:
consistency_level = headers.get(http_constants.HttpHeaders.ConsistencyLevel)
# Figure out if consistency level for this request is session
is_session_consistency = consistency_level == documents.ConsistencyLevel.Session

# Verify that it is not a metadata request, and that it is either a read request, batch request, or an account
# configured to use multiple write regions
return (is_session_consistency is True and cosmos_client_connection.session is not None
and not IsMasterResource(request_object.resource_type)
and (documents._OperationType.IsReadOnlyOperation(request_object.operation_type)
or request_object.operation_type == "Batch"
or cosmos_client_connection._global_endpoint_manager.can_use_multiple_write_locations(request_object)))


def set_session_token_header(
cosmos_client_connection: Union["CosmosClientConnection", "AsyncClientConnection"],
headers: dict,
path: str,
request_object: "RequestObject",
options: Mapping[str, Any],
partition_key_range_id: Optional[str] = None) -> None:
# set session token if required
if _is_session_token_request(cosmos_client_connection, headers, request_object):
# if there is a token set via option, then use it to override default
if options.get("sessionToken"):
headers[http_constants.HttpHeaders.SessionToken] = options["sessionToken"]
else:
# check if the client's default consistency is session (and request consistency level is same),
# then update from session container
if headers[http_constants.HttpHeaders.ConsistencyLevel] == documents.ConsistencyLevel.Session and \
cosmos_client_connection.session:
# populate session token from the client's session container
session_token = (
cosmos_client_connection.session.get_session_token(path,
options.get('partitionKey'),
cosmos_client_connection._container_properties_cache,
cosmos_client_connection._routing_map_provider,
partition_key_range_id))
if session_token != "":
headers[http_constants.HttpHeaders.SessionToken] = session_token

async def set_session_token_header_async(
cosmos_client_connection: Union["CosmosClientConnection", "AsyncClientConnection"],
headers: dict,
path: str,
request_object: "RequestObject",
options: Mapping[str, Any],
partition_key_range_id: Optional[str] = None) -> None:
# set session token if required
if _is_session_token_request(cosmos_client_connection, headers, request_object):
# if there is a token set via option, then use it to override default
if options.get("sessionToken"):
headers[http_constants.HttpHeaders.SessionToken] = options["sessionToken"]
else:
# check if the client's default consistency is session (and request consistency level is same),
# then update from session container
if headers[http_constants.HttpHeaders.ConsistencyLevel] == documents.ConsistencyLevel.Session and \
cosmos_client_connection.session:
# populate session token from the client's session container
session_token = \
await cosmos_client_connection.session.get_session_token_async(path,
options.get('partitionKey'),
cosmos_client_connection._container_properties_cache,
cosmos_client_connection._routing_map_provider,
partition_key_range_id)
if session_token != "":
headers[http_constants.HttpHeaders.SessionToken] = session_token

def GetResourceIdOrFullNameFromLink(resource_link: str) -> str:
"""Gets resource id or full name from resource link.
Expand Down
Loading
Loading