From fa0b600755668220d366f84e18f58406801e95ce Mon Sep 17 00:00:00 2001 From: Prashant Sasatte Date: Thu, 11 Jun 2026 12:31:09 +0530 Subject: [PATCH] Add Cosmos DB user agent suffixes Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- api/sitecustomize.py | 34 +++++++++++++++++++ connectors/ingestion/dotnet/Program.cs | 4 +++ .../dotnet/Services/SourceWatcher.cs | 3 ++ .../Destinations/CosmosDbDestinationWriter.cs | 3 ++ 4 files changed, 44 insertions(+) create mode 100644 api/sitecustomize.py diff --git a/api/sitecustomize.py b/api/sitecustomize.py new file mode 100644 index 0000000..f52d96a --- /dev/null +++ b/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 diff --git a/connectors/ingestion/dotnet/Program.cs b/connectors/ingestion/dotnet/Program.cs index 56ff232..e05fe58 100644 --- a/connectors/ingestion/dotnet/Program.cs +++ b/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), diff --git a/connectors/ingestion/dotnet/Services/SourceWatcher.cs b/connectors/ingestion/dotnet/Services/SourceWatcher.cs index 2610c34..084af8c 100644 --- a/connectors/ingestion/dotnet/Services/SourceWatcher.cs +++ b/connectors/ingestion/dotnet/Services/SourceWatcher.cs @@ -16,6 +16,8 @@ namespace OmniVec.ChangeFeed.Services; /// 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 @@ public async Task StartAsync(CancellationToken ct) new DefaultAzureCredential(), new CosmosClientOptions { + ApplicationName = CosmosDataUserAgent, ConnectionMode = ConnectionMode.Direct, ConsistencyLevel = ConsistencyLevel.Eventual, MaxRetryAttemptsOnRateLimitedRequests = int.MaxValue, diff --git a/connectors/worker/dotnet/Destinations/CosmosDbDestinationWriter.cs b/connectors/worker/dotnet/Destinations/CosmosDbDestinationWriter.cs index 0d31a21..58b6097 100644 --- a/connectors/worker/dotnet/Destinations/CosmosDbDestinationWriter.cs +++ b/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 _logger; private static readonly ConcurrentDictionary _clients = new(); private static readonly ConcurrentDictionary _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),