Add Cosmos DB user agent suffixes

api/sitecustomize.py

@@ -0,0 +1,34 @@
+"""Process-wide Cosmos DB user-agent suffixes for OmniVec API services."""
+
+from __future__ import annotations
+
+import inspect
+import os
+
+import azure.cosmos
+import azure.cosmos.cosmos_client as _cosmos_client
+from azure.cosmos import CosmosClient as _CosmosClient
+
+
+COSMOS_METADATA_USER_AGENT = "OmniVec-MetadataCosmos/1.0"
+COSMOS_DATA_USER_AGENT = "OmniVec-DataCosmos/1.0"
+
+
+def _cosmos_user_agent_suffix() -> str:
+    for frame in inspect.stack(context=0)[2:]:
+        filename = os.path.basename(frame.filename)
+        if filename == "store.py":
+            return COSMOS_METADATA_USER_AGENT
+        if filename == "api.py" and frame.function == "get_changefeed_leases":
+            return COSMOS_METADATA_USER_AGENT
+    return COSMOS_DATA_USER_AGENT
+
+
+class OmniVecCosmosClient(_CosmosClient):
+    def __init__(self, *args, **kwargs):
+        kwargs.setdefault("user_agent_suffix", _cosmos_user_agent_suffix())
+        super().__init__(*args, **kwargs)
+
+
+azure.cosmos.CosmosClient = OmniVecCosmosClient
+_cosmos_client.CosmosClient = OmniVecCosmosClient

connectors/ingestion/dotnet/Program.cs

@@ -9,6 +9,8 @@
 // 500 threads for 100K RU/s provisioned throughput
 ThreadPool.SetMinThreads(1000, 1000);
 
+const string CosmosMetadataUserAgent = "OmniVec-MetadataCosmos/1.0";
+
 var builder = Host.CreateApplicationBuilder(args);
 
 // Bind configuration from env vars: ChangeFeed__OmniVecCosmosEndpoint, etc.
@@ -23,6 +25,7 @@
     return new CosmosClient(opts.OmniVecCosmosEndpoint, new DefaultAzureCredential(),
         new CosmosClientOptions
         {
+            ApplicationName = CosmosMetadataUserAgent,
             ConnectionMode = ConnectionMode.Direct,
             MaxRetryAttemptsOnRateLimitedRequests = int.MaxValue,
             MaxRetryWaitTimeOnRateLimitedRequests = TimeSpan.FromSeconds(300),
@@ -40,6 +43,7 @@
     return new CosmosClient(leaseEndpoint, new DefaultAzureCredential(),
         new CosmosClientOptions
         {
+            ApplicationName = CosmosMetadataUserAgent,
             ConnectionMode = ConnectionMode.Direct,
             MaxRetryAttemptsOnRateLimitedRequests = int.MaxValue,
             MaxRetryWaitTimeOnRateLimitedRequests = TimeSpan.FromSeconds(300),

connectors/ingestion/dotnet/Services/SourceWatcher.cs

@@ -16,6 +16,8 @@
 /// </summary>
 public class SourceWatcher : ISourceWatcher
 {
+    private const string CosmosDataUserAgent = "OmniVec-DataCosmos/1.0";
+
     private readonly Source _source;
     private readonly ChangeFeedOptions _options;
     private readonly OmniVecApiClient _apiClient;
@@ -106,6 +108,7 @@
             new DefaultAzureCredential(),
             new CosmosClientOptions
             {
+                ApplicationName = CosmosDataUserAgent,
                 ConnectionMode = ConnectionMode.Direct,
                 ConsistencyLevel = ConsistencyLevel.Eventual,
                 MaxRetryAttemptsOnRateLimitedRequests = int.MaxValue,
@@ -385,7 +388,7 @@
                     if (dest is null)
                     {
                         _logger.LogWarning(
                             "Skipping {Count} docs for pipeline {PipelineId}: destination {DestId} not found or disabled. " +
                             "Enable the destination (e.g. `omnivec destination enable {DestId}`) and the change-feed will replay.",
                             pipelineDocs.Count, pipeline.Id, pipeline.DestinationId);
                         continue;

connectors/worker/dotnet/Destinations/CosmosDbDestinationWriter.cs

@@ -6,6 +6,8 @@ namespace OmniVec.Worker.Destinations;
 
 public class CosmosDbDestinationWriter : IDestinationWriter
 {
+    private const string CosmosDataUserAgent = "OmniVec-DataCosmos/1.0";
+
     private readonly ILogger<CosmosDbDestinationWriter> _logger;
     private static readonly ConcurrentDictionary<string, CosmosClient> _clients = new();
     private static readonly ConcurrentDictionary<string, string> _pkPathCache = new();
@@ -241,6 +243,7 @@ private static CosmosClient GetOrCreateClient(string endpoint)
         return _clients.GetOrAdd(endpoint, ep =>
             new CosmosClient(ep, new DefaultAzureCredential(), new CosmosClientOptions
             {
+                ApplicationName = CosmosDataUserAgent,
                 ConnectionMode = ConnectionMode.Direct,
                 MaxRetryAttemptsOnRateLimitedRequests = int.MaxValue,
                 MaxRetryWaitTimeOnRateLimitedRequests = TimeSpan.FromSeconds(300),