whatap-python 2.1.0__py3-none-any.whl

whatap/trace/mod/database/psycopg3.py

@@ -0,0 +1,359 @@
+import inspect
+
+from whatap.trace.trace_handler import trace_handler, async_trace_handler
+from whatap.trace.mod.database.util import (
+    interceptor_db_con, interceptor_db_execute, interceptor_db_close,
+    async_interceptor_db_con, async_interceptor_db_execute, async_interceptor_db_close,
+    interceptor_pool_get, interceptor_pool_release,
+    async_interceptor_pool_get, async_interceptor_pool_release
+)
+
+db_info = {}
+
+
+class BaseCursor:
+    def __init__(self, cursor, db_info):
+        self._cursor = cursor
+        self._db_info = db_info
+        self._is_wrapped = True
+
+    def __getattr__(self, name):
+        return getattr(self._cursor, name)
+
+    def __setattr__(self, name, value):
+        if name.startswith('_'):
+            object.__setattr__(self, name, value)
+        else:
+            setattr(self._cursor, name, value)
+
+    @property
+    def connection(self):
+        return self._cursor.connection
+
+    def _execute_wrapper(self, original_execute_method):
+        owner = getattr(original_execute_method, "__self__", None)
+
+        def safe_execute(*args, **kwargs):
+            if args and owner is not None and args[0] is owner:
+                args = args[1:]
+            return original_execute_method(*args, **kwargs)
+
+        return safe_execute
+
+
+class SyncCursor(BaseCursor):
+    def __enter__(self):
+        self._cursor.__enter__()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        return self._cursor.__exit__(exc_type, exc_val, exc_tb)
+
+    def execute(self, *args, **kwargs):
+        if hasattr(self._cursor, '_is_wrapped'):
+            return self._cursor.execute(*args, **kwargs)
+        real_execute = self._cursor.execute
+        safe_fn = self._execute_wrapper(real_execute)
+        return interceptor_db_execute(safe_fn, self._db_info, self._cursor, *args, **kwargs)
+
+    def executemany(self, *args, **kwargs):
+        if hasattr(self._cursor, '_is_wrapped'):
+            return self._cursor.executemany(*args, **kwargs)
+        real_executemany = getattr(self._cursor, "executemany", None)
+        if real_executemany is None:
+            return self._cursor.executemany(*args, **kwargs)
+        safe_fn = self._execute_wrapper(real_executemany)
+        return interceptor_db_execute(safe_fn, self._db_info, self._cursor, *args, **kwargs)
+
+
+class AsyncCursor(BaseCursor):
+    async def __aenter__(self):
+        await self._cursor.__aenter__()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        return await self._cursor.__aexit__(exc_type, exc_val, exc_tb)
+
+    def _async_execute_wrapper(self, original_execute_method):
+        owner = getattr(original_execute_method, "__self__", None)
+
+        async def async_safe_execute(*args, **kwargs):
+            if args and owner is not None and args[0] is owner:
+                args = args[1:]
+            return await original_execute_method(*args, **kwargs)
+
+        return async_safe_execute
+
+    async def execute(self, *args, **kwargs):
+        if hasattr(self._cursor, '_is_wrapped'):
+            return await self._cursor.execute(*args, **kwargs)
+        real_execute = self._cursor.execute
+        safe_fn = self._async_execute_wrapper(real_execute)
+        return await async_interceptor_db_execute(safe_fn, self._db_info, self._cursor, *args, **kwargs)
+
+    async def executemany(self, *args, **kwargs):
+        if hasattr(self._cursor, '_is_wrapped'):
+            return await self._cursor.executemany(*args, **kwargs)
+        real_executemany = getattr(self._cursor, "executemany", None)
+        if real_executemany is None:
+            return await self._cursor.executemany(*args, **kwargs)
+        safe_fn = self._async_execute_wrapper(real_executemany)
+        return await async_interceptor_db_execute(safe_fn, self._db_info, self._cursor, *args, **kwargs)
+
+
+class BaseConnection:
+    def __init__(self, connection, db_info):
+        self._connection = connection
+        self._db_info = db_info
+        self._is_wrapped = True
+
+    def __getattr__(self, name):
+        return getattr(self._connection, name)
+
+    def __setattr__(self, name, value):
+        if name.startswith('_'):
+            object.__setattr__(self, name, value)
+        else:
+            setattr(self._connection, name, value)
+
+    def _execute_wrapper(self, original_execute_method):
+        owner = getattr(original_execute_method, "__self__", None)
+
+        def safe_execute(*args, **kwargs):
+            if args and owner is not None and args[0] is owner:
+                args = args[1:]
+            return original_execute_method(*args, **kwargs)
+
+        return safe_execute
+
+
+class SyncConnection(BaseConnection):
+    def __enter__(self):
+        self._connection.__enter__()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        return self._connection.__exit__(exc_type, exc_val, exc_tb)
+
+    def close(self, *args, **kwargs):
+        if self._db_info.get("pool"):
+            interceptor_pool_release()
+            return self._connection.close(*args, **kwargs)
+
+        real_close = self._connection.close
+        return interceptor_db_close(real_close, *args, **kwargs)
+
+    def cursor(self, *args, **kwargs):
+        real_cursor = self._connection.cursor(*args, **kwargs)
+        if hasattr(real_cursor, '_is_wrapped'):
+            return real_cursor
+        return SyncCursor(real_cursor, self._db_info)
+
+    def execute(self, *args, **kwargs):
+        if hasattr(self._connection, '_is_wrapped'):
+            return self._connection.execute(*args, **kwargs)
+        real_execute = getattr(self._connection, "execute", None)
+        if real_execute is None:
+            return self._connection.execute(*args, **kwargs)
+        safe_fn = self._execute_wrapper(real_execute)
+        return interceptor_db_execute(safe_fn, self._db_info, self._connection, *args, **kwargs)
+
+
+class AsyncConnection(BaseConnection):
+
+    def __init__(self, connection, db_info, proxy=None):
+        super().__init__(connection, db_info)  # 부모 생성자 호출
+        self._pool_proxy = proxy  # proxy 참조 저장
+
+    async def __aenter__(self):
+        await self._connection.__aenter__()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        if self._db_info.get("pool") and self._pool_proxy:
+
+            # Pool connection이고 proxy가 있으면 proxy의 __aexit__ 호출
+            await async_interceptor_pool_release()
+            return await self._pool_proxy.__aexit__(exc_type, exc_val, exc_tb)
+
+        return await self._connection.__aexit__(exc_type, exc_val, exc_tb)
+
+    def cursor(self, *args, **kwargs):
+        real_cursor = self._connection.cursor(*args, **kwargs)
+        if hasattr(real_cursor, '_is_wrapped'):
+            return real_cursor
+        return AsyncCursor(real_cursor, self._db_info)
+
+    def _async_execute_wrapper(self, original_execute_method):
+        owner = getattr(original_execute_method, "__self__", None)
+
+        async def async_safe_execute(*args, **kwargs):
+            if args and owner is not None and args[0] is owner:
+                args = args[1:]
+            return await original_execute_method(*args, **kwargs)
+
+        return async_safe_execute
+
+    async def execute(self, *args, **kwargs):
+        if hasattr(self._connection, '_is_wrapped'):
+            return await self._connection.execute(*args, **kwargs)
+        real_execute = getattr(self._connection, "execute", None)
+        if real_execute is None:
+            return await self._connection.execute(*args, **kwargs)
+        safe_fn = self._async_execute_wrapper(real_execute)
+        return await async_interceptor_db_execute(safe_fn, self._db_info, self._connection, *args, **kwargs)
+
+    async def close(self, *args, **kwargs):
+        if self._db_info.get("pool"):
+            await async_interceptor_pool_release()
+            return await self._connection.close(*args, **kwargs)
+
+        real_close = self._connection.close
+        return await async_interceptor_db_close(real_close, *args, **kwargs)
+
+
+def _is_called_by_pool():
+    try:
+        # 실제로 psycopg_pool은 보통 3-4 레벨 내에서 호출되므로 깊이 5로 결정.
+        filenames = [f.filename for f in inspect.stack(context=0)[2:7]]
+        return any('psycopg_pool' in f for f in filenames if f)
+    except Exception:
+        return False
+
+
+def _sync_wrapper(fn):
+    @trace_handler(fn)
+    def wrapper(*args, **kwargs):
+        if _is_called_by_pool():
+            return fn(*args, **kwargs)
+
+        db_info = {"type": "postgresql"}
+
+        if args:
+            conn_str = args[0]
+            if isinstance(args[0], str) and '=' in args[0]:
+                parsed_kwargs = dict(
+                    x.split('=') for x in conn_str.split()
+                )
+                kwargs.update(parsed_kwargs)
+                db_info.update(kwargs)
+            else:
+                db_info.update({"db_con_stc": "completed"})
+                db_info.update({"db_str": args[0]})
+
+        connection = interceptor_db_con(fn, db_info, *args, **kwargs)
+        if hasattr(connection, '_is_wrapped'):
+            return connection
+        return SyncConnection(connection, dict(db_info))
+
+    return wrapper
+
+
+def _async_wrapper(fn):
+    @async_trace_handler(fn)
+    async def wrapper(*args, **kwargs):
+        if _is_called_by_pool():
+            return await fn(*args, **kwargs)
+
+        db_info = {"type": "postgresql"}
+
+        if args:
+            conn_str = args[0]
+            if isinstance(args[0], str) and '=' in args[0]:
+                parsed_kwargs = dict(
+                    x.split('=') for x in conn_str.split()
+                )
+                kwargs.update(parsed_kwargs)
+                db_info.update(kwargs)
+            else:
+                db_info.update({"db_con_stc": "completed"})
+                db_info.update({"db_str": args[0]})
+
+
+        connection = await async_interceptor_db_con(fn, db_info, *args, **kwargs)
+        if hasattr(connection, '_is_wrapped'):
+            return connection
+        return AsyncConnection(connection, dict(db_info))
+
+    return wrapper
+
+
+def _get_conn_info(connection):
+    conn_info = {}
+    try:
+        if hasattr(connection, 'info'):
+            conn_info['host'] = connection.info.host
+            conn_info['port'] = connection.info.port
+            conn_info['dbname'] = connection.info.dbname
+            conn_info['user'] = connection.info.user
+        elif hasattr(connection, 'dsn'):
+            import re
+            dsn_params = dict(x.split('=') for x in re.sub(r'\s*=\s*', '=', connection.dsn).split())
+            conn_info.update(dsn_params)
+    except Exception as e:
+        print(f"Failed to extract connection info: {e}")
+    return conn_info
+
+
+def _pool_getconn_wrapper(original_getconn):
+    def wrapper(self, *args, **kwargs):
+        connection = original_getconn(self, *args, **kwargs)
+        conn_info = _get_conn_info(connection)
+        db_info = {"type": "postgresql", "pool": True, **conn_info}
+
+        interceptor_pool_get(db_info)
+
+        return SyncConnection(connection, db_info)
+
+    return wrapper
+
+
+class AsyncConnectionProxy:
+    """AsyncConnectionPool.connection()이 반환하는 proxy 객체를 래핑"""
+
+    def __init__(self, proxy):
+        self._proxy = proxy
+        self._wrapped_connection = None
+
+    async def __aenter__(self):
+        connection = await self._proxy.__aenter__()
+        conn_info = _get_conn_info(connection)
+        db_info = {"type": "postgresql", "pool": True, **conn_info}
+        await async_interceptor_pool_get(db_info)
+
+        self._wrapped_connection = AsyncConnection(connection, db_info, proxy=self._proxy)
+        return self._wrapped_connection
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        if self._wrapped_connection:
+            await async_interceptor_pool_release()
+        return await self._proxy.__aexit__(exc_type, exc_val, exc_tb)
+
+
+def _async_pool_connection_wrapper(original_connection):
+    """AsyncConnectionPool.connection() 메서드를 래핑"""
+
+    def wrapper(self, *args, **kwargs):
+        proxy = original_connection(self, *args, **kwargs)
+        return AsyncConnectionProxy(proxy)
+
+    return wrapper
+
+
+def instrument_psycopg(module):
+    original_connect = module.connect
+    module.connect = _sync_wrapper(original_connect)
+    if hasattr(module, 'AsyncConnection') and hasattr(module.AsyncConnection, 'connect'):
+        original_async_connect = module.AsyncConnection.connect
+        module.AsyncConnection.connect = _async_wrapper(original_async_connect)
+
+
+def instrument_psycopg_pool(pool_module):
+    if hasattr(pool_module, 'ConnectionPool'):
+        if hasattr(pool_module.ConnectionPool, 'getconn'):
+            pool_module.ConnectionPool.getconn = _pool_getconn_wrapper(pool_module.ConnectionPool.getconn)
+
+    if hasattr(pool_module, 'AsyncConnectionPool'):
+        if hasattr(pool_module.AsyncConnectionPool, 'connection'):
+            original_connection = pool_module.AsyncConnectionPool.connection
+            pool_module.AsyncConnectionPool.connection = _async_pool_connection_wrapper(original_connection)

