Clean up public surface (#43124)

* Clean up public surface since list GA

* Fixed return types

* Remove unused imports

* Remove List type

Author: allenkim0129

sdk/cosmos/azure-cosmos/azure/cosmos/_global_endpoint_manager.py

@@ -188,10 +188,10 @@ def _endpoints_health_check(self, **kwargs):
                             self.location_cache.location_unavailability_info_by_endpoint):
                         # if the endpoint is unavailable, we need to lower the timeouts to be more aggressive in the
                         # health check. This helps reduce the time the health check is blocking all requests.
-                        self.client.connection_policy.override_dba_timeouts(constants._Constants
-                                                                            .UnavailableEndpointDBATimeouts,
-                                                                            constants._Constants
-                                                                            .UnavailableEndpointDBATimeouts)
+                        self.client.connection_policy._override_dba_timeouts(constants._Constants
+                                                                             .UnavailableEndpointDBATimeouts,
+                                                                             constants._Constants
+                                                                             .UnavailableEndpointDBATimeouts)
                         self.client._GetDatabaseAccountCheck(endpoint, **kwargs)
                     else:
                         self.client._GetDatabaseAccountCheck(endpoint, **kwargs)
@@ -202,8 +202,8 @@ def _endpoints_health_check(self, **kwargs):
 
                 finally:
                     # after the health check for that endpoint setting the timeouts back to their original values
-                    self.client.connection_policy.override_dba_timeouts(previous_dba_read_timeout,
-                                                                        previous_dba_connection_timeout)
+                    self.client.connection_policy._override_dba_timeouts(previous_dba_read_timeout,
+                                                                         previous_dba_connection_timeout)
         self.location_cache.update_location_cache()
 
     def _GetDatabaseAccountStub(self, endpoint, **kwargs):
@@ -220,15 +220,15 @@ def _GetDatabaseAccountStub(self, endpoint, **kwargs):
             try:
                 # if the endpoint is unavailable, we need to lower the timeouts to be more aggressive in the
                 # health check. This helps reduce the time the health check is blocking all requests.
-                self.client.connection_policy.override_dba_timeouts(constants._Constants
-                                                                    .UnavailableEndpointDBATimeouts,
-                                                                    constants._Constants
-                                                                    .UnavailableEndpointDBATimeouts)
+                self.client.connection_policy._override_dba_timeouts(constants._Constants
+                                                                     .UnavailableEndpointDBATimeouts,
+                                                                     constants._Constants
+                                                                     .UnavailableEndpointDBATimeouts)
                 database_account = self.client.GetDatabaseAccount(endpoint, **kwargs)
             finally:
                 # after the health check for that endpoint setting the timeouts back to their original values
-                self.client.connection_policy.override_dba_timeouts(previous_dba_read_timeout,
-                                                                    previous_dba_connection_timeout)
+                self.client.connection_policy._override_dba_timeouts(previous_dba_read_timeout,
+                                                                     previous_dba_connection_timeout)
         else:
             database_account = self.client.GetDatabaseAccount(endpoint, **kwargs)
         return database_account

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

@@ -65,7 +65,7 @@
 # pylint: disable=missing-client-constructor-parameter-credential,missing-client-constructor-parameter-kwargs
 # pylint: disable=docstring-keyword-should-match-keyword-only
 
