From 40676cda4ef8252fdedc33afabc817c7bb2c54b0 Mon Sep 17 00:00:00 2001 From: Stas Moreinis Date: Fri, 19 Dec 2025 21:48:51 -0800 Subject: [PATCH] Add OpenTelemetry instrumentation to SQLAlchemy engines Instrument both main and middleware async database engines with OpenTelemetry to trace SQL queries. The instrumentation operates at the SQLAlchemy level (not asyncpg driver level) for fine-grained control without conflicting with ddtrace's existing instrumentation. --- agentex/pyproject.toml | 3 + agentex/src/config/dependencies.py | 10 ++++ agentex/src/config/otel_instrumentation.py | 57 +++++++++++++++++++ uv.lock | 66 ++++++++++++++++++++++ 4 files changed, 136 insertions(+) create mode 100644 agentex/src/config/otel_instrumentation.py diff --git a/agentex/pyproject.toml b/agentex/pyproject.toml index 2cf7b83..7cb1f83 100644 --- a/agentex/pyproject.toml +++ b/agentex/pyproject.toml @@ -27,6 +27,9 @@ dependencies = [ "ddtrace>=3.13.0", "json_log_formatter>=1.1.1", "datadog>=0.52.1", + "opentelemetry-api>=1.27.0", + "opentelemetry-sdk>=1.27.0", + "opentelemetry-instrumentation-sqlalchemy>=0.48b0", ] [dependency-groups] diff --git a/agentex/src/config/dependencies.py b/agentex/src/config/dependencies.py index 9ddf92a..369d357 100644 --- a/agentex/src/config/dependencies.py +++ b/agentex/src/config/dependencies.py @@ -18,6 +18,7 @@ from temporalio.client import Client as TemporalClient from src.config.environment_variables import Environment, EnvironmentVariables +from src.config.otel_instrumentation import instrument_engine, uninstrument_all from src.utils.database import async_db_engine_creator from src.utils.logging import make_logger @@ -113,6 +114,12 @@ async def load(self): pool_recycle=3600, # Recycle connections after 1 hour ) + # Instrument SQLAlchemy engines with OpenTelemetry + instrument_engine(self.database_async_read_write_engine, "main") + instrument_engine( + self.database_async_middleware_read_write_engine, "middleware" + ) + # Initialize MongoDB client and database try: mongodb_uri = self.environment_variables.MONGODB_URI @@ -222,6 +229,9 @@ def shutdown(): async def async_shutdown(): + # Uninstrument SQLAlchemy engines before disposal + uninstrument_all() + global_dependencies = GlobalDependencies() run_concurrently = [] # if global_dependencies.database_async_read_only_engine: diff --git a/agentex/src/config/otel_instrumentation.py b/agentex/src/config/otel_instrumentation.py new file mode 100644 index 0000000..94b2a20 --- /dev/null +++ b/agentex/src/config/otel_instrumentation.py @@ -0,0 +1,57 @@ +""" +OpenTelemetry instrumentation for SQLAlchemy engines. +""" + +from opentelemetry.instrumentation.sqlalchemy import SQLAlchemyInstrumentor +from sqlalchemy.ext.asyncio import AsyncEngine + +from src.utils.logging import make_logger + +logger = make_logger(__name__) + +_instrumented_engines: set[int] = set() + + +def instrument_engine( + engine: AsyncEngine | None, + engine_name: str = "default", +) -> None: + """ + Instrument a SQLAlchemy async engine with OpenTelemetry. + + For async engines, we instrument the underlying sync_engine since + SQLAlchemyInstrumentor works at the sync engine level. + + Args: + engine: The async SQLAlchemy engine to instrument + engine_name: A name identifier for logging + """ + if engine is None: + logger.warning(f"Cannot instrument {engine_name}: engine is None") + return + + engine_id = id(engine) + if engine_id in _instrumented_engines: + logger.debug(f"Engine {engine_name} already instrumented") + return + + try: + SQLAlchemyInstrumentor().instrument( + engine=engine.sync_engine, + ) + _instrumented_engines.add(engine_id) + logger.info( + f"Instrumented SQLAlchemy engine '{engine_name}' with OpenTelemetry" + ) + except Exception as e: + logger.error(f"Failed to instrument SQLAlchemy engine '{engine_name}': {e}") + + +def uninstrument_all() -> None: + """Remove instrumentation from all engines.""" + try: + SQLAlchemyInstrumentor().uninstrument() + _instrumented_engines.clear() + logger.info("Uninstrumented all SQLAlchemy engines") + except Exception as e: + logger.error(f"Failed to uninstrument SQLAlchemy engines: {e}") diff --git a/uv.lock b/uv.lock index 708b274..0f2d037 100644 --- a/uv.lock +++ b/uv.lock @@ -58,6 +58,9 @@ dependencies = [ { name = "json-log-formatter" }, { name = "kubernetes-asyncio" }, { name = "litellm" }, + { name = "opentelemetry-api" }, + { name = "opentelemetry-instrumentation-sqlalchemy" }, + { name = "opentelemetry-sdk" }, { name = "psycopg2-binary" }, { name = "pymongo" }, { name = "python-dotenv" }, @@ -113,6 +116,9 @@ requires-dist = [ { name = "json-log-formatter", specifier = ">=1.1.1" }, { name = "kubernetes-asyncio", specifier = ">=31.1.0,<32" }, { name = "litellm", specifier = ">=1.48.2,<2" }, + { name = "opentelemetry-api", specifier = ">=1.27.0" }, + { name = "opentelemetry-instrumentation-sqlalchemy", specifier = ">=0.48b0" }, + { name = "opentelemetry-sdk", specifier = ">=1.27.0" }, { name = "psycopg2-binary", specifier = ">=2.9.9,<3" }, { name = "pymongo", specifier = ">=4.11.2,<5" }, { name = "python-dotenv", specifier = ">=1.0.1,<2" }, @@ -767,6 +773,8 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/19/0d/6660d55f7373b2ff8152401a83e02084956da23ae58cddbfb0b330978fe9/greenlet-3.2.4-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:3b3812d8d0c9579967815af437d96623f45c0f2ae5f04e366de62a12d83a8fb0", size = 607586, upload-time = "2025-08-07T13:18:28.544Z" }, { url = "https://files.pythonhosted.org/packages/8e/1a/c953fdedd22d81ee4629afbb38d2f9d71e37d23caace44775a3a969147d4/greenlet-3.2.4-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:abbf57b5a870d30c4675928c37278493044d7c14378350b3aa5d484fa65575f0", size = 1123281, upload-time = "2025-08-07T13:42:39.858Z" }, { url = "https://files.pythonhosted.org/packages/3f/c7/12381b18e21aef2c6bd3a636da1088b888b97b7a0362fac2e4de92405f97/greenlet-3.2.4-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:20fb936b4652b6e307b8f347665e2c615540d4b42b3b4c8a321d8286da7e520f", size = 1151142, upload-time = "2025-08-07T13:18:22.981Z" }, + { url = "https://files.pythonhosted.org/packages/27/45/80935968b53cfd3f33cf99ea5f08227f2646e044568c9b1555b58ffd61c2/greenlet-3.2.4-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:ee7a6ec486883397d70eec05059353b8e83eca9168b9f3f9a361971e77e0bcd0", size = 1564846, upload-time = "2025-11-04T12:42:15.191Z" }, + { url = "https://files.pythonhosted.org/packages/69/02/b7c30e5e04752cb4db6202a3858b149c0710e5453b71a3b2aec5d78a1aab/greenlet-3.2.4-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:326d234cbf337c9c3def0676412eb7040a35a768efc92504b947b3e9cfc7543d", size = 1633814, upload-time = "2025-11-04T12:42:17.175Z" }, { url = "https://files.pythonhosted.org/packages/e9/08/b0814846b79399e585f974bbeebf5580fbe59e258ea7be64d9dfb253c84f/greenlet-3.2.4-cp312-cp312-win_amd64.whl", hash = "sha256:a7d4e128405eea3814a12cc2605e0e6aedb4035bf32697f72deca74de4105e02", size = 299899, upload-time = "2025-08-07T13:38:53.448Z" }, ] @@ -1569,6 +1577,64 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/91/48/28ed9e55dcf2f453128df738210a980e09f4e468a456fa3c763dbc8be70a/opentelemetry_api-1.37.0-py3-none-any.whl", hash = "sha256:accf2024d3e89faec14302213bc39550ec0f4095d1cf5ca688e1bfb1c8612f47", size = 65732, upload-time = "2025-09-11T10:28:41.826Z" }, ] +[[package]] +name = "opentelemetry-instrumentation" +version = "0.58b0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "opentelemetry-api" }, + { name = "opentelemetry-semantic-conventions" }, + { name = "packaging" }, + { name = "wrapt" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/f6/36/7c307d9be8ce4ee7beb86d7f1d31027f2a6a89228240405a858d6e4d64f9/opentelemetry_instrumentation-0.58b0.tar.gz", hash = "sha256:df640f3ac715a3e05af145c18f527f4422c6ab6c467e40bd24d2ad75a00cb705", size = 31549, upload-time = "2025-09-11T11:42:14.084Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d4/db/5ff1cd6c5ca1d12ecf1b73be16fbb2a8af2114ee46d4b0e6d4b23f4f4db7/opentelemetry_instrumentation-0.58b0-py3-none-any.whl", hash = "sha256:50f97ac03100676c9f7fc28197f8240c7290ca1baa12da8bfbb9a1de4f34cc45", size = 33019, upload-time = "2025-09-11T11:41:00.624Z" }, +] + +[[package]] +name = "opentelemetry-instrumentation-sqlalchemy" +version = "0.58b0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "opentelemetry-api" }, + { name = "opentelemetry-instrumentation" }, + { name = "opentelemetry-semantic-conventions" }, + { name = "packaging" }, + { name = "wrapt" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/91/6f/fa2c45d5dfb8da0bcc337a421fd12946994cb5e9782c40a21c669e78460d/opentelemetry_instrumentation_sqlalchemy-0.58b0.tar.gz", hash = "sha256:3e4b444a05088ba473710df9d5c730bb08969c8ea71e04f2886a0f7efee22c12", size = 14993, upload-time = "2025-09-11T11:42:52.053Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/42/26/7a86285adf7131801fe297bdb756b00bc3f43f78bc7e18668d88f7308e53/opentelemetry_instrumentation_sqlalchemy-0.58b0-py3-none-any.whl", hash = "sha256:7c11d11887b3a4dbc3fccdd58a24903e306052ff7a59685caea1dd7797f82b0a", size = 14212, upload-time = "2025-09-11T11:41:53.76Z" }, +] + +[[package]] +name = "opentelemetry-sdk" +version = "1.37.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "opentelemetry-api" }, + { name = "opentelemetry-semantic-conventions" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/f4/62/2e0ca80d7fe94f0b193135375da92c640d15fe81f636658d2acf373086bc/opentelemetry_sdk-1.37.0.tar.gz", hash = "sha256:cc8e089c10953ded765b5ab5669b198bbe0af1b3f89f1007d19acd32dc46dda5", size = 170404, upload-time = "2025-09-11T10:29:11.779Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/9f/62/9f4ad6a54126fb00f7ed4bb5034964c6e4f00fcd5a905e115bd22707e20d/opentelemetry_sdk-1.37.0-py3-none-any.whl", hash = "sha256:8f3c3c22063e52475c5dbced7209495c2c16723d016d39287dfc215d1771257c", size = 131941, upload-time = "2025-09-11T10:28:57.83Z" }, +] + +[[package]] +name = "opentelemetry-semantic-conventions" +version = "0.58b0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "opentelemetry-api" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/aa/1b/90701d91e6300d9f2fb352153fb1721ed99ed1f6ea14fa992c756016e63a/opentelemetry_semantic_conventions-0.58b0.tar.gz", hash = "sha256:6bd46f51264279c433755767bb44ad00f1c9e2367e1b42af563372c5a6fa0c25", size = 129867, upload-time = "2025-09-11T10:29:12.597Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/07/90/68152b7465f50285d3ce2481b3aec2f82822e3f52e5152eeeaf516bab841/opentelemetry_semantic_conventions-0.58b0-py3-none-any.whl", hash = "sha256:5564905ab1458b96684db1340232729fce3b5375a06e140e8904c78e4f815b28", size = 207954, upload-time = "2025-09-11T10:28:59.218Z" }, +] + [[package]] name = "packaging" version = "25.0"