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 9a21678..111d16e 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 @@ -231,6 +238,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"