SQLAlchemy Instrumentation (#102)

pglombardo · web-flow · commit a1b631c3db82 · 2018-10-22T19:29:49.000+02:00
* Initial instrumentation, tests and infra.

* Centralize test env configuration

* Fix whitespace

* SQLAlchemy KVs, span conversion

* Moar tests; Add error logging

* Add test for transactions

* Remove auth from connect strings
diff --git a/instana/__init__.py b/instana/__init__.py
@@ -60,9 +60,10 @@ def load_instrumentation():
     if "INSTANA_DISABLE_AUTO_INSTR" not in os.environ:
         # Import & initialize instrumentation
         from .instrumentation import asynqp  # noqa
-        from .instrumentation import urllib3  # noqa
-        from .instrumentation import sudsjurko  # noqa
         from .instrumentation import mysqlpython  # noqa
+        from .instrumentation import sqlalchemy  # noqa
+        from .instrumentation import sudsjurko  # noqa
+        from .instrumentation import urllib3  # noqa
         from .instrumentation.django import middleware  # noqa
 
 
diff --git a/instana/instrumentation/sqlalchemy.py b/instana/instrumentation/sqlalchemy.py
@@ -0,0 +1,71 @@
+from __future__ import absolute_import
+
+import opentracing
+import opentracing.ext.tags as ext
+import wrapt
+import re
+
+from ..log import logger
+from ..singletons import tracer
+
+try:
+    import sqlalchemy
+    from sqlalchemy import event
+    from sqlalchemy.engine import Engine
+
+    url_regexp = re.compile('\/\/(\S+@)')
+
+    @event.listens_for(Engine, 'before_cursor_execute', named=True)
+    def receive_before_cursor_execute(**kw):
+        try:
+            parent_span = tracer.active_span
+
+            # If we're not tracing, just return
+            if parent_span is None:
+                return
+
+            scope = tracer.start_active_span("sqlalchemy", child_of=parent_span)
+            context = kw['context']
+            context._stan_scope = scope
+
+            conn = kw['conn']
+            url = str(conn.engine.url)
+            scope.span.set_tag('sqlalchemy.sql', kw['statement'])
+            scope.span.set_tag('sqlalchemy.eng', conn.engine.name)
+            scope.span.set_tag('sqlalchemy.url', url_regexp.sub('//', url))
+        except Exception as e:
+            logger.debug(e)
+        finally:
+            return
+
+    @event.listens_for(Engine, 'after_cursor_execute', named=True)
+    def receive_after_cursor_execute(**kw):
+        context = kw['context']
+
+        if context is not None and hasattr(context, '_stan_scope'):
+            this_scope = context._stan_scope
+            if this_scope is not None:
+                this_scope.close()
+
+    @event.listens_for(Engine, 'dbapi_error', named=True)
+    def receive_dbapi_error(**kw):
+        context = kw['context']
+
+        if context is not None and hasattr(context, '_stan_scope'):
+            this_scope = context._stan_scope
+            if this_scope is not None:
+                this_scope.span.set_tag("error", True)
+                ec = this_scope.span.tags.get('ec', 0)
+                this_scope.span.set_tag("ec", ec+1)
+
+                if 'exception' in kw:
+                    e = kw['exception']
+                    this_scope.span.set_tag('sqlalchemy.err', str(e))
+                else:
+                    this_scope.span.set_tag('sqlalchemy.err', "No dbapi error specified.")
+                this_scope.close()
+
+
+    logger.debug("Instrumenting sqlalchemy")
+except ImportError:
+    pass
diff --git a/instana/json_span.py b/instana/json_span.py
@@ -26,13 +26,14 @@ def __init__(self, **kwds):
 
 
 class Data(object):
-    service = None
-    http = None
     baggage = None
     custom = None
+    http = None
+    rabbitmq = None
     sdk = None
+    service = None
+    sqlalchemy = None
     soap = None
-    rabbitmq = None
 
     def __init__(self, **kwds):
         self.__dict__.update(kwds)
@@ -70,6 +71,15 @@ class RabbitmqData(object):
     def __init__(self, **kwds):
         self.__dict__.update(kwds)
 
+class SQLAlchemyData(object):
+    sql = None
+    url = None
+    eng = None
+    error = None
+
+    def __init__(self, **kwds):
+        self.__dict__.update(kwds)
+
 
 class SoapData(object):
     action = None
diff --git a/instana/recorder.py b/instana/recorder.py
@@ -12,7 +12,7 @@
 import instana.singletons
 
 from .json_span import (CustomData, Data, HttpData, JsonSpan, MySQLData,
-                        RabbitmqData, SDKData, SoapData)
+                        RabbitmqData, SDKData, SoapData, SQLAlchemyData)
 from .log import logger
 
 if sys.version_info.major is 2:
@@ -23,10 +23,11 @@
 
 class InstanaRecorder(SpanRecorder):
     registered_spans = ("django", "memcache", "mysql", "rabbitmq", "rpc-client",
-                        "rpc-server", "soap", "urllib3", "wsgi")
+                        "rpc-server", "sqlalchemy", "soap", "urllib3", "wsgi")
     http_spans = ("django", "wsgi", "urllib3", "soap")
 
-    exit_spans = ("memcache", "mysql", "rabbitmq", "rpc-client", "soap", "urllib3")
+    exit_spans = ("memcache", "mysql", "rabbitmq", "rpc-client", "sqlalchemy",
+                  "soap", "urllib3")
     entry_spans = ("django", "wsgi", "rabbitmq", "rpc-server")
 
     entry_kind = ["entry", "server", "consumer"]
@@ -113,6 +114,13 @@ def build_registered_span(self, span):
                                          address=span.tags.pop('address', None),
                                          key=span.tags.pop('key', None))
 
+        if span.operation_name == "sqlalchemy":
+            data.sqlalchemy = SQLAlchemyData(sql=span.tags.pop('sqlalchemy.sql', None),
+                                             eng=span.tags.pop('sqlalchemy.eng', None),
+                                             url=span.tags.pop('sqlalchemy.url', None),
+                                             err=span.tags.pop('sqlalchemy.err', None))
+
+
         if span.operation_name == "soap":
             data.soap = SoapData(action=span.tags.pop('soap.action', None))
 
@@ -125,11 +133,6 @@ def build_registered_span(self, span):
                 tskey = list(data.custom.logs.keys())[0]
                 data.mysql.error = data.custom.logs[tskey]['message']
 
-        if len(span.tags) > 0:
-            if data.custom is None:
-                data.custom = CustomData()
-            data.custom.tags = span.tags
-
         entityFrom = {'e': instana.singletons.agent.from_.pid,
                       'h': instana.singletons.agent.from_.agentUuid}
 
@@ -152,6 +155,11 @@ def build_registered_span(self, span):
             json_span.error = error
             json_span.ec = ec
 
+        if len(span.tags) > 0:
+            if data.custom is None:
+                data.custom = CustomData()
+            data.custom.tags = span.tags
+
         return json_span
 
     def build_sdk_span(self, span):
diff --git a/setup.py b/setup.py
@@ -54,12 +54,14 @@ def check_setuptools():
             'lxml>=3.4',
             'mock>=2.0.0',
             'MySQL-python>=1.2.5;python_version<="2.7"',
+            'psycopg2>=2.7.1',
             'pyOpenSSL>=16.1.0;python_version<="2.7"',
             'pytest>=3.0.1',
             'requests>=2.17.1',
-            'urllib3[secure]>=1.15',
+            'sqlalchemy>=1.1.15',
             'spyne>=2.9',
-            'suds-jurko>=0.6'
+            'suds-jurko>=0.6',
+            'urllib3[secure]>=1.15'
         ],
       },
       test_suite='nose.collector',
diff --git a/tests/helpers.py b/tests/helpers.py
@@ -0,0 +1,45 @@
+import os
+
+testenv = {}
+
+"""
+MySQL Environment
+"""
+if 'MYSQL_HOST' in os.environ:
+    testenv['mysql_host']= os.environ['MYSQL_HOST']
+elif 'TRAVIS_MYSQL_HOST' in os.environ:
+    testenv['mysql_host'] = os.environ['TRAVIS_MYSQL_HOST']
+else:
+    testenv['mysql_host'] = '127.0.0.1'
+
+testenv['mysql_port'] = int(os.environ.get('MYSQL_PORT', '3306'))
+testenv['mysql_db']   = os.environ.get('MYSQL_DB', 'travis_ci_test')
+testenv['mysql_user'] = os.environ.get('MYSQL_USER', 'root')
+
+if 'MYSQL_PW' in os.environ:
+    testenv['mysql_pw'] = os.environ['MYSQL_PW']
+elif 'TRAVIS_MYSQL_PASS' in os.environ:
+    testenv['mysql_pw'] = os.environ['TRAVIS_MYSQL_PASS']
+else:
+    testenv['mysql_pw'] = ''
+
+"""
+PostgreSQL Environment
+"""
+if 'POSTGRESQL_HOST' in os.environ:
+    testenv['postgresql_host']= os.environ['POSTGRESQL_HOST']
+elif 'TRAVIS_POSTGRESQL_HOST' in os.environ:
+    testenv['postgresql_host'] = os.environ['TRAVIS_POSTGRESQL_HOST']
+else:
+    testenv['postgresql_host'] = '127.0.0.1'
+
+testenv['postgresql_port'] = int(os.environ.get('POSTGRESQL_PORT', '3306'))
+testenv['postgresql_db']   = os.environ.get('POSTGRESQL_DB', 'travis_ci_test')
+testenv['postgresql_user'] = os.environ.get('POSTGRESQL_USER', 'root')
+
+if 'POSTGRESQL_PW' in os.environ:
+    testenv['postgresql_pw'] = os.environ['POSTGRESQL_PW']
+elif 'TRAVIS_POSTGRESQL_PASS' in os.environ:
+    testenv['postgresql_pw'] = os.environ['TRAVIS_POSTGRESQL_PASS']
+else:
+    testenv['postgresql_pw'] = ''
diff --git a/tests/test_mysql-python.py b/tests/test_mysql-python.py
@@ -9,6 +9,8 @@
 
 from instana.singletons import tracer
 
+from .helpers import testenv
+
 if sys.version_info < (3, 0):
     import MySQLdb
 else:
@@ -17,36 +19,6 @@
 
 logger = logging.getLogger(__name__)
 
-
-if 'MYSQL_HOST' in os.environ:
-    mysql_host = os.environ['MYSQL_HOST']
-elif 'TRAVIS_MYSQL_HOST' in os.environ:
-    mysql_host = os.environ['TRAVIS_MYSQL_HOST']
-else:
-    mysql_host = '127.0.0.1'
-
-if 'MYSQL_PORT' in os.environ:
-    mysql_port = int(os.environ['MYSQL_PORT'])
-else:
-    mysql_port = 3306
-
-if 'MYSQL_DB' in os.environ:
-    mysql_db = os.environ['MYSQL_DB']
-else:
-    mysql_db = "travis_ci_test"
-
-if 'MYSQL_USER' in os.environ:
-    mysql_user = os.environ['MYSQL_USER']
-else:
-    mysql_user = "root"
-
-if 'MYSQL_PW' in os.environ:
-    mysql_pw = os.environ['MYSQL_PW']
-elif 'TRAVIS_MYSQL_PASS' in os.environ:
-    mysql_pw = os.environ['TRAVIS_MYSQL_PASS']
-else:
-    mysql_pw = ''
-
 create_table_query = 'CREATE TABLE IF NOT EXISTS users(id serial primary key, \
                       name varchar(40) NOT NULL, email varchar(40) NOT NULL)'
 
@@ -57,9 +29,9 @@
 END
 """
 
-db = MySQLdb.connect(host=mysql_host, port=mysql_port,
-                     user=mysql_user, passwd=mysql_pw,
-                     db=mysql_db)
+db = MySQLdb.connect(host=testenv['mysql_host'], port=testenv['mysql_port'],
+                     user=testenv['mysql_user'], passwd=testenv['mysql_pw'],
+                     db=testenv['mysql_db'])
 
 cursor = db.cursor()
 cursor.execute(create_table_query)
@@ -83,10 +55,10 @@
 
 class TestMySQLPython:
     def setUp(self):
-        logger.warn("MySQL connecting: %s:<pass>@%s:3306/%s", mysql_user, mysql_host, mysql_db)
-        self.db = MySQLdb.connect(host=mysql_host, port=mysql_port,
-                                  user=mysql_user, passwd=mysql_pw,
-                                  db=mysql_db)
+        logger.warn("MySQL connecting: %s:<pass>@%s:3306/%s", testenv['mysql_user'], testenv['mysql_host'], testenv['mysql_db'])
+        self.db = MySQLdb.connect(host=testenv['mysql_host'], port=testenv['mysql_port'],
+                                  user=testenv['mysql_user'], passwd=testenv['mysql_pw'],
+                                  db=testenv['mysql_db'])
         self.cursor = self.db.cursor()
         self.recorder = tracer.recorder
         self.recorder.clear_spans()
@@ -126,10 +98,10 @@ def test_basic_query(self):
         assert_equals(None, db_span.ec)
 
         assert_equals(db_span.n, "mysql")
-        assert_equals(db_span.data.mysql.db, mysql_db)
-        assert_equals(db_span.data.mysql.user, mysql_user)
+        assert_equals(db_span.data.mysql.db, testenv['mysql_db'])
+        assert_equals(db_span.data.mysql.user, testenv['mysql_user'])
         assert_equals(db_span.data.mysql.stmt, 'SELECT * from users')
-        assert_equals(db_span.data.mysql.host, "%s:3306" % mysql_host)
+        assert_equals(db_span.data.mysql.host, "%s:3306" % testenv['mysql_host'])
 
     def test_basic_insert(self):
         result = None
@@ -154,10 +126,10 @@ def test_basic_insert(self):
         assert_equals(None, db_span.ec)
 
         assert_equals(db_span.n, "mysql")
-        assert_equals(db_span.data.mysql.db, mysql_db)
-        assert_equals(db_span.data.mysql.user, mysql_user)
+        assert_equals(db_span.data.mysql.db, testenv['mysql_db'])
+        assert_equals(db_span.data.mysql.user, testenv['mysql_user'])
         assert_equals(db_span.data.mysql.stmt, 'INSERT INTO users(name, email) VALUES(%s, %s)')
-        assert_equals(db_span.data.mysql.host, "%s:3306" % mysql_host)
+        assert_equals(db_span.data.mysql.host, "%s:3306" % testenv['mysql_host'])
 
     def test_executemany(self):
         result = None
@@ -182,10 +154,10 @@ def test_executemany(self):
         assert_equals(None, db_span.ec)
 
         assert_equals(db_span.n, "mysql")
-        assert_equals(db_span.data.mysql.db, mysql_db)
-        assert_equals(db_span.data.mysql.user, mysql_user)
+        assert_equals(db_span.data.mysql.db, testenv['mysql_db'])
+        assert_equals(db_span.data.mysql.user, testenv['mysql_user'])
         assert_equals(db_span.data.mysql.stmt, 'INSERT INTO users(name, email) VALUES(%s, %s)')
-        assert_equals(db_span.data.mysql.host, "%s:3306" % mysql_host)
+        assert_equals(db_span.data.mysql.host, "%s:3306" % testenv['mysql_host'])
 
     def test_call_proc(self):
         result = None
@@ -208,10 +180,10 @@ def test_call_proc(self):
         assert_equals(None, db_span.ec)
 
         assert_equals(db_span.n, "mysql")
-        assert_equals(db_span.data.mysql.db, mysql_db)
-        assert_equals(db_span.data.mysql.user, mysql_user)
+        assert_equals(db_span.data.mysql.db, testenv['mysql_db'])
+        assert_equals(db_span.data.mysql.user, testenv['mysql_user'])
         assert_equals(db_span.data.mysql.stmt, 'test_proc')
-        assert_equals(db_span.data.mysql.host, "%s:3306" % mysql_host)
+        assert_equals(db_span.data.mysql.host, "%s:3306" % testenv['mysql_host'])
 
     def test_error_capture(self):
         result = None
@@ -240,10 +212,10 @@ def test_error_capture(self):
 
         assert_equals(True, db_span.error)
         assert_equals(1, db_span.ec)
-        assert_equals(db_span.data.mysql.error, '(1146, "Table \'%s.blah\' doesn\'t exist")' % mysql_db)
+        assert_equals(db_span.data.mysql.error, '(1146, "Table \'%s.blah\' doesn\'t exist")' % testenv['mysql_db'])
 
         assert_equals(db_span.n, "mysql")
-        assert_equals(db_span.data.mysql.db, mysql_db)
-        assert_equals(db_span.data.mysql.user, mysql_user)
+        assert_equals(db_span.data.mysql.db, testenv['mysql_db'])
+        assert_equals(db_span.data.mysql.user, testenv['mysql_user'])
         assert_equals(db_span.data.mysql.stmt, 'SELECT * from blah')
-        assert_equals(db_span.data.mysql.host, "%s:3306" % mysql_host)
+        assert_equals(db_span.data.mysql.host, "%s:3306" % testenv['mysql_host'])
diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py
