Skip to content

Commit

Permalink
add possibility to instrument mutiple engines
Browse files Browse the repository at this point in the history
  • Loading branch information
hemma committed Jun 14, 2022
1 parent 81fcad0 commit 75fb63b
Show file tree
Hide file tree
Showing 2 changed files with 9 additions and 9 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -73,9 +73,6 @@ class SQLAlchemyInstrumentor(BaseInstrumentor):
See `BaseInstrumentor`
"""

def __new__(cls, *args, **kwargs):
return object.__new__(cls, *args, **kwargs)

def instrumentation_dependencies(self) -> Collection[str]:
return _instruments

Expand All @@ -86,6 +83,7 @@ def _instrument(self, **kwargs):
Args:
**kwargs: Optional arguments
``engine``: a SQLAlchemy engine instance
``engines``: a list of SQLAlchemy engine instances
``tracer_provider``: a TracerProvider, defaults to global
Returns:
Expand All @@ -111,6 +109,12 @@ def _instrument(self, **kwargs):
kwargs.get("engine"),
kwargs.get("enable_commenter", False),
)
if kwargs.get("engines") is not None and isinstance(kwargs.get("engines"), list):
return [EngineTracer(
_get_tracer(tracer_provider),
engine,
kwargs.get("enable_commenter", False),) for engine in kwargs.get("engines")]

return None

def _uninstrument(self, **kwargs):
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -55,16 +55,12 @@ def test_instrument_two_engines(self):
engine_2 = create_engine("sqlite:///:memory:")

SQLAlchemyInstrumentor().instrument(
engine=engine_1,
engines=[engine_1, engine_2],
tracer_provider=self.tracer_provider,
)

cnx_1 = engine_1.connect()
cnx_1.execute("SELECT 1 + 1;").fetchall()

SQLAlchemyInstrumentor().instrument(
engine=engine_2,
tracer_provider=self.tracer_provider,
)
cnx_2 = engine_2.connect()
cnx_2.execute("SELECT 1 + 1;").fetchall()

Expand Down

0 comments on commit 75fb63b

Please sign in to comment.