New couchbase instrumentation

.circleci/config.yml

@@ -15,6 +15,7 @@ jobs:
       - image: circleci/mariadb:10.1-ram
       - image: circleci/redis:5.0.4
       - image: rabbitmq:3.5.4
+      - image: couchbase/server-sandbox:5.5.0
 
     working_directory: ~/repo
 
@@ -31,6 +32,12 @@ jobs:
       - run:
           name: install dependencies
           command: |
+            sudo apt-get update
+            sudo apt install lsb-release -y
+            curl -O https://packages.couchbase.com/releases/couchbase-release/couchbase-release-1.0-6-amd64.deb
+            sudo dpkg -i ./couchbase-release-1.0-6-amd64.deb
+            sudo apt-get update
+            sudo apt install libcouchbase-dev -y
             rm -rf venv
             export PATH=/home/circleci/.local/bin:$PATH
             pip install --user -U pip setuptools virtualenv
@@ -66,6 +73,7 @@ jobs:
       - image: circleci/mariadb:10-ram
       - image: circleci/redis:5.0.4
       - image: rabbitmq:3.5.4
+      - image: couchbase/server-sandbox:5.5.0
 
     working_directory: ~/repo
 
@@ -82,6 +90,12 @@ jobs:
       - run:
           name: install dependencies
           command: |
+            sudo apt-get update
+            sudo apt install lsb-release -y
+            curl -O https://packages.couchbase.com/releases/couchbase-release/couchbase-release-1.0-6-amd64.deb
+            sudo dpkg -i ./couchbase-release-1.0-6-amd64.deb
+            sudo apt-get update
+            sudo apt install libcouchbase-dev -y
             python -m venv venv
             . venv/bin/activate
             pip install -U pip
@@ -115,6 +129,7 @@ jobs:
       - image: circleci/mariadb:10-ram
       - image: circleci/redis:5.0.4
       - image: rabbitmq:3.5.4
+      - image: couchbase/server-sandbox:5.5.0
 
     working_directory: ~/repo
 
@@ -131,6 +146,12 @@ jobs:
       - run:
           name: install dependencies
           command: |
+            sudo apt-get update
+            sudo apt install lsb-release -y
+            curl -O https://packages.couchbase.com/releases/couchbase-release/couchbase-release-1.0-6-amd64.deb
+            sudo dpkg -i ./couchbase-release-1.0-6-amd64.deb
+            sudo apt-get update
+            sudo apt install libcouchbase-dev -y
             python -m venv venv
             . venv/bin/activate
             pip install -U pip

instana/__init__.py

@@ -67,6 +67,7 @@ def boot_agent():
         else:
             from .instrumentation import mysqlclient
 
+        from .instrumentation import couchbase_inst
         from .instrumentation import flask
         from .instrumentation import grpcio
         from .instrumentation.tornado import client

instana/instrumentation/couchbase_inst.py

@@ -0,0 +1,89 @@
+"""
+couchbase instrumentation - This instrumentation supports the Python CouchBase 2.3.4 --> 2.5.x SDK currently:
+https://docs.couchbase.com/python-sdk/2.5/start-using-sdk.html
+"""
+from __future__ import absolute_import
+
+from distutils.version import LooseVersion
+import wrapt
+
+from ..log import logger
+from ..singletons import tracer
+
+try:
+    import couchbase
+    from couchbase.n1ql import N1QLQuery
+
+    # List of operations to instrument
+    # incr, incr_multi, decr, decr_multi, retrieve_in are wrappers around operations above
+    operations = ['upsert', 'insert', 'replace', 'append', 'prepend', 'get', 'rget',
+                  'touch', 'lock', 'unlock', 'remove', 'counter', 'mutate_in', 'lookup_in',
+                  'stats', 'ping', 'diagnostics', 'observe',
+
+                  'upsert_multi', 'insert_multi', 'replace_multi', 'append_multi',
+                  'prepend_multi', 'get_multi', 'touch_multi', 'lock_multi', 'unlock_multi',
+                  'observe_multi', 'endure_multi', 'remove_multi', 'counter_multi']
+
+    def capture_kvs(scope, instance, query_arg, op):
+        try:
+            scope.span.set_tag('couchbase.hostname', instance.server_nodes[0])
+            scope.span.set_tag('couchbase.bucket', instance.bucket)
+            scope.span.set_tag('couchbase.type', op)
+
+            if query_arg is not None:
+                query = None
+                if type(query_arg) is N1QLQuery:
+                    query = query_arg.statement
+                else:
+                    query = query_arg
+
+                scope.span.set_tag('couchbase.sql', query)
+        except:
+            # No fail on key capture - best effort
+            pass
+
+    def make_wrapper(op):
+        def wrapper(wrapped, instance, args, kwargs):
+            parent_span = tracer.active_span
+
+            # If we're not tracing, just return
+            if parent_span is None:
+                return wrapped(*args, **kwargs)
+
+            with tracer.start_active_span("couchbase", child_of=parent_span) as scope:
+                capture_kvs(scope, instance, None, op)
+                try:
+                    return wrapped(*args, **kwargs)
+                except Exception as e:
+                    scope.span.log_exception(e)
+                    scope.span.set_tag('couchbase.error', repr(e))
+                    raise
+        return wrapper
+
+    def query_with_instana(wrapped, instance, args, kwargs):
+        parent_span = tracer.active_span
+
+        # If we're not tracing, just return
+        if parent_span is None:
+            return wrapped(*args, **kwargs)
+
+        with tracer.start_active_span("couchbase", child_of=parent_span) as scope:
+            capture_kvs(scope, instance, args[0], 'n1ql_query')
+            try:
+                return wrapped(*args, **kwargs)
+            except Exception as e:
+                scope.span.log_exception(e)
+                scope.span.set_tag('couchbase.error', repr(e))
+                raise
+
+    if hasattr(couchbase, '__version__') \
+            and (LooseVersion(couchbase.__version__) >= LooseVersion('2.3.4')) \
+            and (LooseVersion(couchbase.__version__) < LooseVersion('3.0.0')):
+        logger.debug("Instrumenting couchbase")
+        wrapt.wrap_function_wrapper('couchbase.bucket', 'Bucket.n1ql_query', query_with_instana)
+        for op in operations:
+            f = make_wrapper(op)
+            wrapt.wrap_function_wrapper('couchbase.bucket', 'Bucket.%s' % op, f)
+
+except ImportError:
+    pass

