whatap-python 2.1.0__py3-none-any.whl

whatap/trace/mod/database/mongo.py

@@ -0,0 +1,169 @@
+import pymongo.monitoring
+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 os, sys, logging
+import traceback
+
+from whatap.util.frame_util import get_current_frame
+
+from pprint import pprint, pformat
+
+class WhatapMongoCmdEventListener(pymongo.monitoring.CommandListener):
+    def __init__(self):
+        pass
+
+    def started(self, event):
+        try:
+            self.__started(event)
+        except Exception as e:
+            ctx = TraceContextManager.getLocalContext()
+            if ctx:
+                logging.debug(e, extra={'id': 'MON22', 'ctx.id':ctx.id}, exc_info=True)
+            else:
+                logging.debug(e, extra={'id': 'MON22'}, exc_info=True)
+
+    def __started(self, event):
+        ctx = TraceContextManager.getLocalContext()
+        if not ctx:
+            return
+        start_time = DateUtil.nowSystem()
+        ctx.start_time = start_time
+
+        port_or_path = event.database_name
+        text = 'mongodb://'
+        if hasattr(event, 'connection_id') and len(event.connection_id) == 2:
+            text += '{}:{}'.format(*event.connection_id)
+
+        text += '/'
+        text += event.database_name
+        ctx.active_dbc = text
+        ctx.lctx['dbc'] = text
+
+        ctx.active_dbc = 0
+        ctx.db_opening = True
+        datas = [text]
+        ctx.elapsed = 0
+        async_sender.send_packet(PacketTypeEnum.TX_DB_CONN, ctx, datas)
+
+        start_time = DateUtil.nowSystem()
+        ctx.start_time = start_time
+
+    def succeeded(self, event):
+        try:
+            self.__succeeded(event)
+        except Exception as e:
+            ctx = TraceContextManager.getLocalContext()
+            if ctx:
+                logging.debug(e, extra={'id': 'MON60', 'ctx.id':ctx.id}, exc_info=True)
+            else:
+                logging.debug(e, extra={'id': 'MON60'}, exc_info=True)
+
+
+    def __succeeded(self, event):
+        ctx = TraceContextManager.getLocalContext()
+        if not ctx:
+            return
+        query = 'command: {}, operation_id: {} '.format(event.command_name, event.operation_id)
+        if ctx.active_sqlhash:
+            query +=' query: {}'.format(ctx.active_sqlhash)
+
+        ctx.db_opening = False
+        datas = [ctx.lctx.get('dbc', ''), query, 0]
+        ctx.elapsed = int(event.duration_micros / 1000)
+        async_sender.send_packet(PacketTypeEnum.TX_SQL, ctx,
+                               datas)
+
+    def failed(self, event):
+        try:
+            self.__failed(event)
+        except Exception as e:
+            ctx = TraceContextManager.getLocalContext()
+            if ctx:
+                logging.debug(e, extra={'id': 'MON85', 'ctx.id':ctx.id}, exc_info=True)
+            else:
+                logging.debug(e, extra={'id': 'MON85'}, exc_info=True)
+
+    def __failed(self, event):
+        ctx = TraceContextManager.getLocalContext()
+        if not ctx:
+            return
+
+        if hasattr(event, 'errtype'):
+            ctx.error_step = event.errtype
+        if not ctx.error:
+            ctx.error = 1
+
+        errors = []
+        if hasattr(event, 'errtype'):
+            errors.append(event.errtype)
+        else:
+            errors.append(str())
+        errors.append(str())
+
+        error = 'command: {}, operation_id: {}'.format(event.command_name, event.operation_id)
+        if ctx.active_sqlhash:
+            error +=' query: {}'.format(ctx.active_sqlhash)
+        frame = get_current_frame(ctx)
+        if not frame:
+            return
+
+        for stack in traceback.extract_stack(frame):
+            line = stack[0]
+            line_num = stack[1]
+            method_name = stack[2]
+
+            if 'whatap' + os.sep + 'trace' in line or 'threading.py' in line:
+                continue
+            error += '{} ({}:{})\n'.format(method_name,line, line_num)
+
+        errors.append(error)
+        ctx.elapsed = int(event.duration_micros / 1000)
+        async_sender.send_packet(PacketTypeEnum.TX_ERROR, ctx, errors)
+
+def instrument_mongo_client(module):
+
+    def wrapper_init(fn):
+        def trace(*args, **kwargs):
+            _kwargs = dict()
+            if kwargs:
+                _kwargs.update(kwargs)
+
+            event_listeners = [WhatapMongoCmdEventListener()]
+            if 'event_listeners' in _kwargs:
+                if isinstance(_kwargs['event_listeners'], list):
+                    event_listeners += _kwargs['event_listeners']
+                else:
+                    event_listeners = _kwargs['event_listeners']
+            _kwargs['event_listeners'] = event_listeners
+
+            callback = fn(*args, **_kwargs)
+            return callback
+
+        return trace
+
+    def wrapper_find(fn):
+        def trace(*args, **kwargs):
+            ctx = TraceContextManager.getLocalContext()
+            if ctx:
+                if len(args) > 1 and isinstance(args[1], dict):
+                    query = pformat(args[1], compact=True)
+
+                    ctx.active_sqlhash = query
+            callback = fn(*args, **kwargs)
+            return callback
+        return trace
+
+    if hasattr(module, 'version_tuple'):
+        (major, minor, buildno) = module.version_tuple
+        if major >= 3:
+            if hasattr(module, 'monitoring'):
+                # module.MongoClient.__init__ = wrapper_init(module.MongoClient.__init__)
+                module.monitoring.register(WhatapMongoCmdEventListener())
+            if hasattr(module, 'collection') and hasattr(module.collection, 'Collection'):
+                module.collection.Collection.find = wrapper_find(module.collection.Collection.find)

