Build AI document stores with Azure Cosmos DB for NoSQL
Design partitions and throughput, connect securely with the Python SDK, implement CRUD and optimistic concurrency, and write efficient SQL queries for recommendation and RAG workloads.
Suggested study time: 95 minutes • Intermediate • Complete original rewrite with a concise version of every topic, assessment review, and guided RAG lab
By João Ricardo Dutra••Complete original content
1. Match a flexible document store to an AI access pattern
Recommendation engines and retrieval-augmented generation systems often store catalogs, preferences, interaction history, model outputs, and document chunks as JSON. for NoSQL combines a flexible schema, automatic indexing, global distribution, and throughput that can scale independently from storage. Those qualities help when an AI product evolves faster than a rigid relational schema or receives unpredictable traffic.
The design is still driven by access patterns. A responsive solution must decide how accounts, databases, containers, and items are organized; select a partition key that distributes work; choose manual or autoscale throughput; authenticate securely; and prefer point reads or targeted queries when possible. Indexing is automatic by default, but specialized ORDER BY patterns can require a composite index.
Explain the resource hierarchy and its configuration boundaries.
Implement secure SDK access and CRUD operations.
Choose point reads or queries from known identifiers and filters.
Write SQL-like queries that project, filter, sort, aggregate, and control RU consumption.
Topic summary
Start from the AI workload’s read and write patterns, then align data shape, partitioning, throughput, authentication, and query strategy.
2. Navigate the account, database, container, and item hierarchy
An account is the top-level management boundary and exposes a unique DNS endpoint to SDKs and APIs. Account settings include the default consistency level, network access policy, and replicated regions. Separate accounts can isolate production, staging, and development.
Databases are logical namespaces for related containers and can host shared throughput. Containers store JSON items and are the primary scalability boundary: each defines a partition key and can contain documents with different structures. Items are the application records. Product, user, inference-cache, and document-chunk containers can therefore coexist in one database while retaining different access and scaling policies.
Configuration flows from the account boundary to databases and containers, while items carry the application data.
Topic summary
The account supplies the endpoint and global settings, databases group containers, containers scale and partition, and items hold JSON documents.
3. Select a partition key that distributes data and serves queries
The partition-key path identifies a JSON property, while each value creates a logical partition. hashes those values and maps logical partitions to managed physical partitions. The combination of item id and partition-key value identifies an item. Because the partition key cannot be changed in place after container creation, it deserves early workload analysis.
Choose a stable property present in every item.
Favor high cardinality and even storage and RU distribution.
Align the key with common equality filters and point reads.
Use userId or tenantId when activity naturally groups by user or tenant.
Avoid booleans, skewed categories, or time-only values that create hot partitions.
A recommendation catalog might use categoryId only if categories are balanced; interaction logs commonly benefit from userId. Nested paths such as /metadata/region are valid. If the chosen key must change, create another container and migrate or copy the data.
Topic summary
A good partition key is immutable, high-cardinality, evenly used, and present in the operations that dominate the workload.
4. Provision manual or autoscale throughput in RU/s
Request Units normalize CPU, memory, and I/O consumed by reads, writes, queries, and stored procedures. Provisioned throughput is measured in RU/s. Container-level throughput reserves capacity for a critical workload; database-level throughput lets several containers share capacity when their peaks occur at different times.
Throughput choices.
Mode
Behavior
Useful when
Manual
Keeps a fixed RU/s value; dedicated containers start at 400 RU/s.
Demand is predictable and stable.
Autoscale
Moves from 10% of the configured maximum to that maximum; the initial maximum is at least 1,000 RU/s.
Inference, promotions, or ingestion create variable peaks.
Shared database throughput
Distributes one pool across multiple containers.
Several containers have complementary usage patterns.
Balanced keys let storage and throughput scale horizontally without concentrating requests in one hot partition.
Topic summary
Select throughput scope and mode from demand variability, isolation needs, and cost, then watch for 429 throttling and hot partitions.
5. Understand items, system properties, indexing, and request cost
Every item must supply an id. Within a logical partition, id must be unique; id plus the partition-key value identifies the record across the container. adds _rid for internal identity, _self for its resource URI, _etag for optimistic concurrency, _ts for the last-update Unix time, and the legacy _attachments path.
Automatic indexing speeds flexible querying but adds write overhead. Item size, property count, indexing policy, consistency level, filters, sorting, aggregation, projected fields, and the number of partitions all affect RU charge. A point read of a 1-KB item is about 1 RU, while large cross-partition aggregates can cost far more. Read x-ms-request-charge and use or Cosmos DB insights to track the real workload.
Topic summary
Use meaningful ids and system metadata deliberately, and measure RU charge instead of estimating query or indexing cost by intuition.
6. Connect with the SDK and choose production authentication
Official SDKs exist for .NET, Python, JavaScript, Java, and Go. CosmosClient is the entry point and manages connections, routing, failover, and endpoint refresh. The Python azure-cosmos package follows the same resource concepts as the other language SDKs.
Account keys are shared secrets with broad account access. Primary and secondary keys permit rotation, but leaked keys are difficult to scope and audit. is the production preference: grant least-privilege RBAC roles to a user, group, service principal, or managed identity. Cosmos DB Built-in Data Reader supports reads; Cosmos DB Built-in Data Contributor supports data read/write. DefaultAzureCredential can use a developer’s Azure CLI or identity locally and managed identity after deployment.
Create CosmosClient with the account endpoint and prefer plus scoped RBAC over distributing account keys.
7. Reuse clients and create resources idempotently
Keep one CosmosClient for the application lifetime. Recreating it per request discards connection pools and cached routing information, increases latency variance, and can exhaust connections. In Flask or FastAPI, initialize it once at startup and retain frequently used database and container clients.
get_database_client() and get_container_client() return lightweight handles without checking the network. create_database() and create_container() fail if the identifier already exists, which is useful when duplicates signal an error. Their *_if_not_exists variants make repeated startup and test setup idempotent. Containers require a PartitionKey; offer_throughput sets dedicated manual capacity, while ThroughputProperties configures autoscale.
Reuse a singleton client, understand that resource handles are lazy, and choose strict create or idempotent create-if-missing methods intentionally.
8. Create, upsert, replace, and protect concurrent updates
create_item() inserts a new item and returns HTTP 409 when the same id and partition key already exist. upsert_item() inserts or replaces, which suits cache refreshes and external synchronization. replace_item() requires an existing record and is the better choice when absence should remain an error.
_etag changes whenever an item changes. Supply the previously read value through if_match during replacement; a mismatch raises CosmosAccessConditionFailedError, proving another process changed the record. Re-read, reconsider the update, and retry according to application policy rather than silently overwriting newer data.
from azure.cosmos import exceptions
item = chunks.read_item(item="policy-42-chunk-3", partition_key="contoso")
item["reviewed"] = True
try:
chunks.replace_item(
item=item["id"],
body=item,
if_match=item["_etag"]
)
except exceptions.CosmosAccessConditionFailedError:
print("The item changed; read it again before retrying.")
Topic summary
Choose create for uniqueness, upsert for insert-or-replace, and replace with _etag when concurrent writers must not lose updates.
9. Prefer point reads and capture response metadata
read_item() is the lowest-latency and lowest-RU way to retrieve one known document because id and partition key route directly to its logical partition. Design identifiers so frequent user profiles, cached inferences, and configuration records can use this path. Handle CosmosResourceNotFoundError to compute a cache miss, return a default, or report absence.
delete_item() also requires id and partition key. The deleted item stops consuming storage, although deletion itself consumes RUs. After SDK operations, capture x-ms-request-charge and x-ms-activity-id from response metadata. Activity IDs are especially useful when correlating failures with Microsoft support.
Topic summary
When id and partition key are known, use read_item or delete_item directly and log RU charge plus activity ID for cost and diagnostics.
10. Build SELECT and WHERE queries for JSON documents
The NoSQL query language resembles SQL but operates inside one container and traverses JSON properties. FROM introduces the item alias, SELECT chooses the result shape, and WHERE filters it. Project only the fields an AI caller needs to reduce response size and downstream processing.
Use =, !=, <, >, <=, and >= for comparison.
Combine predicates with AND, OR, and NOT.
Use CONTAINS, STARTSWITH, ENDSWITH, UPPER, and LOWER for text patterns.
Use BETWEEN for ranges and IN or NOT IN for fixed sets.
Remember that string comparison is case-sensitive by default.
The query iterator can issue several requests as results are consumed. A broad SELECT * is convenient for inspection but is rarely the most efficient contract for an inference or retrieval endpoint.
Topic summary
Use the SQL-like language to filter JSON, but project a deliberate response and account for the iterator’s potentially paged execution.
11. Parameterize values and route to one partition
Never concatenate user or external values into query text. @parameters keep structure separate from data, prevent injection, and allow query-plan reuse. Provide the partition_key option when the target is known; an equality predicate on that key also allows efficient routing.
query = """
SELECT c.id, c.sourceId, c.text
FROM c
WHERE c.tenantId = @tenant AND c.sourceId = @source
ORDER BY c.position
"""
parameters = [
{"name": "@tenant", "value": "contoso"},
{"name": "@source", "value": "policy-42"}
]
results = chunks.query_items(
query=query,
parameters=parameters,
partition_key="contoso",
max_item_count=25
)
A cross-partition query fans out when the partition cannot be inferred. It is valid for global searches, but enable_cross_partition_query=True explicitly and watch its latency and RU charge. Increasing provisioned throughput can reduce throttling, but it does not repair a query that ignores an available partition key.
Topic summary
Parameterize all external values and target a single partition whenever the access pattern supplies its key.
12. Sort, paginate, shape, aggregate, and inspect query cost
ORDER BY sorts ascending or descending and can require an index that matches the query shape. For large result sets, set max_item_count and iterate pages; a web API can return the SDK continuation token as an opaque bookmark. Larger pages reduce round trips, while smaller pages reduce memory pressure.
Projections can rename properties, compute expressions, create nested JSON, or use VALUE to unwrap scalars and arrays. COUNT, SUM, AVG, MIN, and MAX calculate summaries but may scan many matches. ARRAY_CONTAINS tests membership; JOIN ... IN flattens array elements for filtering.
SELECT VALUE {
"chunkId": c.id,
"content": c.text,
"hasEmbedding": IS_DEFINED(c.embedding)
}
FROM c
WHERE c.tenantId = @tenant
SELECT COUNT(1) AS totalChunks, MAX(c.position) AS lastPosition
FROM c
WHERE c.tenantId = @tenant AND c.sourceId = @source
SELECT c.id, tag
FROM c
JOIN tag IN c.tags
WHERE tag IN ("security", "governance")
Filter as early and narrowly as possible.
Return only the fields the caller consumes.
Use the partition key and TOP when the request permits.
Align indexing, including composite indexes, with measured query patterns.
Read x-ms-request-charge per page and examine query and index metrics for expensive operations.
Topic summary
Control result size, routing, page size, projection, aggregation, array traversal, and indexing while measuring the RU cost of every important query.
13. Guided lab: build a Cosmos DB RAG document store
The approximately 30-minute exercise provisions an for NoSQL account, database, and container for chunked organizational documents. Each chunk carries metadata that a retrieval layer can query before passing grounded context to a language model. Python functions store and retrieve chunks, and a Flask application validates the workflow through the Cosmos DB SQL API.
Download the starter project and review the deployment configuration.
Deploy the account, database, container, and partition strategy.
Implement Python functions that write and retrieve document chunks.
Query the relevant context and pass it to the RAG application.
Run the Flask interface, verify results, and remove billable Azure resources.
Prerequisites are an Azure subscription with deployment permission, , the latest Azure CLI, and Python 3.12 or later.
The database is the retrieval layer in the RAG flow; relevance and partition design determine how efficiently context is assembled.
Topic summary
The lab joins schema design, deployment, Python SDK operations, SQL retrieval, and a Flask client into a working RAG document-store pattern.
14. Assessment decisions and official references
Assessment review.
Scenario
Best answer
Why
Retrieve all interaction logs for one user
Use userId as the partition key.
It aligns distribution and routing with the main access pattern.
Cache may or may not already exist
Use upsert_item().
It inserts or replaces without a preliminary existence check.
Known recommendation id and category
Use read_item() with id and partition key.
A point read is more efficient than a query.
Filters come from a user
Use parameterized queries.
They prevent injection and support plan caching.
Price query is costly and categoryId is the key
Include categoryId in WHERE or specify partition_key.
Exam questions reward access-pattern alignment: partition by the dominant scope, use the narrowest SDK operation, parameterize values, and avoid unnecessary fan-out.