whatap/trace/mod/database/redis.py

@@ -0,0 +1,122 @@
+from whatap.trace import get_dict
+from whatap.trace.trace_handler import trace_handler
+from whatap.trace.trace_error import interceptor_step_error
+from whatap.trace.trace_context_manager import TraceContextManager
+import whatap.net.async_sender as async_sender
+from whatap.net.packet_type_enum import PacketTypeEnum
+from whatap.util.date_util import DateUtil
+import numbers
+
+_current_command = None
+
+
+def intercept_send_command(fn, instance, *args, **kwargs):
+    global _current_command
+    command = str(args[0]).upper() if args else None
+    _current_command = command
+
+    if command not in {'SET', 'GET', 'DEL', 'HSET', 'HGET'}:
+        return fn(instance, *args, **kwargs)
+
+    ctx = TraceContextManager.getLocalContext()
+    if not ctx:
+        return fn(instance, *args, **kwargs)
+
+    start_time = DateUtil.nowSystem()
+    ctx.start_time = start_time
+
+    port_or_path = getattr(instance, 'port', getattr(instance, 'path', None))
+    text = 'redis://'
+    text += getattr(instance, 'host', 'localhost')
+    text += ':{}'.format(port_or_path)
+    text += '/'
+    text += str(getattr(instance, 'db', 0))
+    ctx.active_dbc = text
+    ctx.lctx['dbc'] = text
+
+    ctx.active_dbc = 0
+    ctx.db_opening = True
+    datas = [text]
+    ctx.elapsed = DateUtil.nowSystem() - start_time
+    async_sender.send_packet(PacketTypeEnum.TX_DB_CONN, ctx, datas)
+    query = f"{command} " + " ".join([str(arg)[:20] for arg in args[1:]])
+
+    start_time = DateUtil.nowSystem()
+    ctx.start_time = start_time
+    ctx.active_sqlhash = query
+
+    try:
+        callback = fn(instance, *args, **kwargs)
+        return callback
+    except Exception as e:
+        interceptor_step_error(e)
+    finally:
+        ctx.db_opening = False
+        datas = [ctx.lctx.get('dbc', ''), query, 0]
+        ctx.elapsed = DateUtil.nowSystem() - start_time
+        async_sender.send_packet(PacketTypeEnum.TX_SQL, ctx, datas)
+
+
+def interceptor_read_response(fn, *args, **kwargs):
+    global _current_command
+
+    if _current_command not in {'SET', 'GET', 'DEL', 'HSET', 'HGET'}:
+        _current_command = None
+        return fn(*args, **kwargs)
+
+    ctx = TraceContextManager.getLocalContext()
+    if not ctx:
+        _current_command = None
+        return fn(*args, **kwargs)
+
+    start_time = DateUtil.nowSystem()
+    ctx.start_time = start_time
+
+    try:
+        callback = fn(*args, **kwargs)
+
+        if callback:
+            count = -1
+            if isinstance(callback, int):
+                count = callback
+            elif isinstance(callback, list):
+                count = len(callback)
+
+            if count > -1:
+                msg = '{0}: {1}'.format('Fetch bytes', count)
+                datas = [msg, ' ', ' ']
+                ctx.elapsed = 0
+                async_sender.send_packet(PacketTypeEnum.TX_MSG, ctx, datas)
+
+        ctx.active_sqlhash = 0
+        return callback
+    finally:
+        _current_command = None
+
+
+def instrument_redis_connection(module):
+    def wrapper(fn):
+        @trace_handler(fn)
+        def trace(*args, **kwargs):
+            callback = intercept_send_command(fn, args[0], *args[1:], **kwargs)
+            return callback
+
+        return trace
+
+    if hasattr(module, 'Connection') and hasattr(module.Connection, 'send_command'):
+        if not getattr(module.Connection.send_command, '_wrapped', False):
+            module.Connection.send_command = wrapper(module.Connection.send_command)
+            module.Connection.send_command._wrapped = True
+
+    def wrapper(fn):
+        @trace_handler(fn)
+        def trace(*args, **kwargs):
+            callback = interceptor_read_response(fn, *args, **kwargs)
+            return callback
+
+        return trace
+
+    if hasattr(module, 'Connection') and hasattr(module.Connection, 'read_response'):
+        if not getattr(module.Connection.read_response, '_wrapped', False):
+            module.Connection.read_response = wrapper(module.Connection.read_response)
+            module.Connection.read_response._wrapped = True