whatap/trace/mod/database/mysql.py

@@ -0,0 +1,80 @@
+from whatap.trace import get_dict
+from whatap.trace.trace_handler import trace_handler
+from whatap.trace.mod.database.util import interceptor_db_con, interceptor_db_execute, interceptor_db_close
+db_info = {}
+
+def _preload_mysql_db_info(*args, **kwargs):
+    global db_info
+    db_info = {'type': 'mysql'}
+    db_info.update(kwargs)
+
+def instrument_MySQLdb(module):
+    def wrapper(fn):
+        @trace_handler(fn, preload=_preload_mysql_db_info)
+        def trace(*args, **kwargs):
+            callback = interceptor_db_con(fn, db_info, *args, **kwargs)
+            return callback
+
+        return trace
+
+    if hasattr(module, "connect"):
+        module.connect = wrapper(module.connect)
+
+    def wrapper(fn):
+        @trace_handler(fn)
+        def trace(*args, **kwargs):
+            callback = interceptor_db_close(fn, *args, **kwargs)
+            return callback
+
+        return trace
+
+    if hasattr(module, "connection") and hasattr(module.connection, "close"):
+        get_dict(module.connection)['close'] = wrapper(
+            module.connection.close)
+
+
+def instrument_MySQLdb_cursors(module):
+    def wrapper(fn):
+        @trace_handler(fn)
+        def trace(*args, **kwargs):
+            callback = interceptor_db_execute(fn, db_info, *args, **kwargs)
+            return callback
+
+        return trace
+
+    module.BaseCursor.execute = wrapper(module.BaseCursor.execute)
+    if hasattr(module.BaseCursor, 'callproc'):
+        module.BaseCursor.callproc = wrapper(module.BaseCursor.callproc)
+
+def instrument_pymysql(module):
+    def wrapper(fn):
+        @trace_handler(fn, preload=_preload_mysql_db_info)
+        def trace(*args, **kwargs):
+            callback = interceptor_db_con(fn, db_info, *args, **kwargs)
+            return callback
+
+        return trace
+
+    module.connect = wrapper(module.connect)
+
+    def wrapper(fn):
+        @trace_handler(fn)
+        def trace(*args, **kwargs):
+            callback = interceptor_db_close(fn, *args, **kwargs)
+            return callback
+
+        return trace
+
+    module.connections.Connection.close = wrapper(
+        module.connections.Connection.close)
+
+def instrument_pymysql_cursors(module):
+    def wrapper(fn):
+        @trace_handler(fn)
+        def trace(*args, **kwargs):
+            callback = interceptor_db_execute(fn, db_info, *args, **kwargs)
+            return callback
+
+        return trace
+
+    module.Cursor.execute = wrapper(module.Cursor.execute)

whatap/trace/mod/database/neo4j.py

