Query pipeline changes and Offset/Limit/Distinct support

.gitignore

@@ -90,3 +90,4 @@ sdk/storage/azure-storage-blob/tests/settings_real.py
 sdk/storage/azure-storage-queue/tests/settings_real.py
 sdk/storage/azure-storage-file/tests/settings_real.py
 *.code-workspace
+sdk/cosmos/azure-cosmos/test/test_config.py

sdk/cosmos/azure-cosmos/HISTORY.md

@@ -1,4 +1,46 @@
 # Change Log azure-cosmos
+
+## Version 4.0.0b2:
+
+Version 4.0.0b2 is the second iteration in our efforts to build a more Pythonic client library.
+
+**Breaking changes**
+
+- The client connection has been adapted to consume the HTTP pipeline defined in `azure.core.pipeline`.
+- Interactive objects have now been renamed as proxies. This includes:
+    - `Database` -> `DatabaseProxy`
+    - `User` -> `UserProxy`
+    - `Container` -> `ContainerProxy`
+    - `Scripts` -> `ScriptsProxy`
+- The constructor of `CosmosClient` has been updated:
+    - The `auth` parameter has been renamed to `credential` and will now take an authentication type directly. This means the master key value, a dictionary of resource tokens, or a list of permissions can be passed in. However the old dictionary format is still supported.
+    - The `connection_policy` parameter has been made a keyword only parameter, and while it is still supported, each of the individual attributes of the policy can now be passed in as explicit keyword arguments:
+        - `request_timeout`
+        - `media_request_timeout`
+        - `connection_mode`
+        - `media_read_mode`
+        - `proxy_config`
+        - `enable_endpoint_discovery`
+        - `preferred_locations`
+        - `multiple_write_locations`
+- A new classmethod constructor has been added to `CosmosClient` to enable creation via a connection string retrieved from the Azure portal.
+- Some `read_all` operations have been renamed to `list` operations:
+    - `CosmosClient.read_all_databases` -> `CosmosClient.list_databases`
+    - `Container.read_all_conflicts` -> `ContainerProxy.list_conflicts`
+    - `Database.read_all_containers` -> `DatabaseProxy.list_containers`
+    - `Database.read_all_users` -> `DatabaseProxy.list_users`
+    - `User.read_all_permissions` -> `UserProxy.list_permissions`
+- All operations that take `request_options` or `feed_options` parameters, these have been moved to keyword only parameters. In addition, while these options dictionaries are still supported, each of the individual options within the dictionary are now supported as explicit keyword arguments.
+- The error heirarchy is now inherited from `azure.core.AzureError` instead of `CosmosError` which has been removed.
+    - `HTTPFailure` has been renamed to `CosmosHttpResponseError`
+    - `JSONParseFailure` has been removed and replaced by `azure.core.DecodeError`
+    - Added additional errors for specific response codes:
+        - `CosmosResourceNotFoundError` for status 404
+        - `CosmosResourceExistsError` for status 409
+        - `CosmosAccessConditionFailedError` for status 412
+- `CosmosClient` can now be run in a context manager to handle closing the client connection.
+- Iterable responses (e.g. query responses and list responses) are now of type `azure.core.paging.ItemPaged`. The method `fetch_next_block` has been replaced by a secondary iterator, accessed by the `by_page` method.
+
 ## Version 4.0.0b1:
 
 Version 4.0.0b1 is the first preview of our efforts to create a user-friendly and Pythonic client library for Azure Cosmos. For more information about this, and preview releases of other Azure SDK libraries, please visit https://aka.ms/azure-sdk-preview1-python.

sdk/cosmos/azure-cosmos/README.md

