Fix Client Settings issue with missing available regions (#40763)

* Fixed Client Available Regions in logging policy

* Added emulator tests for client setting in logs

* Disabled pylint errors

Author: allenkim0129

sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_http_logging_policy.py

@@ -34,9 +34,10 @@
 from azure.core.pipeline import PipelineRequest, PipelineResponse
 from azure.core.pipeline.policies import HttpLoggingPolicy
 
+from ._location_cache import LocationCache
 from .http_constants import HttpHeaders
 from ._global_endpoint_manager import _GlobalEndpointManager
-from .documents import DatabaseAccount
+from .documents import DatabaseAccount, ConnectionPolicy
 
 if TYPE_CHECKING:
     from azure.core.rest import HttpRequest, HttpResponse, AsyncHttpResponse
@@ -215,7 +216,7 @@ def on_request(
                 logger.info(log_string, extra=cosmos_logger_attributes)
 
             except Exception as err:  # pylint: disable=broad-except
-                logger.warning("Failed to log request: %s", repr(err))
+                logger.warning("Failed to log request: %s", repr(err)) #pylint: disable=do-not-log-exceptions-if-not-debug
             return
         super().on_request(request)
 
@@ -299,25 +300,32 @@ def on_response(  # pylint: disable=too-many-statements, too-many-branches
                     log_string += "\nResponse error message: {}".format(_format_error(http_response.text()))
                 logger.info(log_string, extra=log_data)
             except Exception as err:  # pylint: disable=broad-except
-                logger.warning("Failed to log response: %s", repr(err), extra=log_data)
+                logger.warning("Failed to log response: %s", repr(err), extra=log_data) #pylint: disable=do-not-log-exceptions-if-not-debug
             return
         super().on_response(request, response)
 
     def __get_client_settings(self) -> Optional[Dict[str, Any]]:
         # Place any client settings we want to log here
-        client_info: Dict[str, Any] = {"Client Preferred Regions": [], "Client Available Read Regions": [],
-                                       "Client Available Write Regions": []}
+        client_preferred_regions = []
+        client_excluded_regions = []
+        client_account_read_regions = []
+        client_account_write_regions = []
+
         if self.__global_endpoint_manager:
-            client_info['Client Preferred Regions'] = self.__global_endpoint_manager.PreferredLocations \
-                if hasattr(self.__global_endpoint_manager, "PreferredLocations") else []
-            try:
-                location_cache = self.__global_endpoint_manager.location_cache
-                client_info["Client Available Read Regions"] = location_cache.available_read_locations
-                client_info["Client Available Write Regions"] = location_cache.available_write_locations
-            except AttributeError:
-                client_info["Client Available Read Regions"] = []
-                client_info["Client Available Write Regions"] = []
-        return client_info
+            if self.__global_endpoint_manager.Client and self.__global_endpoint_manager.Client.connection_policy:
+                connection_policy: ConnectionPolicy = self.__global_endpoint_manager.Client.connection_policy
+                client_preferred_regions = connection_policy.PreferredLocations
+                client_excluded_regions = connection_policy.ExcludedLocations
+
+            if self.__global_endpoint_manager.location_cache:
+                location_cache: LocationCache = self.__global_endpoint_manager.location_cache
+                client_account_read_regions = location_cache.account_read_locations
+                client_account_write_regions = location_cache.account_write_locations
+
+        return {"Client Preferred Regions": client_preferred_regions,
+                "Client Excluded Regions": client_excluded_regions,
+                "Client Account Read Regions": client_account_read_regions,
+                "Client Account Write Regions": client_account_write_regions}
 
     def __get_database_account_settings(self) -> Optional[DatabaseAccount]:
         if self.__global_endpoint_manager and hasattr(self.__global_endpoint_manager, '_database_account_cache'):

sdk/cosmos/azure-cosmos/azure/cosmos/aio/_global_endpoint_manager_async.py

@@ -103,7 +103,7 @@ async def refresh_endpoint_list(self, database_account, **kwargs):
                 await self.refresh_task
                 self.refresh_task = None
             except (Exception, asyncio.CancelledError) as exception: #pylint: disable=broad-exception-caught
-                logger.exception("Health check task failed: %s", exception)
+                logger.exception("Health check task failed: %s", exception) #pylint: disable=do-not-use-logging-exception
         if self.location_cache.current_time_millis() - self.last_refresh_time > self.refresh_time_interval_in_ms:
             self.refresh_needed = True
         if self.refresh_needed:

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

@@ -10,7 +10,12 @@
 import pytest
 
 import azure.cosmos.cosmos_client as cosmos_client
+from azure.cosmos.container import ContainerProxy
 import test_config
+from _fault_injection_transport import FaultInjectionTransport
+from test_fault_injection_transport import TestFaultInjectionTransport
+from typing import List, Callable
+from azure.core.rest import HttpRequest
 
 try:
     from unittest.mock import Mock
@@ -30,8 +35,32 @@ def reset(self):
     def emit(self, record):
         self.messages.append(record)
 
-
-
+CONFIG = test_config.TestConfig
+L1 = "Location1"
+L2 = "Location2"
+L1_URL = test_config.TestConfig.local_host
+L2_URL = L1_URL.replace("localhost", "127.0.0.1")
+URL_TO_LOCATIONS = {
+    L1_URL: L1,
+    L2_URL: L2}
+
+
+def create_logger(name: str, mock_handler: MockHandler, level: int = logging.INFO) -> logging.Logger:
+    logger = logging.getLogger(name)
+    logger.addHandler(mock_handler)
+    logger.setLevel(level)
+
+    return logger
+
+def get_locations_list(msg: str) -> List[str]:
+    msg = msg.replace(' ', '')
+    msg = msg.replace('\'', '')
+    # Find the substring between the first '[' and the last ']'
+    start = msg.find('[') + 1
+    end = msg.rfind(']')
+    # Extract the substring and convert it to a list using ast.literal_eval
+    msg = msg[start:end]
+    return msg.split(',')
 
 @pytest.mark.cosmosEmulator
 class TestCosmosHttpLogger(unittest.TestCase):
@@ -54,12 +83,8 @@ def setUpClass(cls):
                 "tests.")
         cls.mock_handler_default = MockHandler()
         cls.mock_handler_diagnostic = MockHandler()
-        cls.logger_default = logging.getLogger("testloggerdefault")
-        cls.logger_default.addHandler(cls.mock_handler_default)
-        cls.logger_default.setLevel(logging.INFO)
-        cls.logger_diagnostic = logging.getLogger("testloggerdiagnostic")
-        cls.logger_diagnostic.addHandler(cls.mock_handler_diagnostic)
-        cls.logger_diagnostic.setLevel(logging.INFO)
+        cls.logger_default = create_logger("testloggerdefault", cls.mock_handler_default)
+        cls.logger_diagnostic = create_logger("testloggerdiagnostic", cls.mock_handler_diagnostic)
         cls.client_default = cosmos_client.CosmosClient(cls.host, cls.masterKey,
                                                         consistency_level="Session",
                                                         connection_policy=cls.connectionPolicy,
@@ -136,6 +161,65 @@ def test_cosmos_http_logging_policy(self):
 
         self.mock_handler_diagnostic.reset()
 
+    def test_client_settings(self):
+        # Test data
+        all_locations = [L1, L2]
+        client_excluded_locations = [L1]
+        multiple_write_locations = True
+
+        # Client setup
+        mock_handler = MockHandler()
+        logger = create_logger("test_logger_client_settings", mock_handler)
+
+        custom_transport = FaultInjectionTransport()
+        is_get_account_predicate: Callable[[HttpRequest], bool] = lambda \
+                r: FaultInjectionTransport.predicate_is_database_account_call(r)
+        emulator_as_multi_write_region_account_transformation = \
+            lambda r, inner: FaultInjectionTransport.transform_topology_mwr(
+                first_region_name=L1,
+                second_region_name=L2,
+                inner=inner,
+                first_region_url=L1_URL,
+                second_region_url=L2_URL,
+            )
+        custom_transport.add_response_transformation(
+            is_get_account_predicate,
+            emulator_as_multi_write_region_account_transformation)
+
+        initialized_objects = TestFaultInjectionTransport.setup_method_with_custom_transport(
+            custom_transport,
+            default_endpoint=CONFIG.host,
+            key=CONFIG.masterKey,
+            database_id=CONFIG.TEST_DATABASE_ID,
+            container_id=CONFIG.TEST_SINGLE_PARTITION_CONTAINER_ID,
+            preferred_locations=all_locations,
+            excluded_locations=client_excluded_locations,
+            multiple_write_locations=multiple_write_locations,
+            custom_logger=logger
+        )
+        mock_handler.reset()
+
+        # create an item
+        id_value: str = str(uuid.uuid4())
+        document_definition = {'id': id_value, 'pk': id_value}
+        container: ContainerProxy = initialized_objects["col"]
+        container.create_item(body=document_definition)
+
+        # Verify endpoint locations
+        messages_split = mock_handler.messages[1].message.split("\n")
+        for message in messages_split:
+            if "Client Preferred Regions:" in message:
+                locations = get_locations_list(message)
+                assert all_locations == locations
+            elif "Client Excluded Regions:" in message:
+                locations = get_locations_list(message)
+                assert client_excluded_locations == locations
+            elif "Client Account Read Regions:" in message:
+                locations = get_locations_list(message)
+                assert all_locations == locations
+            elif "Client Account Write Regions:" in message:
+                locations = get_locations_list(message)
+                assert all_locations == locations
 
 if __name__ == "__main__":
     unittest.main()

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

@@ -68,9 +68,10 @@ def setup_method_with_custom_transport(
             key: str = master_key,
             database_id: str = TEST_DATABASE_ID,
             container_id: str = SINGLE_PARTITION_CONTAINER_NAME,
+            custom_logger = logger,
             **kwargs):
         client = CosmosClient(default_endpoint, key, consistency_level="Session",
-                              transport=custom_transport, logger=logger, enable_diagnostics_logging=True, **kwargs)
+                              transport=custom_transport, logger=custom_logger, enable_diagnostics_logging=True, **kwargs)
         db: DatabaseProxy = client.get_database_client(database_id)
         container: ContainerProxy = db.get_container_client(container_id)
         return {"client": client, "db": db, "col": container}