RD-3763 - sqlalchemy (#159)

* support sqlalchemy

Author: saartochner-lumigo

requirements.txt

@@ -8,3 +8,4 @@ attrs==19.1.0
 requests==2.24.0
 pymongo==3.11.0
 redispy==3.0.0
+sqlalchemy==1.3.20

src/lumigo_tracer/wrappers/__init__.py

@@ -1,6 +1,7 @@
 from .http.sync_http_wrappers import wrap_http_calls
 from .pymongo.pymongo_wrapper import wrap_pymongo
 from .redis.redis_wrapper import wrap_redis
+from .sql.sqlalchemy_wrapper import wrap_sqlalchemy
 
 
 already_wrapped = False
@@ -14,4 +15,5 @@ def wrap(force: bool = False):
     if force or not already_wrapped:
         wrap_pymongo()
         wrap_redis()
+        wrap_sqlalchemy()
         already_wrapped = True

src/lumigo_tracer/wrappers/sql/__init__.py

@@ -0,0 +1,84 @@
+import importlib
+import uuid
+
+from lumigo_tracer.libs.wrapt import wrap_function_wrapper
+from lumigo_tracer.lumigo_utils import (
+    lumigo_safe_execute,
+    get_logger,
+    lumigo_dumps,
+    get_current_ms_time,
+)
+from lumigo_tracer.spans_container import SpansContainer
+
+try:
+    from sqlalchemy.event import listen
+except Exception:
+    listen = None
+
+
+SQL_SPAN = "mySql"
+
+
+def _before_cursor_execute(conn, cursor, statement, parameters, context, executemany):
+    with lumigo_safe_execute("handle sqlalchemy before execute"):
+        SpansContainer.get_span().add_span(
+            {
+                "id": str(uuid.uuid4()),
+                "type": SQL_SPAN,
+                "started": get_current_ms_time(),
+                "connectionParameters": {
+                    "host": conn.engine.url.host or conn.engine.url.database,
+                    "port": conn.engine.url.port,
+                    "database": conn.engine.url.database,
+                    "user": conn.engine.url.username,
+                },
+                "query": lumigo_dumps(statement),
+                "values": lumigo_dumps(parameters),
+            }
+        )
+
+
+def _after_cursor_execute(conn, cursor, statement, parameters, context, executemany):
+    with lumigo_safe_execute("handle sqlalchemy after execute"):
+        span = SpansContainer.get_span().get_last_span()
+        if not span:
+            get_logger().warning("Redis span ended without a record on its start")
+            return
+        span.update({"ended": get_current_ms_time(), "response": ""})
+
+
+def _handle_error(context):
+    with lumigo_safe_execute("handle sqlalchemy error"):
+        span = SpansContainer.get_span().get_last_span()
+        if not span:
+            get_logger().warning("Redis span ended without a record on its start")
+            return
+        span.update(
+            {
+                "ended": get_current_ms_time(),
+                "error": lumigo_dumps(
+                    {
+                        "type": context.original_exception.__class__.__name__,
+                        "args": context.original_exception.args,
+                    }
+                ),
+            }
+        )
+
+
+def execute_wrapper(func, instance, args, kwargs):
+    result = func(*args, **kwargs)
+    with lumigo_safe_execute("sqlalchemy: listen to engine"):
+        listen(result, "before_cursor_execute", _before_cursor_execute)
+        listen(result, "after_cursor_execute", _after_cursor_execute)
+        listen(result, "handle_error", _handle_error)
+    return result
+
+
+def wrap_sqlalchemy():
+    with lumigo_safe_execute("wrap sqlalchemy"):
+        if importlib.util.find_spec("sqlalchemy") and listen:
+            get_logger().debug("wrapping sqlalchemy")
+            wrap_function_wrapper(
+                "sqlalchemy.engine.strategies", "DefaultEngineStrategy.create", execute_wrapper
+            )

src/lumigo_tracer/wrappers/sql/sqlalchemy_wrapper.py

@@ -0,0 +1,95 @@
+import uuid
+
+import pytest
+
+from sqlalchemy.exc import OperationalError
+from sqlalchemy import create_engine, Table, Column, Integer, String, MetaData
+from sqlalchemy.sql import select
+
+from lumigo_tracer.lumigo_utils import DEFAULT_MAX_ENTRY_SIZE
+from lumigo_tracer.spans_container import SpansContainer
+from lumigo_tracer.tracer import lumigo_tracer
+
+
+md = MetaData()
+Users = Table("users", md, Column("id", Integer, primary_key=True), Column("name", String))
+
+
+@pytest.fixture
+def db(tmp_path):
+    path = tmp_path / "file.db"
+    engine = create_engine(fr"sqlite:///{path}")
+    md.create_all(engine)
+    yield f"sqlite:///{path}"
+
+
+def test_happy_flow(context, db):
+    @lumigo_tracer(token="123")
+    def lambda_test_function(event, context):
+        engine = create_engine(db)
+        conn = engine.connect()
+        conn.execute(Users.insert().values(name="saart"))
+        result = conn.execute(select([Users]))
+        return result.fetchone()
+
+    assert lambda_test_function({}, context) == (1, "saart")
+    http_spans = SpansContainer.get_span().spans
+
+    assert len(http_spans) == 2
+    assert http_spans[0]["query"] == '"INSERT INTO users (name) VALUES (?)"'
+    assert http_spans[0]["values"] == '["saart"]'
+    assert http_spans[0]["ended"] >= http_spans[0]["started"]
+
+    assert http_spans[1]["query"] == '"SELECT users.id, users.name \\nFROM users"'
+    assert http_spans[1]["values"] == "[]"
+    assert http_spans[0]["ended"] >= http_spans[0]["started"]
+
+
+def test_non_existing_table(context, db):
+    @lumigo_tracer(token="123")
+    def lambda_test_function(event, context):
+        others = Table("others", md, Column("id", Integer, primary_key=True))
+        engine = create_engine(db)
+        conn = engine.connect()
+        result = conn.execute(select([others]))
+        return result.fetchone()
+
+    with pytest.raises(OperationalError):
+        lambda_test_function({}, context)
+
+    http_spans = SpansContainer.get_span().spans
+
+    assert len(http_spans) == 1
+    assert http_spans[0]["query"] == '"SELECT others.id \\nFROM others"'
+    assert (
+        http_spans[0]["error"] == '{"type": "OperationalError", "args": ["no such table: others"]}'
+    )
+    assert http_spans[0]["ended"] >= http_spans[0]["started"]
+
+
+def test_pruning_long_strings(context, db):
+    @lumigo_tracer(token="123")
+    def lambda_test_function(event, context):
+        engine = create_engine(db)
+        conn = engine.connect()
+        conn.execute(Users.insert().values(name="a" * (DEFAULT_MAX_ENTRY_SIZE * 5)))
+
+    lambda_test_function({}, context)
+    http_spans = SpansContainer.get_span().spans
+
+    assert len(http_spans) == 1
+    assert len(http_spans[0]["values"]) <= DEFAULT_MAX_ENTRY_SIZE * 2
+
+
+def test_exception_in_wrapper(context, db, monkeypatch):
+    @lumigo_tracer(token="123")
+    def lambda_test_function(event, context):
+        engine = create_engine(db)
+        conn = engine.connect()
+        conn.execute(Users.insert().values(name="a" * (DEFAULT_MAX_ENTRY_SIZE * 5)))
+
+    monkeypatch.setattr(uuid, "uuid4", lambda: 1 / 0)
+
+    lambda_test_function({}, context)  # No exception
+
+    assert not SpansContainer.get_span().spans

src/test/unit/wrappers/sql/test_sqlalchemy_wrapper.py

@@ -13,8 +13,12 @@ def test_wrapping_without_libraries(monkeypatch):
     monkeypatch.setitem(sys.modules, "redis", None)
     importlib.reload(lumigo_tracer.wrappers.redis.redis_wrapper)
 
+    monkeypatch.setitem(sys.modules, "sqlalchemy", None)
+    importlib.reload(lumigo_tracer.wrappers.sql.sqlalchemy_wrapper)
+
     lumigo_tracer.wrappers.wrap(force=True)  # should succeed
 
     monkeypatch.undo()
     importlib.reload(lumigo_tracer.wrappers.pymongo.pymongo_wrapper)
     importlib.reload(lumigo_tracer.wrappers.redis.redis_wrapper)
+    importlib.reload(lumigo_tracer.wrappers.sql.sqlalchemy_wrapper)