-def get_epk_range_for_partition_key(
+def _get_epk_range_for_partition_key(
         container_properties: Dict[str, Any],
         partition_key_value: _PartitionKeyType) -> Range:
     partition_key_obj: PartitionKey = _build_partition_key_from_properties(container_properties)
@@ -657,7 +657,7 @@ def query_items_change_feed(
             partition_key = kwargs.pop("partition_key")
             change_feed_state_context["partitionKey"] = self._set_partition_key(cast(_PartitionKeyType, partition_key))
             change_feed_state_context["partitionKeyFeedRange"] = \
-                get_epk_range_for_partition_key(container_properties, partition_key)
+                _get_epk_range_for_partition_key(container_properties, partition_key)
         if "feed_range" in kwargs:
             change_feed_state_context["feedRange"] = kwargs.pop('feed_range')
         if "continuation" in feed_options:
@@ -1910,7 +1910,7 @@ def feed_range_from_partition_key(self, partition_key: _PartitionKeyType) -> Dic
         """
         container_properties = self._get_properties()
         partition_key_value = self._set_partition_key(partition_key)
-        epk_range_for_partition_key = get_epk_range_for_partition_key(container_properties, partition_key_value)
+        epk_range_for_partition_key = _get_epk_range_for_partition_key(container_properties, partition_key_value)
         return FeedRangeInternalEpk(epk_range_for_partition_key).to_dict()
 
     def is_feed_range_subset(self, parent_feed_range: Dict[str, Any], child_feed_range: Dict[str, Any]) -> bool:

sdk/cosmos/azure-cosmos/azure/cosmos/documents.py

@@ -22,7 +22,7 @@
 """Classes and enums for documents in the Azure Cosmos database service.
 """
 
-from typing import List, Optional, TYPE_CHECKING, Union, Dict, Any
+from typing import Optional, TYPE_CHECKING, Union
 
 from typing_extensions import Literal, TypedDict
 
@@ -75,23 +75,23 @@ def __init__(self) -> None:
         self.ReservedDocumentStorageInMB: int = 0
         self.ProvisionedDocumentStorageInMB: int = 0
         self.ConsistencyPolicy: Optional[UserConsistencyPolicy] = None
-        self._WritableLocations: List[dict] = []
-        self._ReadableLocations: List[dict] = []
+        self._WritableLocations: list[dict[str, str]] = []
+        self._ReadableLocations: list[dict[str, str]] = []
         self._EnableMultipleWritableLocations = False
 
     @property
-    def WritableLocations(self) -> List[Dict[Any, Any]]:
+    def WritableLocations(self) -> list[dict[str, str]]:
         """The list of writable locations for a geo-replicated database account.
         :returns: List of writable locations for the database account.
-        :rtype: List[str]
+        :rtype: list[dict[str, str]]
         """
         return self._WritableLocations
 
     @property
-    def ReadableLocations(self) -> List[Dict[Any, Any]]:
+    def ReadableLocations(self) -> list[dict[str, str]]:
         """The list of readable locations for a geo-replicated database account.
         :returns: List of readable locations for the database account.
-        :rtype: List[str]
+        :rtype: list[dict[str, str]]
         """
         return self._ReadableLocations
 
@@ -307,14 +307,14 @@ class ConnectionPolicy:  # pylint: disable=too-many-instance-attributes
         taking into consideration the order specified in PreferredLocations. The
         locations in this list are specified as the names of the azure Cosmos
         locations like, 'West US', 'East US', 'Central India' and so on.
-    :vartype PreferredLocations: List[str]
+    :vartype PreferredLocations: list[str]
     :ivar ExcludedLocations:
         Gets or sets the excluded locations for geo-replicated database
         accounts. When ExcludedLocations is non-empty, the client will skip this
         set of locations from the final location evaluation. The locations in
         this list are specified as the names of the azure Cosmos locations like,
         'West US', 'East US', 'Central India' and so on.
-    :vartype ExcludedLocations: List[str]
+    :vartype ExcludedLocations: list[str]
     :ivar RetryOptions:
         Gets or sets the retry options to be applied to all requests when
         retrying.
@@ -359,16 +359,16 @@ def __init__(self) -> None:
         self.SSLConfiguration: Optional[SSLConfiguration] = None
         self.ProxyConfiguration: Optional[ProxyConfiguration] = None
         self.EnableEndpointDiscovery: bool = True
-        self.PreferredLocations: List[str] = []
-        self.ExcludedLocations: Optional[List[str]] = None
+        self.PreferredLocations: list[str] = []
+        self.ExcludedLocations: Optional[list[str]] = None
         self.RetryOptions: RetryOptions = RetryOptions()
         self.DisableSSLVerification: bool = False
         self.UseMultipleWriteLocations: bool = False
         self.ConnectionRetryConfiguration: Optional["ConnectionRetryPolicy"] = None
         self.ResponsePayloadOnWriteDisabled: bool = False
         self.RetryNonIdempotentWrites: bool = False
 
-    def override_dba_timeouts(
+    def _override_dba_timeouts(
             self,
             connection_timeout: Optional[int] = None,
             read_timeout: Optional[int] = None

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

@@ -8,7 +8,7 @@
 import azure.cosmos.cosmos_client as cosmos_client
 import test_config
 from azure.cosmos.partition_key import PartitionKey, _get_partition_key_from_partition_key_definition
-from azure.cosmos.container import get_epk_range_for_partition_key
+from azure.cosmos.container import _get_epk_range_for_partition_key
 
 @pytest.mark.cosmosEmulator
 class TestChangeFeedPKVariation(unittest.TestCase):
@@ -263,7 +263,7 @@ def _get_properties_override():
 
             for item in items:
                 try:
-                    epk_range = get_epk_range_for_partition_key(container_properties, item["pk"])
+                    epk_range = _get_epk_range_for_partition_key(container_properties, item["pk"])
                     assert epk_range is not None, f"EPK range should not be None for partition key {item['pk']}."
                 except Exception as e:
                     assert False, f"Failed to get EPK range for partition key {item['pk']}: {str(e)}"