Skip to content

Commit 79d62b3

Browse files
tammy-baylis-swishalevr
tammy-baylis-swi
and
shalevr
authored
sqlalchemy wrap_create_engine now accepts sqlcommenter options (#1873)
* sqlalchemy wrap_create_engine accepts sqlcommenter options * Changelog * Lint * Fix default val * Add sqlalchemy tests * Change a default in _instrument get * Lint * More lint * Update default Co-authored-by: Shalev Roda <[email protected]> * Update args doc * lintttt --------- Co-authored-by: Shalev Roda <[email protected]>
1 parent 2e49ba1 commit 79d62b3

File tree

5 files changed

+146
-7
lines changed

5 files changed

+146
-7
lines changed

‎CHANGELOG.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
1515
([#1870](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1870))
1616
- Update falcon instrumentation to follow semantic conventions
1717
([#1824](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1824))
18+
- Fix sqlalchemy instrumentation wrap methods to accept sqlcommenter options([#1873](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/1873))
1819

1920
### Added
2021

‎instrumentation/opentelemetry-instrumentation-sqlalchemy/src/opentelemetry/instrumentation/sqlalchemy/__init__.py

Lines changed: 14 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -134,6 +134,9 @@ def _instrument(self, **kwargs):
134134
``engine``: a SQLAlchemy engine instance
135135
``engines``: a list of SQLAlchemy engine instances
136136
``tracer_provider``: a TracerProvider, defaults to global
137+
``meter_provider``: a MeterProvider, defaults to global
138+
``enable_commenter``: bool to enable sqlcommenter, defaults to False
139+
``commenter_options``: dict of sqlcommenter config, defaults to None
137140
138141
Returns:
139142
An instrumented engine if passed in as an argument or list of instrumented engines, None otherwise.
@@ -151,16 +154,21 @@ def _instrument(self, **kwargs):
151154
)
152155

153156
enable_commenter = kwargs.get("enable_commenter", False)
157+
commenter_options = kwargs.get("commenter_options", {})
154158

155159
_w(
156160
"sqlalchemy",
157161
"create_engine",
158-
_wrap_create_engine(tracer, connections_usage, enable_commenter),
162+
_wrap_create_engine(
163+
tracer, connections_usage, enable_commenter, commenter_options
164+
),
159165
)
160166
_w(
161167
"sqlalchemy.engine",
162168
"create_engine",
163-
_wrap_create_engine(tracer, connections_usage, enable_commenter),
169+
_wrap_create_engine(
170+
tracer, connections_usage, enable_commenter, commenter_options
171+
),
164172
)
165173
_w(
166174
"sqlalchemy.engine.base",
@@ -172,7 +180,10 @@ def _instrument(self, **kwargs):
172180
"sqlalchemy.ext.asyncio",
173181
"create_async_engine",
174182
_wrap_create_async_engine(
175-
tracer, connections_usage, enable_commenter
183+
tracer,
184+
connections_usage,
185+
enable_commenter,
186+
commenter_options,
176187
),
177188
)
178189
if kwargs.get("engine") is not None:

‎instrumentation/opentelemetry-instrumentation-sqlalchemy/src/opentelemetry/instrumentation/sqlalchemy/engine.py

Lines changed: 16 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,7 @@ def _normalize_vendor(vendor):
4343

4444

4545
def _wrap_create_async_engine(
46-
tracer, connections_usage, enable_commenter=False
46+
tracer, connections_usage, enable_commenter=False, commenter_options=None
4747
):
4848
# pylint: disable=unused-argument
4949
def _wrap_create_async_engine_internal(func, module, args, kwargs):
@@ -52,20 +52,32 @@ def _wrap_create_async_engine_internal(func, module, args, kwargs):
5252
"""
5353
engine = func(*args, **kwargs)
5454
EngineTracer(
55-
tracer, engine.sync_engine, connections_usage, enable_commenter
55+
tracer,
56+
engine.sync_engine,
57+
connections_usage,
58+
enable_commenter,
59+
commenter_options,
5660
)
5761
return engine
5862

5963
return _wrap_create_async_engine_internal
6064

6165

62-
def _wrap_create_engine(tracer, connections_usage, enable_commenter=False):
66+
def _wrap_create_engine(
67+
tracer, connections_usage, enable_commenter=False, commenter_options=None
68+
):
6369
def _wrap_create_engine_internal(func, _module, args, kwargs):
6470
"""Trace the SQLAlchemy engine, creating an `EngineTracer`
6571
object that will listen to SQLAlchemy events.
6672
"""
6773
engine = func(*args, **kwargs)
68-
EngineTracer(tracer, engine, connections_usage, enable_commenter)
74+
EngineTracer(
75+
tracer,
76+
engine,
77+
connections_usage,
78+
enable_commenter,
79+
commenter_options,
80+
)
6981
return engine
7082

7183
return _wrap_create_engine_internal

‎instrumentation/opentelemetry-instrumentation-sqlalchemy/tests/test_sqlalchemy.py

Lines changed: 97 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
1414
import asyncio
15+
import logging
1516
from unittest import mock
1617

1718
import pytest
@@ -176,6 +177,43 @@ def test_create_engine_wrapper(self):
176177
"opentelemetry.instrumentation.sqlalchemy",
177178
)
178179

180+
def test_create_engine_wrapper_enable_commenter(self):
181+
logging.getLogger("sqlalchemy.engine").setLevel(logging.INFO)
182+
SQLAlchemyInstrumentor().instrument(
183+
enable_commenter=True,
184+
commenter_options={"db_framework": False},
185+
)
186+
from sqlalchemy import create_engine # pylint: disable-all
187+
188+
engine = create_engine("sqlite:///:memory:")
189+
cnx = engine.connect()
190+
cnx.execute("SELECT 1;").fetchall()
191+
# sqlcommenter
192+
self.assertRegex(
193+
self.caplog.records[-2].getMessage(),
194+
r"SELECT 1 /\*db_driver='(.*)',traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;",
195+
)
196+
197+
def test_create_engine_wrapper_enable_commenter_otel_values_false(self):
198+
logging.getLogger("sqlalchemy.engine").setLevel(logging.INFO)
199+
SQLAlchemyInstrumentor().instrument(
200+
enable_commenter=True,
201+
commenter_options={
202+
"db_framework": False,
203+
"opentelemetry_values": False,
204+
},
205+
)
206+
from sqlalchemy import create_engine # pylint: disable-all
207+
208+
engine = create_engine("sqlite:///:memory:")
209+
cnx = engine.connect()
210+
cnx.execute("SELECT 1;").fetchall()
211+
# sqlcommenter
212+
self.assertRegex(
213+
self.caplog.records[-2].getMessage(),
214+
r"SELECT 1 /\*db_driver='(.*)'\*/;",
215+
)
216+
179217
def test_custom_tracer_provider(self):
180218
provider = TracerProvider(
181219
resource=Resource.create(
@@ -242,6 +280,65 @@ async def run():
242280

243281
asyncio.get_event_loop().run_until_complete(run())
244282

283+
@pytest.mark.skipif(
284+
not sqlalchemy.__version__.startswith("1.4"),
285+
reason="only run async tests for 1.4",
286+
)
287+
def test_create_async_engine_wrapper_enable_commenter(self):
288+
async def run():
289+
logging.getLogger("sqlalchemy.engine").setLevel(logging.INFO)
290+
SQLAlchemyInstrumentor().instrument(
291+
enable_commenter=True,
292+
commenter_options={
293+
"db_framework": False,
294+
},
295+
)
296+
from sqlalchemy.ext.asyncio import ( # pylint: disable-all
297+
create_async_engine,
298+
)
299+
300+
engine = create_async_engine("sqlite+aiosqlite:///:memory:")
301+
async with engine.connect() as cnx:
302+
await cnx.execute(sqlalchemy.text("SELECT 1;"))
303+
# sqlcommenter
304+
self.assertRegex(
305+
self.caplog.records[1].getMessage(),
306+
r"SELECT 1 /\*db_driver='(.*)',traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;",
307+
)
308+
309+
asyncio.get_event_loop().run_until_complete(run())
310+
311+
@pytest.mark.skipif(
312+
not sqlalchemy.__version__.startswith("1.4"),
313+
reason="only run async tests for 1.4",
314+
)
315+
def test_create_async_engine_wrapper_enable_commenter_otel_values_false(
316+
self,
317+
):
318+
async def run():
319+
logging.getLogger("sqlalchemy.engine").setLevel(logging.INFO)
320+
SQLAlchemyInstrumentor().instrument(
321+
enable_commenter=True,
322+
commenter_options={
323+
"db_framework": False,
324+
"opentelemetry_values": False,
325+
},
326+
)
327+
from sqlalchemy.ext.asyncio import ( # pylint: disable-all
328+
create_async_engine,
329+
)
330+
331+
engine = create_async_engine("sqlite+aiosqlite:///:memory:")
332+
async with engine.connect() as cnx:
333+
await cnx.execute(sqlalchemy.text("SELECT 1;"))
334+
# sqlcommenter
335+
self.assertRegex(
336+
self.caplog.records[1].getMessage(),
337+
r"SELECT 1 /\*db_driver='(.*)'\*/;",
338+
)
339+
340+
asyncio.get_event_loop().run_until_complete(run())
341+
245342
def test_uninstrument(self):
246343
engine = create_engine("sqlite:///:memory:")
247344
SQLAlchemyInstrumentor().instrument(

‎instrumentation/opentelemetry-instrumentation-sqlalchemy/tests/test_sqlcommenter.py

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,24 @@ def test_sqlcommenter_enabled(self):
5656
r"SELECT 1 /\*db_driver='(.*)',traceparent='\d{1,2}-[a-zA-Z0-9_]{32}-[a-zA-Z0-9_]{16}-\d{1,2}'\*/;",
5757
)
5858

59+
def test_sqlcommenter_enabled_otel_values_false(self):
60+
engine = create_engine("sqlite:///:memory:")
61+
SQLAlchemyInstrumentor().instrument(
62+
engine=engine,
63+
tracer_provider=self.tracer_provider,
64+
enable_commenter=True,
65+
commenter_options={
66+
"db_framework": False,
67+
"opentelemetry_values": False,
68+
},
69+
)
70+
cnx = engine.connect()
71+
cnx.execute("SELECT 1;").fetchall()
72+
self.assertRegex(
73+
self.caplog.records[-2].getMessage(),
74+
r"SELECT 1 /\*db_driver='(.*)'\*/;",
75+
)
76+
5977
def test_sqlcommenter_flask_integration(self):
6078
engine = create_engine("sqlite:///:memory:")
6179
SQLAlchemyInstrumentor().instrument(

0 commit comments

Comments
 (0)