@@ -0,0 +1,90 @@
+from whatap.trace.trace_handler import trace_handler
+from whatap.trace.mod.database.util import interceptor_db_con, interceptor_db_execute, interceptor_db_close
+
+db_info = {}
+
+
+
+
+
+def instrument_neo4j(module):
+
+    if hasattr(module, 'GraphDatabase'):
+        orig_driver = module.GraphDatabase.driver
+        def wrapper(fn):
+            @trace_handler(fn)
+            def trace(*args, **kwargs):
+                db_info = {'type': 'neo4j'}
+
+                if args:
+                    db_info['uri'] = args[0]
+
+                auth = kwargs.get('auth') or (args[1] if len(args) > 1 else None)
+                if auth:
+                    user = getattr(auth, 'principal', None) or (auth[0] if isinstance(auth, tuple) else None)
+                    if user:
+                        db_info['user'] = user
+
+                callback = interceptor_db_con(fn, db_info, *args, **kwargs)
+                return callback
+
+            return trace
+        module.GraphDatabase.driver = wrapper(orig_driver)
+
+
+
+    if hasattr(module, 'Driver'):
+        orig_driver_close = module.Driver.close
+        def wrapper(fn):
+            @trace_handler(fn)
+            def trace(driver, *args, **kwargs):
+                callback = interceptor_db_close(fn, driver, *args, **kwargs)
+                return callback
+
+            return trace
+
+        module.Driver.close = wrapper(orig_driver_close)
+
+    orig = module.Session.run
+    def wrapper(fn):
+        @trace_handler(fn)
+        def trace(session, *args, **kwargs):
+            db_info = {'type': 'neo4j'}
+            try:
+                setattr(session, 'rowcount', -1)
+            except Exception as e:
+                raise e
+            callback = interceptor_db_execute(fn, db_info, session, *args, **kwargs)
+
+            return callback
+        return trace
+    module.Session.run = wrapper(orig)
+
+    def wrapper(fn):
+        @trace_handler(fn)
+        def trace(tx, *args, **kwargs):
+            session = None
+            try:
+                session = tx._on_closed.__self__
+            except AttributeError:
+                return fn(tx, *args, **kwargs)
+            db_info = {'type': 'neo4j'}
+            try:
+                setattr(tx, 'rowcount', -1)
+            except Exception:
+                pass
+            callback = interceptor_db_execute(fn, db_info, tx, *args, **kwargs)
+            return callback
+
+        return trace
+
+    tx_classes_to_patch = ['Transaction', 'ManagedTransaction', 'BoltTransaction']
+    for class_name in tx_classes_to_patch:
+        if hasattr(module, class_name):
+            TxClass = getattr(module, class_name)
+            if hasattr(TxClass, 'run'):
+                original_run = getattr(TxClass, 'run')
+                setattr(TxClass, 'run', wrapper(original_run))
+
+
+

whatap/trace/mod/database/psycopg2.py

@@ -0,0 +1,45 @@
+from whatap.trace import get_dict
+from whatap.trace.trace_handler import trace_handler
+from whatap.trace.mod.database.util import interceptor_db_con, interceptor_db_execute, interceptor_db_close
+
+db_info = {}
+
+def _preload_pg_db_info(*args, **kwargs):
+    global db_info
+    db_info = {'type': 'postgresql'}
+    db_info.update(kwargs)
+
+def instrument_psycopg2(module):
+    def wrapper(fn):
+        @trace_handler(fn, preload=_preload_pg_db_info)
+        def trace(*args, **kwargs):
+            callback = interceptor_db_con(fn, db_info, *args, **kwargs)
+            return callback
+
+        return trace
+
+    module.connect = wrapper(module.connect)
+
+
+def instrument_psycopg2_connection(module):
+    def wrapper(fn):
+        @trace_handler(fn)
+        def trace(*args, **kwargs):
+            callback = interceptor_db_close(fn, *args, **kwargs)
+            return callback
+
+        return trace
+
+    get_dict(module.connection)['close'] = wrapper(module.connection.close)
+
+
+def instrument_psycopg2_extensions(module):
+    def wrapper(fn):
+        @trace_handler(fn)
+        def trace(*args, **kwargs):
+            callback = interceptor_db_execute(fn, db_info, *args, **kwargs)
+            return callback
+
+        return trace
+
+    get_dict(module.cursor)['execute'] = wrapper(module.cursor.execute)