whatap/trace/mod/database/sqlalchemy.py

@@ -0,0 +1,213 @@
+from whatap.net.packet_type_enum import PacketTypeEnum
+import whatap.net.async_sender as async_sender
+from whatap.trace.trace_handler import trace_handler
+from whatap.trace.trace_error import interceptor_step_error, sendDebugProfile
+from whatap.util.date_util import DateUtil
+from whatap.trace.trace_context_manager import TraceContextManager
+import sys
+from functools import wraps
+
+def instrument_sqlalchemy(module):
+    def wrapper(fn):
+        @trace_handler(fn)
+        def trace(*args, **kwargs):
+            ctx = TraceContextManager.getLocalContext()
+            
+            start_time = DateUtil.nowSystem()
+            ctx.start_time = start_time
+            
+            text = args[0].bind.url.__to_string__().replace('***', '')
+            ctx.lctx['dbc'] = text
+            
+            datas = [' ', ' ', 'DB SESSION INFO: ' + text]
+            ctx.elapsed = DateUtil.nowSystem() - start_time
+            async_sender.send_packet(PacketTypeEnum.TX_MSG, ctx, datas)
+            
+            callback = fn(*args, **kwargs)
+            return callback
+        
+        return trace
+    if sys.modules.get("flask_sqlalchemy") is None:
+        module.Session.get_bind = wrapper(module.Session.get_bind)
+
+def addQuote(arg_dict):
+    quoted_dict = dict()
+
+    for k, v in arg_dict.items():
+        if isinstance(v, str):
+            quoted_dict[k] = "'" + v.replace("'", "\\'") + "'"
+        else:
+            quoted_dict[k] = v
+
+    return quoted_dict
+
+def instrument_sqlalchemy_engine(module):
+    def wrapper(fn):
+        @trace_handler(fn)
+        def trace(*args, **kwargs):
+            ctx = TraceContextManager.getLocalContext()
+            cursor = args[1]
+            query = None
+            if len(args) > 3 and args[3]:
+                try:
+                    ##oracle을 사용하는 경우
+                    if (type(args[3]) == dict) and (":" in args[2]) and ("oracle" in str(args[0])):
+                        oracle_sql_query = args[2]
+                        for k, v in args[3].items():
+                            replaced_key = f":{k}"
+                            replaced_value = f"'{v}'"
+                            oracle_sql_query = oracle_sql_query.replace(replaced_key,
+                                                                        replaced_value) if replaced_key in oracle_sql_query else None
+                        query = oracle_sql_query
+
+                    ##mysql를 사용하는 경우
+                    elif (type(args[3]) == tuple) and ("%s" in args[2]) and ("mysql" in str(args[0])):
+                        my_sql_query = args[2]
+                        for v in args[3]:
+                            replaced_value = f"'{v}'"
+                            my_sql_query = my_sql_query.replace("%s", replaced_value,
+                                                                1) if "%s" in my_sql_query else None
+                        query = my_sql_query
+                    else:
+                        query = args[2] % addQuote(args[3])
+                except Exception as e:
+                    pass
+            # print('instrument_sqlalchemy_engine 2:', query)
+            try:
+                if not query:
+                    query = args[2].decode()
+            except Exception as e:
+                query = str(args[2])
+            # print('instrument_sqlalchemy_engine 3:', query)
+            sqlalchemy_track_skip = False
+            except_track_module = ['pymysql', 'mysqldb', 'psycopg2','psycopg','sqlite3']
+
+            # 이미 DB드라이브 추적 모듈이 존재하면 sqlAlhechmy 의 추적 중지
+            if any(keyword in str(args[0]) for keyword in except_track_module):
+                sqlalchemy_track_skip = True
+
+            if (not query) or (sqlalchemy_track_skip):
+                return fn(*args, **kwargs)
+
+            start_time = DateUtil.nowSystem()
+            ctx.start_time = start_time
+            ctx.active_sqlhash = query
+            # print('instrument_sqlalchemy_engine 4')
+            try:
+                callback = fn(*args, **kwargs)
+                # print('instrument_sqlalchemy_engine 5')
+                return callback
+            except Exception as e:
+                interceptor_step_error(e)
+            finally:
+                # print('instrument_sqlalchemy_engine 6')
+                ctx = TraceContextManager.getLocalContext()
+                if ctx:
+                    datas = [ctx.lctx.get('dbc', ''), query, 0]
+                    ctx.elapsed = DateUtil.nowSystem() - start_time
+                    # print('instrument_sqlalchemy_engine 6.1', datas)
+                    async_sender.send_packet(PacketTypeEnum.TX_SQL, ctx,
+                                           datas)
+                    # print('instrument_sqlalchemy_engine 7')
+                    if hasattr(cursor, 'rowcount'):
+                        count = cursor.rowcount
+                        # print('instrument_sqlalchemy_engine 8')
+                        if count > -1:
+                            desc = '{0}: {1}'.format('Fetch count', count)
+                            datas = [' ', ' ', desc]
+                            ctx.elapsed = 0
+                            async_sender.send_packet(PacketTypeEnum.TX_MSG, ctx, datas)
+                    ctx.active_sqlhash = 0
+
+                # print('instrument_sqlalchemy_engine 9')
+
+        return trace
+
+    # print('database_toolkit.instrument_sqlalchemy_engine step -1')
+    if hasattr(module, 'DefaultDialect') and hasattr(module.DefaultDialect, 'do_execute'):
+        # print('database_toolkit.instrument_sqlalchemy_engine step -2')
+        module.DefaultDialect.do_execute = wrapper(module.DefaultDialect.do_execute)
+        module.DefaultDialect.do_executemany = wrapper(module.DefaultDialect.do_executemany)
+        module.DefaultDialect.do_execute_no_params = wrapper(module.DefaultDialect.do_execute_no_params)
+        # print('database_toolkit.instrument_sqlalchemy_engine step -3')
+
+class _urlbase(object):
+    def __init__(self, url):
+        self.url = url
+
+class connect(_urlbase):
+    def __call__(self,*args, **kwargs):
+        TraceContextManager.addDBPoolIdle(self.url)
+        # import os
+        # print("connect for ",self.url, os.getpid())
+        
+class checkin(_urlbase):
+    def __call__(self,*args, **kwargs):
+        TraceContextManager.addDBPoolIdle(self.url, isCheckIn=True)
+        #print("checkin for ",self.url)
+        
+class checkout(_urlbase):
+    def __call__(self,*args, **kwargs):
+        TraceContextManager.addDBPoolActive(self.url)
+        #print("checkout for ",self.url)
+
+class reset(_urlbase):
+    def __call__(self,*args, **kwargs):
+        #print("reset for ",self.url)
+        pass
+
+class invalidate(_urlbase):
+    def __call__(self,*args, **kwargs):
+        #print("invalidate for ",self.url)
+        TraceContextManager.removeDBPoolIdle(self.url)
+
+class close(_urlbase):
+    def __call__(self,*args, **kwargs):
+        TraceContextManager.removeDBPoolIdle(self.url)
+        
+class detach(_urlbase):
+    def __call__(self,*args, **kwargs):
+        TraceContextManager.removeDBPoolIdle(self.url)
+
+class close_detached(_urlbase):
+    def __call__(self,*args, **kwargs):
+        #print("close_detached for ",self.url)
+        pass
+
+
+def instrument_sqlalchemy_engine_basic(module):
+    def wrapper(fn):
+        @wraps(fn)
+        def trace(*args, **kwargs):
+            engine = fn(*args, **kwargs)
+            if len(args) > 0:
+                url = str(args[0])
+                try:
+                    import sqlalchemy.engine.event as event
+                    
+                    event.listen(engine, "connect", connect(url))
+                    event.listen(engine, "checkin", checkin(url))
+                    event.listen(engine, "checkout", checkout(url))
+                    event.listen(engine, "close", close(url))
+                except Exception as e:
+                    pass
+                try:
+                    import sqlalchemy.event.api as event
+                    
+                    event.listen(engine, "connect", connect(url))
+                    event.listen(engine, "checkin", checkin(url))
+                    event.listen(engine, "checkout", checkout(url))
+                    
+                    #event.listen(engine, "reset", reset(url))
+                    #event.listen(engine, "invalidate", invalidate(url))
+                    event.listen(engine, "close", close(url))
+                    #event.listen(engine, "detach", detach(url))
+                    #event.listen(engine, "close_detached", close_detached(url))
+
+                except Exception as e:
+                    pass
+            return engine
+        
+        return trace
+    if hasattr(module, "create_engine"):
+        module.create_engine = wrapper(module.create_engine)