instana/json_span.py

@@ -48,6 +48,15 @@ class Data(BaseSpan):
     log = None
 
 
+class CouchbaseData(BaseSpan):
+    hostname = None
+    bucket = None
+    type = None
+    error = None
+    error_code = None
+    sql = None
+
+
 class HttpData(BaseSpan):
     host = None
     url = None

instana/recorder.py

@@ -9,7 +9,7 @@
 
 import instana.singletons
 
-from .json_span import (CustomData, Data, HttpData, JsonSpan, LogData, MySQLData, PostgresData,
+from .json_span import (CouchbaseData, CustomData, Data, HttpData, JsonSpan, LogData, MySQLData, PostgresData,
                         RabbitmqData, RedisData, RenderData, RPCData, SDKData, SoapData,
                         SQLAlchemyData)
 from .log import logger
@@ -23,15 +23,18 @@
 
 class InstanaRecorder(SpanRecorder):
     THREAD_NAME = "Instana Span Reporting"
-    registered_spans = ("aiohttp-client", "aiohttp-server", "django", "log", "memcache", "mysql",
+    registered_spans = ("aiohttp-client", "aiohttp-server", "couchbase", "django", "log", "memcache", "mysql",
                         "postgres", "rabbitmq", "redis", "render", "rpc-client", "rpc-server", "sqlalchemy", "soap",
                         "tornado-client", "tornado-server", "urllib3", "wsgi")
+
     http_spans = ("aiohttp-client", "aiohttp-server", "django", "http", "soap", "tornado-client",
                   "tornado-server", "urllib3", "wsgi")
 
-    exit_spans = ("aiohttp-client", "log", "memcache", "mysql", "postgres", "rabbitmq", "redis", "rpc-client",
+    exit_spans = ("aiohttp-client", "couchbase", "log", "memcache", "mysql", "postgres", "rabbitmq", "redis", "rpc-client",
                   "sqlalchemy", "soap", "tornado-client", "urllib3")
+
     entry_spans = ("aiohttp-server", "django", "wsgi", "rabbitmq", "rpc-server", "tornado-server")
+
     local_spans = ("log", "render")
 
     entry_kind = ["entry", "server", "consumer"]
@@ -161,6 +164,14 @@ def build_registered_span(self, span):
             if data.rabbitmq.sort == 'consume':
                 kind = 1 # entry
 
+        if span.operation_name == "couchbase":
+            data.couchbase = CouchbaseData(hostname=span.tags.pop('couchbase.hostname', None),
+                                           bucket=span.tags.pop('couchbase.bucket', None),
+                                           type=span.tags.pop('couchbase.type', None),
+                                           error=span.tags.pop('couchbase.error', None),
+                                           error_type=span.tags.pop('couchbase.error_type', None),
+                                           sql=span.tags.pop('couchbase.sql', None))
+
         if span.operation_name == "redis":
             data.redis = RedisData(connection=span.tags.pop('connection', None),
                                    driver=span.tags.pop('driver', None),

setup.py

@@ -69,6 +69,7 @@ def check_setuptools():
           'test': [
               'aiohttp>=3.5.4;python_version>="3.5"',
               'asynqp>=0.4;python_version>="3.5"',
+              'couchbase==2.5.9',
               'django>=1.11,<2.2',
               'nose>=1.0',
               'flask>=0.12.2',

tests/helpers.py

@@ -2,6 +2,13 @@
 
 testenv = {}
 
+"""
+CouchDB Environment
+"""
+testenv['couchdb_host'] = os.environ.get('COUCHDB_HOST', '127.0.0.1')
+testenv['couchdb_username'] = os.environ.get('COUCHDB_USERNAME', 'Administrator')
+testenv['couchdb_password'] = os.environ.get('COUCHDB_PASSWORD', 'password')
+
 """
 MySQL Environment
 """