|
45 | 45 | if TYPE_CHECKING: |
46 | 46 | from ._cosmos_client_connection import CosmosClientConnection |
47 | 47 | from .aio._cosmos_client_connection_async import CosmosClientConnection as AsyncClientConnection |
| 48 | + from ._request_object import RequestObject |
48 | 49 |
|
| 50 | +# pylint: disable=protected-access |
49 | 51 |
|
50 | 52 | _COMMON_OPTIONS = { |
51 | 53 | 'initial_headers': 'initialHeaders', |
@@ -174,37 +176,9 @@ def GetHeaders( # pylint: disable=too-many-statements,too-many-branches |
174 | 176 | if options.get("indexingDirective"): |
175 | 177 | headers[http_constants.HttpHeaders.IndexingDirective] = options["indexingDirective"] |
176 | 178 |
|
177 | | - consistency_level = None |
178 | | - |
179 | | - # get default client consistency level |
180 | | - default_client_consistency_level = headers.get(http_constants.HttpHeaders.ConsistencyLevel) |
181 | | - |
182 | | - # set consistency level. check if set via options, this will override the default |
| 179 | + # set request consistency level - if session consistency, the client should be setting this on its own |
183 | 180 | if options.get("consistencyLevel"): |
184 | | - consistency_level = options["consistencyLevel"] |
185 | | - # TODO: move this line outside of if-else cause to remove the code duplication |
186 | | - headers[http_constants.HttpHeaders.ConsistencyLevel] = consistency_level |
187 | | - elif default_client_consistency_level is not None: |
188 | | - consistency_level = default_client_consistency_level |
189 | | - headers[http_constants.HttpHeaders.ConsistencyLevel] = consistency_level |
190 | | - |
191 | | - # figure out if consistency level for this request is session |
192 | | - is_session_consistency = consistency_level == documents.ConsistencyLevel.Session |
193 | | - |
194 | | - # set session token if required |
195 | | - if is_session_consistency is True and not IsMasterResource(resource_type): |
196 | | - # if there is a token set via option, then use it to override default |
197 | | - if options.get("sessionToken"): |
198 | | - headers[http_constants.HttpHeaders.SessionToken] = options["sessionToken"] |
199 | | - else: |
200 | | - # check if the client's default consistency is session (and request consistency level is same), |
201 | | - # then update from session container |
202 | | - if default_client_consistency_level == documents.ConsistencyLevel.Session and \ |
203 | | - cosmos_client_connection.session: |
204 | | - # populate session token from the client's session container |
205 | | - headers[http_constants.HttpHeaders.SessionToken] = cosmos_client_connection.session.get_session_token( |
206 | | - path |
207 | | - ) |
| 181 | + headers[http_constants.HttpHeaders.ConsistencyLevel] = options["consistencyLevel"] |
208 | 182 |
|
209 | 183 | if options.get("enableScanInQuery"): |
210 | 184 | headers[http_constants.HttpHeaders.EnableScanInQuery] = options["enableScanInQuery"] |
@@ -349,6 +323,77 @@ def GetHeaders( # pylint: disable=too-many-statements,too-many-branches |
349 | 323 |
|
350 | 324 | return headers |
351 | 325 |
|
| 326 | +def _is_session_token_request( |
| 327 | + cosmos_client_connection: Union["CosmosClientConnection", "AsyncClientConnection"], |
| 328 | + headers: dict, |
| 329 | + request_object: "RequestObject") -> bool: |
| 330 | + consistency_level = headers.get(http_constants.HttpHeaders.ConsistencyLevel) |
| 331 | + # Figure out if consistency level for this request is session |
| 332 | + is_session_consistency = consistency_level == documents.ConsistencyLevel.Session |
| 333 | + |
| 334 | + # Verify that it is not a metadata request, and that it is either a read request, batch request, or an account |
| 335 | + # configured to use multiple write regions. Batch requests are special-cased because they can contain both read and |
| 336 | + # write operations, and we want to use session consistency for the read operations. |
| 337 | + return (is_session_consistency is True and cosmos_client_connection.session is not None |
| 338 | + and not IsMasterResource(request_object.resource_type) |
| 339 | + and (documents._OperationType.IsReadOnlyOperation(request_object.operation_type) |
| 340 | + or request_object.operation_type == "Batch" |
| 341 | + or cosmos_client_connection._global_endpoint_manager.can_use_multiple_write_locations(request_object))) |
| 342 | + |
| 343 | + |
| 344 | +def set_session_token_header( |
| 345 | + cosmos_client_connection: Union["CosmosClientConnection", "AsyncClientConnection"], |
| 346 | + headers: dict, |
| 347 | + path: str, |
| 348 | + request_object: "RequestObject", |
| 349 | + options: Mapping[str, Any], |
| 350 | + partition_key_range_id: Optional[str] = None) -> None: |
| 351 | + # set session token if required |
| 352 | + if _is_session_token_request(cosmos_client_connection, headers, request_object): |
| 353 | + # if there is a token set via option, then use it to override default |
| 354 | + if options.get("sessionToken"): |
| 355 | + headers[http_constants.HttpHeaders.SessionToken] = options["sessionToken"] |
| 356 | + else: |
| 357 | + # check if the client's default consistency is session (and request consistency level is same), |
| 358 | + # then update from session container |
| 359 | + if headers[http_constants.HttpHeaders.ConsistencyLevel] == documents.ConsistencyLevel.Session and \ |
| 360 | + cosmos_client_connection.session: |
| 361 | + # populate session token from the client's session container |
| 362 | + session_token = ( |
| 363 | + cosmos_client_connection.session.get_session_token(path, |
| 364 | + options.get('partitionKey'), |
| 365 | + cosmos_client_connection._container_properties_cache, |
| 366 | + cosmos_client_connection._routing_map_provider, |
| 367 | + partition_key_range_id)) |
| 368 | + if session_token != "": |
| 369 | + headers[http_constants.HttpHeaders.SessionToken] = session_token |
| 370 | + |
| 371 | +async def set_session_token_header_async( |
| 372 | + cosmos_client_connection: Union["CosmosClientConnection", "AsyncClientConnection"], |
| 373 | + headers: dict, |
| 374 | + path: str, |
| 375 | + request_object: "RequestObject", |
| 376 | + options: Mapping[str, Any], |
| 377 | + partition_key_range_id: Optional[str] = None) -> None: |
| 378 | + # set session token if required |
| 379 | + if _is_session_token_request(cosmos_client_connection, headers, request_object): |
| 380 | + # if there is a token set via option, then use it to override default |
| 381 | + if options.get("sessionToken"): |
| 382 | + headers[http_constants.HttpHeaders.SessionToken] = options["sessionToken"] |
| 383 | + else: |
| 384 | + # check if the client's default consistency is session (and request consistency level is same), |
| 385 | + # then update from session container |
| 386 | + if headers[http_constants.HttpHeaders.ConsistencyLevel] == documents.ConsistencyLevel.Session and \ |
| 387 | + cosmos_client_connection.session: |
| 388 | + # populate session token from the client's session container |
| 389 | + session_token = \ |
| 390 | + await cosmos_client_connection.session.get_session_token_async(path, |
| 391 | + options.get('partitionKey'), |
| 392 | + cosmos_client_connection._container_properties_cache, |
| 393 | + cosmos_client_connection._routing_map_provider, |
| 394 | + partition_key_range_id) |
| 395 | + if session_token != "": |
| 396 | + headers[http_constants.HttpHeaders.SessionToken] = session_token |
352 | 397 |
|
353 | 398 | def GetResourceIdOrFullNameFromLink(resource_link: str) -> str: |
354 | 399 | """Gets resource id or full name from resource link. |
|
0 commit comments