fix: fixing a bug where exclusion list was not honored before falling back to global endpoint. (#43297)

* fix: fixing a bug where exclusion list was not honored before falling back to global endpoint.

* fix: updated changelog

* fix: addressing comments

Author: dibahlfi

sdk/cosmos/azure-cosmos/CHANGELOG.md

@@ -10,7 +10,8 @@
 * Changed `retry_write` from `bool` to `int` to match other retryable options. See [PR 43341](https://github.com/Azure/azure-sdk-for-python/pull/43341).
 
 #### Bugs Fixed
-
+* Fixed bug where exclusion list was not honored before falling back to global endpoint for multi-write region accounts. See[PR 43297](https://github.com/Azure/azure-sdk-for-python/pull/43297)
+ 
 #### Other Changes
 * Removed dual endpoint tracking from the sdk. See [PR 40451](https://github.com/Azure/azure-sdk-for-python/pull/40451).
 * Reverted typehints to fix the mismatch issue. See [PR 43124](https://github.com/Azure/azure-sdk-for-python/pull/43124)

sdk/cosmos/azure-cosmos/azure/cosmos/_location_cache.py

@@ -105,7 +105,7 @@ def _get_applicable_regional_routing_contexts(regional_routing_contexts: list[Re
     if base.IsMasterResource(resource_type):
         applicable_regional_routing_contexts.extend(excluded_regional_routing_contexts)
 
-    # If applicable_regional_routing_contexts is empty add fallback regional routing context
+    # If all preferred locations are excluded, use the fallback endpoint.
     if not applicable_regional_routing_contexts:
         applicable_regional_routing_contexts.append(fall_back_regional_routing_context)
 
@@ -397,19 +397,22 @@ def get_preferred_regional_routing_contexts(
                     # can use multiple write locations, preferred locations list should be
                     # used for determining both read and write endpoints order.
                     for location in self.effective_preferred_locations:
-                        regional_endpoint = endpoints_by_location[location] if location in endpoints_by_location \
-                            else None
+                        regional_endpoint = endpoints_by_location.get(location)
                         if regional_endpoint:
                             if self.is_endpoint_unavailable(regional_endpoint.get_primary(),
                                                             expected_available_operation):
                                 unavailable_endpoints.append(regional_endpoint)
                             else:
                                 regional_endpoints.append(regional_endpoint)
 
+                # If all preferred locations are unavailable, honor the preferred list by trying them anyway.
+                if not regional_endpoints and unavailable_endpoints:
+                    regional_endpoints.extend(unavailable_endpoints)
+
+                # If there are no preferred locations or none of the preferred locations are in the account,
+                # add the fallback endpoint.
                 if not regional_endpoints:
                     regional_endpoints.append(fallback_endpoint)
-
-                regional_endpoints.extend(unavailable_endpoints)
             else:
                 for location in orderedLocations:
                     if location and location in endpoints_by_location:

sdk/cosmos/azure-cosmos/tests/test_location_cache.py

@@ -146,8 +146,10 @@ def test_resolve_request_endpoint_preferred_regions(self):
         lc.mark_endpoint_unavailable_for_write(location3_endpoint, True)
         read_resolved = lc.resolve_service_endpoint(read_doc_request)
         write_resolved = lc.resolve_service_endpoint(write_doc_request)
-        assert read_resolved == write_resolved
-        assert read_resolved == default_endpoint
+        # With the updated logic, we should retry the unavailable preferred locations
+        # instead of falling back to the default endpoint.
+        assert read_resolved == location1_endpoint
+        assert write_resolved == location1_endpoint
 
     @pytest.mark.parametrize("test_type",["OnClient", "OnRequest", "OnBoth"])
     def test_get_applicable_regional_endpoints_excluded_regions(self, test_type):
@@ -266,6 +268,92 @@ def test_set_excluded_locations_for_requests(self):
         assert str(
             e.value) == expected_error_message
 
+    def test_resolve_endpoint_unavailable_and_excluded_preferred_regions(self):
+        # Scenario: All preferred read locations are unavailable AND in the excluded list.
+        # Expected: Fallback to the primary write region.
+        connection_policy = documents.ConnectionPolicy()
+        connection_policy.ExcludedLocations = [location1_name, location4_name]
+        lc = refresh_location_cache([location1_name, location4_name], True, connection_policy)
+        db_acc = create_database_account(True)
+        lc.perform_on_database_account_read(db_acc)
+
+        # Mark all preferred read locations as unavailable
+        lc.mark_endpoint_unavailable_for_read(location1_endpoint, True)
+        lc.mark_endpoint_unavailable_for_read(location4_endpoint, True)
+
+        # Create a read request
+        read_doc_request = RequestObject(ResourceType.Document, _OperationType.Read, None)
+
+        # Resolve the endpoint for the read request
+        read_doc_resolved = lc.resolve_service_endpoint(read_doc_request)
+
+        # All preferred read locations ([loc1, loc4]) are excluded.
+        # The fallback for read is the primary write region, which is loc1.
+        assert read_doc_resolved == location1_endpoint
+
+        # Scenario: All preferred write locations are unavailable AND in the excluded list.
+        # Expected: Fallback to the default endpoint.
+        connection_policy.ExcludedLocations = [location1_name, location2_name]
+        lc = refresh_location_cache([location1_name, location2_name], True, connection_policy)
+        db_acc = create_database_account(True)
+        lc.perform_on_database_account_read(db_acc)
+
+        # Mark preferred write locations as unavailable
+        lc.mark_endpoint_unavailable_for_write(location1_endpoint, True)
+        lc.mark_endpoint_unavailable_for_write(location2_endpoint, True)
+
+        # Create a write request
+        write_doc_request = RequestObject(ResourceType.Document, _OperationType.Create, None)
+
+        # Resolve the endpoint for the write request
+        write_doc_resolved = lc.resolve_service_endpoint(write_doc_request)
+
+        # All preferred write locations ([loc1, loc2]) are excluded.
+        # The fallback for write is the default_endpoint.
+        assert write_doc_resolved == default_endpoint
+
+    def test_resolve_endpoint_unavailable_and_excluded_on_request(self):
+        # Scenario: All preferred read locations are unavailable AND in the excluded list on the request.
+        # Expected: Fallback to the primary write region.
+        lc = refresh_location_cache([location1_name, location4_name], True)
+        db_acc = create_database_account(True)
+        lc.perform_on_database_account_read(db_acc)
+
+        # Mark all preferred read locations as unavailable
+        lc.mark_endpoint_unavailable_for_read(location1_endpoint, True)
+        lc.mark_endpoint_unavailable_for_read(location4_endpoint, True)
+
+        # Create a read request and set excluded locations
+        read_doc_request = RequestObject(ResourceType.Document, _OperationType.Read, None)
+        read_doc_request.excluded_locations = [location1_name, location4_name]
+
+        # Resolve the endpoint for the read request
+        read_doc_resolved = lc.resolve_service_endpoint(read_doc_request)
+
+        # All preferred read locations ([loc1, loc4]) are excluded.
+        # The fallback for read is the primary write region, which is loc1.
+        assert read_doc_resolved == location1_endpoint
+
+        # Scenario: All preferred write locations are unavailable AND in the excluded list on the request.
+        # Expected: Fallback to the default endpoint.
+        lc = refresh_location_cache([location1_name, location2_name], True)
+        db_acc = create_database_account(True)
+        lc.perform_on_database_account_read(db_acc)
+
+        # Mark preferred write locations as unavailable
+        lc.mark_endpoint_unavailable_for_write(location1_endpoint, True)
+        lc.mark_endpoint_unavailable_for_write(location2_endpoint, True)
+
+        # Create a write request and set excluded locations
+        write_doc_request = RequestObject(ResourceType.Document, _OperationType.Create, None)
+        write_doc_request.excluded_locations = [location1_name, location2_name]
+
+        # Resolve the endpoint for the write request
+        write_doc_resolved = lc.resolve_service_endpoint(write_doc_request)
+
+        # All preferred write locations ([loc1, loc2]) are excluded.
+        # The fallback for write is the default_endpoint.
+        assert write_doc_resolved == default_endpoint
 
 if __name__ == "__main__":
     unittest.main()