@@ -63,23 +63,21 @@ export ACCOUNT_KEY=$(az cosmosdb list-keys --resource-group $RES_GROUP --name $A
 Once you've populated the `ACCOUNT_URI` and `ACCOUNT_KEY` environment variables, you can create the [CosmosClient][ref_cosmosclient].
 
 ```Python
-from azure.cosmos import HTTPFailure, CosmosClient, Container, Database, PartitionKey
+from azure.cosmos import CosmosClient, Container, Database, PartitionKey, errors
 
 import os
 url = os.environ['ACCOUNT_URI']
 key = os.environ['ACCOUNT_KEY']
-client = CosmosClient(url, auth = {
-    'masterKey': key
-})
+client = CosmosClient(url, credential=key)
 ```
 
 ## Usage
 
 Once you've initialized a [CosmosClient][ref_cosmosclient], you can interact with the primary resource types in Cosmos DB:
 
-* [Database][ref_database]: A Cosmos DB account can contain multiple databases. When you create a database, you specify the API you'd like to use when interacting with its documents: SQL, MongoDB, Gremlin, Cassandra, or Azure Table. Use the [Database][ref_database] object to manage its containers.
+* [Database][ref_database]: A Cosmos DB account can contain multiple databases. When you create a database, you specify the API you'd like to use when interacting with its documents: SQL, MongoDB, Gremlin, Cassandra, or Azure Table. Use the [DatabaseProxy][ref_database] object to manage its containers.
 
-* [Container][ref_container]: A container is a collection of JSON documents. You create (insert), read, update, and delete items in a container by using methods on the [Container][ref_container] object.
+* [Container][ref_container]: A container is a collection of JSON documents. You create (insert), read, update, and delete items in a container by using methods on the [ContainerProxy][ref_container] object.
 
 * [Item][ref_item]: An Item is the dictionary-like representation of a JSON document stored in a container. Each Item you add to a container must include an `id` key with a value that uniquely identifies the item within the container.
 
@@ -106,9 +104,7 @@ After authenticating your [CosmosClient][ref_cosmosclient], you can work with an
 database_name = 'testDatabase'
 try:
     database = client.create_database(database_name)
-except HTTPFailure as e:
-    if e.status_code != 409:
-        raise
+except errors.CosmosResourceExistsError:
     database = client.get_database_client(database_name)
 ```
 
@@ -120,13 +116,13 @@ This example creates a container with default settings. If a container with the
 container_name = 'products'
 try:
     container = database.create_container(id=container_name, partition_key=PartitionKey(path="/productName"))
-except HTTPFailure as e:
-    if e.status_code != 409:
-        raise
+except errors.CosmosResourceExistsError:
     container = database.get_container_client(container_name)
+except errors.CosmosHttpResponseError:
+    raise
 ```
 
-The preceding snippet also handles the [HTTPFailure][ref_httpfailure] exception if the container creation failed. For more information on error handling and troubleshooting, see the [Troubleshooting](#troubleshooting) section.
+The preceding snippet also handles the [CosmosHttpResponseError][ref_httpfailure] exception if the container creation failed. For more information on error handling and troubleshooting, see the [Troubleshooting](#troubleshooting) section.
 
 ### Get an existing container
 
@@ -139,7 +135,7 @@ container = database.get_container_client(container_name)
 
 ### Insert data
 
-To insert items into a container, pass a dictionary containing your data to [Container.upsert_item][ref_container_upsert_item]. Each item you add to a container must include an `id` key with a value that uniquely identifies the item within the container.
+To insert items into a container, pass a dictionary containing your data to [ContainerProxy.upsert_item][ref_container_upsert_item]. Each item you add to a container must include an `id` key with a value that uniquely identifies the item within the container.
 
 This example inserts several items into the container, each with a unique `id`:
 
@@ -158,7 +154,7 @@ for i in range(1, 10):
 
 ### Delete data
 
-To delete items from a container, use [Container.delete_item][ref_container_delete_item]. The SQL API in Cosmos DB does not support the SQL `DELETE` statement.
+To delete items from a container, use [ContainerProxy.delete_item][ref_container_delete_item]. The SQL API in Cosmos DB does not support the SQL `DELETE` statement.
 
 ```Python
 for item in container.query_items(query='SELECT * FROM products p WHERE p.productModel = "DISCONTINUED"',
@@ -168,7 +164,7 @@ for item in container.query_items(query='SELECT * FROM products p WHERE p.produc
 
 ### Query the database
 
-A Cosmos DB SQL API database supports querying the items in a container with [Container.query_items][ref_container_query_items] using SQL-like syntax.
+A Cosmos DB SQL API database supports querying the items in a container with [ContainerProxy.query_items][ref_container_query_items] using SQL-like syntax.
 
 This example queries a container for items with a specific `id`:
 
@@ -186,7 +182,7 @@ for item in container.query_items(
 
 > NOTE: Although you can specify any value for the container name in the `FROM` clause, we recommend you use the container name for consistency.
 
-Perform parameterized queries by passing a dictionary containing the parameters and their values to [Container.query_items][ref_container_query_items]:
+Perform parameterized queries by passing a dictionary containing the parameters and their values to [ContainerProxy.query_items][ref_container_query_items]:
 
 ```Python
 discontinued_items = container.query_items(
@@ -243,13 +239,11 @@ For example, if you try to create a container using an ID (name) that's already
 ```Python
 try:
     database.create_container(id=container_name, partition_key=PartitionKey(path="/productName")
-except HTTPFailure as e:
-    if e.status_code == 409:
-        print("""Error creating container.
+except errors.CosmosResourceExistsError:
+    print("""Error creating container
 HTTP status code 409: The ID (name) provided for the container is already in use.
 The container name must be unique within the database.""")
-    else:
-        raise
+
 ```
 
 ## More sample code
@@ -285,7 +279,7 @@ For more extensive documentation on the Cosmos DB service, see the [Azure Cosmos
 [ref_cosmosclient_create_database]: http://cosmosproto.westus.azurecontainer.io/#azure.cosmos.CosmosClient.create_database
 [ref_cosmosclient]: http://cosmosproto.westus.azurecontainer.io/#azure.cosmos.CosmosClient
 [ref_database]: http://cosmosproto.westus.azurecontainer.io/#azure.cosmos.Database
-[ref_httpfailure]: https://docs.microsoft.com/python/api/azure-cosmos/azure.cosmos.errors.httpfailure
+[ref_httpfailure]: https://docs.microsoft.com/python/api/azure-cosmos/azure.cosmos.errors.CosmosHttpResponseError
 [ref_item]: http://cosmosproto.westus.azurecontainer.io/#azure.cosmos.Item
 [sample_database_mgmt]: https://github.com/binderjoe/cosmos-python-prototype/blob/master/examples/databasemanagementsample.py
 [sample_document_mgmt]: https://github.com/binderjoe/cosmos-python-prototype/blob/master/examples/documentmanagementsample.py

sdk/cosmos/azure-cosmos/azure/__init__.py

@@ -1 +1 @@
-__path__ = __import__("pkgutil").extend_path(__path__, __name__)
+__path__ = __import__("pkgutil").extend_path(__path__, __name__)  # type: ignore

sdk/cosmos/azure-cosmos/azure/cosmos/__init__.py

@@ -19,9 +19,11 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
 
-from .container import Container
+from .container import ContainerProxy
 from .cosmos_client import CosmosClient
-from .database import Database
+from .database import DatabaseProxy
+from .user import UserProxy
+from .scripts import ScriptsProxy
 from .documents import (
     ConsistencyLevel,
     DataType,
@@ -35,17 +37,16 @@
 )
 from .partition_key import PartitionKey
 from .permission import Permission
-from .scripts import Scripts
-from .user import User
+from .version import VERSION
 
 __all__ = (
-    "Container",
     "CosmosClient",
-    "Database",
+    "DatabaseProxy",
+    "ContainerProxy",
     "PartitionKey",
     "Permission",
-    "Scripts",
-    "User",
+    "ScriptsProxy",
+    "UserProxy",
     "ConsistencyLevel",
     "DataType",
     "IndexKind",
@@ -56,3 +57,4 @@
     "TriggerOperation",
     "TriggerType",
 )
+__version__ = VERSION

sdk/cosmos/azure-cosmos/azure/cosmos/_base.py

@@ -27,6 +27,7 @@
 import json
 import uuid
 import binascii
+from typing import Dict, Any
 
 import six
 from six.moves.urllib.parse import quote as urllib_quote
@@ -39,6 +40,45 @@
 
 # pylint: disable=protected-access
 
+_COMMON_OPTIONS = {
+    'initial_headers': 'initialHeaders',
+    'pre_trigger_include': 'preTriggerInclude',
+    'post_trigger_include': 'postTriggerInclude',
+    'max_item_count': 'maxItemCount',
+    'access_condition': 'accessCondition',
+    'indexing_directive': 'indexingDirective',
+    'consistency_level': 'consistencyLevel',
+    'session_token': 'sessionToken',
+    'enable_scan_in_query': 'enableScanInQuery',
+    'resource_token_expiry_seconds': 'resourceTokenExpirySeconds',
+    'offer_type': 'offerType',
+    'offer_throughput': 'offerThroughput',
+    'partition_key': 'partitionKey',
+    'enable_cross_partition_query': 'enableCrossPartitionQuery',
+    'populate_query_metrics': 'populateQueryMetrics',
+    'enable_script_logging': 'enableScriptLogging',
+    'offer_enable_ru_per_minute_throughput': 'offerEnableRUPerMinuteThroughput',
+    'disable_ru_per_minute_usage': 'disableRUPerMinuteUsage',
+    'change_feed': 'changeFeed',
+    'continuation': 'continuation',
+    'is_start_from_beginning': 'isStartFromBeginning',
+    'populate_partition_key_range_statistics': 'populatePartitionKeyRangeStatistics',
+    'populate_quota_info': 'populateQuotaInfo'
+}
+
+def build_options(kwargs):
+    # type: (Dict[str, Any]) -> Dict[str, Any]
+    options = kwargs.pop('request_options', kwargs.pop('feed_options', {}))
+    for key, value in _COMMON_OPTIONS.items():
+        if key in kwargs:
+            options[value] = kwargs.pop(key)
+
+    if 'if_match' in kwargs:
+        options['accessCondition'] = {'type': 'IfMatch', 'condition': kwargs.pop('if_match')}
+    if 'if_none_match' in kwargs:
+        options['accessCondition'] = {'type': 'IfNoneMatch', 'condition': kwargs.pop('if_none_match')}
+    return options
+
 
 def GetHeaders(  # pylint: disable=too-many-statements,too-many-branches
     cosmos_client_connection,
@@ -138,6 +178,18 @@ def GetHeaders(  # pylint: disable=too-many-statements,too-many-branches
     if options.get("offerThroughput"):
         headers[http_constants.HttpHeaders.OfferThroughput] = options["offerThroughput"]
 
+    if options.get("contentType"):
+        headers[http_constants.HttpHeaders.ContentType] = options['contentType']
+
+    if options.get("isQueryPlanRequest"):
+        headers[http_constants.HttpHeaders.IsQueryPlanRequest] = options['isQueryPlanRequest']
+
+    if options.get("supportedQueryFeatures"):
+        headers[http_constants.HttpHeaders.SupportedQueryFeatures] = options['supportedQueryFeatures']
+
+    if options.get("queryVersion"):
+        headers[http_constants.HttpHeaders.QueryVersion] = options['queryVersion']
+
     if "partitionKey" in options:
         # if partitionKey value is Undefined, serialize it as [{}] to be consistent with other SDKs.
         if options.get("partitionKey") is partition_key._Undefined: