PyMySQL Instrumentation, support and tests (#166)

* PyMySQL Instrumentation, support and tests

* Add py2 special case

* Add sql_sanitizer and sanitize queries

Author: pglombardo

instana/__init__.py

@@ -65,6 +65,7 @@ def boot_agent():
         from .instrumentation.tornado import server
         from .instrumentation import logging
         from .instrumentation import mysqlpython
+        from .instrumentation import pymysql
         from .instrumentation import redis
         from .instrumentation import sqlalchemy
         from .instrumentation import sudsjurko

instana/instrumentation/pep0249.py

@@ -4,6 +4,7 @@
 
 from ..log import logger
 from ..singletons import tracer
+from ..util import sql_sanitizer
 
 
 class CursorWrapper(wrapt.ObjectProxy):
@@ -20,7 +21,7 @@ def _collect_kvs(self, span, sql):
         try:
             span.set_tag(ext.SPAN_KIND, 'exit')
             span.set_tag(ext.DATABASE_INSTANCE, self._connect_params[1]['db'])
-            span.set_tag(ext.DATABASE_STATEMENT, sql)
+            span.set_tag(ext.DATABASE_STATEMENT, sql_sanitizer(sql))
             span.set_tag(ext.DATABASE_TYPE, 'mysql')
             span.set_tag(ext.DATABASE_USER, self._connect_params[1]['user'])
             span.set_tag('host', "%s:%s" %

instana/instrumentation/pymysql.py

@@ -0,0 +1,17 @@
+from __future__ import absolute_import
+
+from ..log import logger
+from .pep0249 import ConnectionFactory
+
+try:
+    import pymysql #
+
+    cf = ConnectionFactory(connect_func=pymysql.connect, module_name='mysql')
+
+    setattr(pymysql, 'connect', cf)
+    if hasattr(pymysql, 'Connect'):
+        setattr(pymysql, 'Connect', cf)
+
+    logger.debug("Instrumenting pymysql")
+except ImportError:
+    pass

instana/util.py

@@ -182,6 +182,20 @@ def strip_secrets(qp, matcher, kwlist):
         logger.debug("strip_secrets", exc_info=True)
 
 
+def sql_sanitizer(sql):
+    """
+    Removes values from valid SQL statements and returns a stripped version.
+
+    :param sql: The SQL statement to be sanitized
+    :return: String - A sanitized SQL statement without values.
+    """
+    return regexp_sql_values.sub('?', sql)
+
+
+# Used by sql_sanitizer
+regexp_sql_values = re.compile('(\'[\s\S][^\']*\'|\d*\.\d+|\d+|NULL)')
+
+
 def get_default_gateway():
     """
     Attempts to read /proc/self/net/route to determine the default gateway in use.
@@ -230,6 +244,9 @@ def get_py_source(file):
     finally:
         return response
 
+# Used by get_py_source
+regexp_py = re.compile('\.py$')
+
 
 def every(delay, task, name):
     """
@@ -253,5 +270,5 @@ def every(delay, task, name):
         next_time += (time.time() - next_time) // delay * delay + delay
 
 
-# Used by get_py_source
-regexp_py = re.compile('\.py$')
+
+

setup.py

@@ -1,8 +1,8 @@
 # coding: utf-8
-from distutils.version import LooseVersion
-from setuptools import find_packages, setup
 import sys
 from os import path
+from distutils.version import LooseVersion
+from setuptools import find_packages, setup
 
 # Import README.md into long_description
 pwd = path.abspath(path.dirname(__file__))
@@ -16,6 +16,7 @@
 
 
 def check_setuptools():
+    """ Validate that we have min version required of setuptools """
     import pkg_resources
     st_version = pkg_resources.get_distribution('setuptools').version
     if LooseVersion(st_version) < LooseVersion('20.2.2'):
@@ -73,6 +74,7 @@ def check_setuptools():
               'mock>=2.0.0',
               'MySQL-python>=1.2.5;python_version<="2.7"',
               'psycopg2>=2.7.1',
+              'PyMySQL[rsa]>=0.9.1',
               'pyOpenSSL>=16.1.0;python_version<="2.7"',
               'pytest>=3.0.1',
               'redis<3.0.0',
@@ -85,7 +87,8 @@ def check_setuptools():
           ],
       },
       test_suite='nose.collector',
-      keywords=['performance', 'opentracing', 'metrics', 'monitoring', 'tracing', 'distributed-tracing'],
+      keywords=['performance', 'opentracing', 'metrics', 'monitoring',
+                'tracing', 'distributed-tracing'],
       classifiers=[
           'Development Status :: 5 - Production/Stable',
           'Framework :: Django',

tests/helpers.py

@@ -15,13 +15,8 @@
 testenv['mysql_port'] = int(os.environ.get('MYSQL_PORT', '3306'))
 testenv['mysql_db']   = os.environ.get('MYSQL_DB', 'circle_test')
 testenv['mysql_user'] = os.environ.get('MYSQL_USER', 'root')
+testenv['mysql_pw'] = os.environ.get('MYSQL_PW', '')
 
-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

tests/test_mysql-python.py

@@ -1,7 +1,6 @@
 from __future__ import absolute_import
 
 import logging
-import os
 import sys
 from unittest import SkipTest
 

tests/test_pymysql.py

@@ -0,0 +1,249 @@
+from __future__ import absolute_import
+
+import logging
+import sys
+from unittest import SkipTest
+
+from nose.tools import assert_equals
+
+from instana.singletons import tracer
+
+from .helpers import testenv
+
+import pymysql
+
+logger = logging.getLogger(__name__)
+
+create_table_query = 'CREATE TABLE IF NOT EXISTS users(id serial primary key, \
+                      name varchar(40) NOT NULL, email varchar(40) NOT NULL)'
+
+create_proc_query = """
+CREATE PROCEDURE test_proc(IN t VARCHAR(255))
+BEGIN
+    SELECT name FROM users WHERE name = t;
+END
+"""
+
+db = pymysql.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)
+
+while cursor.nextset() is not None:
+    pass
+
+cursor.execute('DROP PROCEDURE IF EXISTS test_proc')
+
+while cursor.nextset() is not None:
+    pass
+
+cursor.execute(create_proc_query)
+
+while cursor.nextset() is not None:
+    pass
+
+cursor.close()
+db.close()
+
+
+class TestPyMySQL:
+    def setUp(self):
+        logger.warn("MySQL connecting: %s:<pass>@%s:3306/%s", testenv['mysql_user'], testenv['mysql_host'], testenv['mysql_db'])
+        self.db = pymysql.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()
+        tracer.cur_ctx = None
+
+    def tearDown(self):
+        """ Do nothing for now """
+        return None
+
+    def test_vanilla_query(self):
+        self.cursor.execute("""SELECT * from users""")
+        result = self.cursor.fetchone()
+        assert_equals(3, len(result))
+
+        spans = self.recorder.queued_spans()
+        assert_equals(0, len(spans))
+
+    def test_basic_query(self):
+        result = None
+        with tracer.start_active_span('test'):
+            result = self.cursor.execute("""SELECT * from users""")
+            self.cursor.fetchone()
+
+        assert(result >= 0)
+
+        spans = self.recorder.queued_spans()
+        assert_equals(2, len(spans))
+
+        db_span = spans[0]
+        test_span = spans[1]
+
+        assert_equals("test", test_span.data.sdk.name)
+        assert_equals(test_span.t, db_span.t)
+        assert_equals(db_span.p, test_span.s)
+
+        assert_equals(None, db_span.error)
+        assert_equals(None, db_span.ec)
+
+        assert_equals(db_span.n, "mysql")
+        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" % testenv['mysql_host'])
+
+    def test_query_with_params(self):
+        result = None
+        with tracer.start_active_span('test'):
+            result = self.cursor.execute("""SELECT * from users where id=1""")
+            self.cursor.fetchone()
+
+        assert(result >= 0)
+
+        spans = self.recorder.queued_spans()
+        assert_equals(2, len(spans))
+
+        db_span = spans[0]
+        test_span = spans[1]
+
+        assert_equals("test", test_span.data.sdk.name)
+        assert_equals(test_span.t, db_span.t)
+        assert_equals(db_span.p, test_span.s)
+
+        assert_equals(None, db_span.error)
+        assert_equals(None, db_span.ec)
+
+        assert_equals(db_span.n, "mysql")
+        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 where id=?')
+        assert_equals(db_span.data.mysql.host, "%s:3306" % testenv['mysql_host'])
+
+    def test_basic_insert(self):
+        result = None
+        with tracer.start_active_span('test'):
+            result = self.cursor.execute(
+                        """INSERT INTO users(name, email) VALUES(%s, %s)""",
+                        ('beaker', '[email protected]'))
+
+        assert_equals(1, result)
+
+        spans = self.recorder.queued_spans()
+        assert_equals(2, len(spans))
+
+        db_span = spans[0]
+        test_span = spans[1]
+
+        assert_equals("test", test_span.data.sdk.name)
+        assert_equals(test_span.t, db_span.t)
+        assert_equals(db_span.p, test_span.s)
+
+        assert_equals(None, db_span.error)
+        assert_equals(None, db_span.ec)
+
+        assert_equals(db_span.n, "mysql")
+        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" % testenv['mysql_host'])
+
+    def test_executemany(self):
+        result = None
+        with tracer.start_active_span('test'):
+            result = self.cursor.executemany("INSERT INTO users(name, email) VALUES(%s, %s)",
+                                             [('beaker', '[email protected]'), ('beaker', '[email protected]')])
+            self.db.commit()
+
+        assert_equals(2, result)
+
+        spans = self.recorder.queued_spans()
+        assert_equals(2, len(spans))
+
+        db_span = spans[0]
+        test_span = spans[1]
+
+        assert_equals("test", test_span.data.sdk.name)
+        assert_equals(test_span.t, db_span.t)
+        assert_equals(db_span.p, test_span.s)
+
+        assert_equals(None, db_span.error)
+        assert_equals(None, db_span.ec)
+
+        assert_equals(db_span.n, "mysql")
+        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" % testenv['mysql_host'])
+
+    def test_call_proc(self):
+        result = None
+        with tracer.start_active_span('test'):
+            result = self.cursor.callproc('test_proc', ('beaker',))
+
+        assert(result)
+
+        spans = self.recorder.queued_spans()
+        assert_equals(2, len(spans))
+
+        db_span = spans[0]
+        test_span = spans[1]
+
+        assert_equals("test", test_span.data.sdk.name)
+        assert_equals(test_span.t, db_span.t)
+        assert_equals(db_span.p, test_span.s)
+
+        assert_equals(None, db_span.error)
+        assert_equals(None, db_span.ec)
+
+        assert_equals(db_span.n, "mysql")
+        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" % testenv['mysql_host'])
+
+    def test_error_capture(self):
+        result = None
+        span = None
+        try:
+            with tracer.start_active_span('test'):
+                result = self.cursor.execute("""SELECT * from blah""")
+                self.cursor.fetchone()
+        except Exception:
+            pass
+        finally:
+            if span:
+                span.finish()
+
+        assert(result is None)
+
+        spans = self.recorder.queued_spans()
+        assert_equals(2, len(spans))
+
+        db_span = spans[0]
+        test_span = spans[1]
+
+        assert_equals("test", test_span.data.sdk.name)
+        assert_equals(test_span.t, db_span.t)
+        assert_equals(db_span.p, test_span.s)
+
+        assert_equals(True, db_span.error)
+        assert_equals(1, db_span.ec)
+
+        if sys.version_info[0] >= 3:
+            # Python 3
+            assert_equals(db_span.data.mysql.error, u'(1146, "Table \'%s.blah\' doesn\'t exist")' % testenv['mysql_db'])
+        else:
+            # Python 2
+            assert_equals(db_span.data.mysql.error, u'(1146, u"Table \'%s.blah\' doesn\'t exist")' % testenv['mysql_db'])
+
+        assert_equals(db_span.n, "mysql")
+        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" % testenv['mysql_host'])