Merge pull request #119 from DataDog/benjamin/pin

LotharSee · web-flow · commit c98c31d7d000 · 2016-11-23T17:39:38.000+01:00
Document `Pin.override()` as the only way to customize a client
diff --git a/ddtrace/contrib/cassandra/__init__.py b/ddtrace/contrib/cassandra/__init__.py
@@ -17,7 +17,7 @@
 
     # Use a pin to specify metadata related to this cluster
     cluster = Cluster(contact_points=['10.1.1.3', '10.1.1.4', '10.1.1.5'], port=9042)
-    Pin.new(service='cassandra-backend').onto(cluster)
+    Pin.override(cluster, service='cassandra-backend')
     session = cluster.connect("my_keyspace")
     session.execute("select id from my_table limit 10;")
 """
diff --git a/ddtrace/contrib/cassandra/session.py b/ddtrace/contrib/cassandra/session.py
@@ -23,7 +23,7 @@ def patch():
     """ patch will add tracing to the cassandra library. """
     setattr(cassandra.cluster.Cluster, 'connect',
             wrapt.FunctionWrapper(_connect, traced_connect))
-    Pin.new(service=SERVICE, app=SERVICE, app_type="db").onto(cassandra.cluster.Cluster)
+    Pin(service=SERVICE, app=SERVICE, app_type="db").onto(cassandra.cluster.Cluster)
 
 def unpatch():
     cassandra.cluster.Cluster.connect = _connect
diff --git a/ddtrace/contrib/dbapi/__init__.py b/ddtrace/contrib/dbapi/__init__.py
@@ -84,7 +84,7 @@ class TracedConnection(wrapt.ObjectProxy):
     def __init__(self, conn):
         super(TracedConnection, self).__init__(conn)
         name = _get_vendor(conn)
-        Pin.new(service=name, app=name).onto(self)
+        Pin(service=name, app=name).onto(self)
 
     def cursor(self, *args, **kwargs):
         cursor = self.__wrapped__.cursor(*args, **kwargs)
diff --git a/ddtrace/contrib/elasticsearch/__init__.py b/ddtrace/contrib/elasticsearch/__init__.py
@@ -16,7 +16,7 @@
 
     # Use a pin to specify metadata related to this client
     es = elasticsearch.Elasticsearch(port=ELASTICSEARCH_CONFIG['port'])
-    Pin(service='elasticsearch-videos').onto(es)
+    Pin.override(es, service='elasticsearch-videos')
     es.indices.create(index='videos', ignore=400)
 """
 from ..util import require_modules
diff --git a/ddtrace/contrib/elasticsearch/patch.py b/ddtrace/contrib/elasticsearch/patch.py
@@ -33,7 +33,7 @@ def __init__(self, *args, **kwargs):
         es = _Elasticsearch(*args, **kwargs)
         super(TracedElasticsearch, self).__init__(es)
 
-        pin = Pin.new(service=DEFAULT_SERVICE, app="elasticsearch", app_type="db")
+        pin = Pin(service=DEFAULT_SERVICE, app="elasticsearch", app_type="db")
         pin.onto(self)
 
         wrapt.wrap_function_wrapper(es.transport, 'perform_request', _perform_request)
diff --git a/ddtrace/contrib/mongoengine/__init__.py b/ddtrace/contrib/mongoengine/__init__.py
@@ -14,9 +14,7 @@
 
     # Use a pin to specify metadata related to this client
     client = mongoengine.connect('db', alias='master')
-    pin = Pin.get_from(client)
-    if pin:
-        pin.clone(service="mongo-master").onto(client)
+    Pin.override(client, service="mongo-master")
 """
 
 from ..util import require_modules
diff --git a/ddtrace/contrib/mongoengine/trace.py b/ddtrace/contrib/mongoengine/trace.py
@@ -17,7 +17,7 @@ class WrappedConnect(wrapt.ObjectProxy):
 
     def __init__(self, connect):
         super(WrappedConnect, self).__init__(connect)
-        ddtrace.Pin.new(service=mongox.TYPE, tracer=ddtrace.tracer).onto(self)
+        ddtrace.Pin(service=mongox.TYPE, tracer=ddtrace.tracer).onto(self)
 
     def __call__(self, *args, **kwargs):
         client = self.__wrapped__(*args, **kwargs)
@@ -33,6 +33,6 @@ def __call__(self, *args, **kwargs):
                 app_type=AppTypes.db,
             )
             client = TracedMongoClient(client)
-            ddtrace.Pin.new(pin.service, tracer=pin.tracer).onto(client)
+            ddtrace.Pin(service=pin.service, tracer=pin.tracer).onto(client)
 
         return client
diff --git a/ddtrace/contrib/mysql/__init__.py b/ddtrace/contrib/mysql/__init__.py
@@ -15,7 +15,7 @@
     cursor.execute("SELECT 6*7 AS the_answer;")
 
     # Use a pin to specify metadata related to this connection
-    Pin.get_from(conn).service = 'mysql-users'
+    Pin.override(conn, service='mysql-users')
 
 This package works for mysql.connector version 2.1.x.
 Only the default full-Python integration works. The binary C connector,
diff --git a/ddtrace/contrib/mysql/patch.py b/ddtrace/contrib/mysql/patch.py
@@ -34,7 +34,7 @@ def _connect(func, instance, args, kwargs):
 def patch_conn(conn):
 
     tags = {t: getattr(conn, a, '') for t, a in CONN_ATTR_BY_TAG.items()}
-    pin = Pin.new(service="mysql", app="mysql", app_type="db", tags=tags)
+    pin = Pin(service="mysql", app="mysql", app_type="db", tags=tags)
 
     # grab the metadata from the conn
     wrapped = TracedConnection(conn)
diff --git a/ddtrace/contrib/psycopg/__init__.py b/ddtrace/contrib/psycopg/__init__.py
@@ -15,7 +15,7 @@
     cursor.execute("select * from users where id = 1")
 
     # Use a pin to specify metadata related to this connection
-    Pin.get_from(db).service = 'postgres-users'
+    Pin.override(db, service='postgres-users')
 """
 from ..util import require_modules
 
diff --git a/ddtrace/contrib/psycopg/patch.py b/ddtrace/contrib/psycopg/patch.py
@@ -38,7 +38,7 @@ def patch_conn(conn):
         "db.application" : dsn.get("application_name"),
     }
 
-    Pin.new(
+    Pin(
         service="postgres",
         app="postgres",
         app_type="db",
diff --git a/ddtrace/contrib/pylibmc/__init__.py b/ddtrace/contrib/pylibmc/__init__.py
@@ -16,9 +16,7 @@
     client.set("key1", "value1")
 
     # Use a pin to specify metadata related to this client
-    pin = Pin.get_from(client)
-    if pin:
-        pin.clone(service="memcached-sessions").onto(client)
+    Pin.override(client, service="memcached-sessions")
 """
 
 from ..util import require_modules
diff --git a/ddtrace/contrib/pylibmc/client.py b/ddtrace/contrib/pylibmc/client.py
@@ -41,7 +41,7 @@ def __init__(self, client=None, service=memcached.SERVICE, tracer=None, *args, *
 
         super(TracedClient, self).__init__(client)
 
-        pin = ddtrace.Pin.new(service=service, tracer=tracer)
+        pin = ddtrace.Pin(service=service, tracer=tracer)
         pin.onto(self)
 
         # attempt to collect the pool of urls this client talks to
diff --git a/ddtrace/contrib/pymongo/__init__.py b/ddtrace/contrib/pymongo/__init__.py
@@ -19,9 +19,7 @@
 
     # Use a pin to specify metadata related to this client
     client = pymongo.MongoClient()
-    pin = Pin.get_from(client)
-    if pin:
-        pin.clone(service="mongo-master").onto(client)
+    pin = Pin.override(client, service="mongo-master")
 """
 from ..util import require_modules
 
diff --git a/ddtrace/contrib/pymongo/client.py b/ddtrace/contrib/pymongo/client.py
@@ -30,7 +30,7 @@ def trace_mongo_client(client, tracer, service=mongox.TYPE):
         app_type=AppTypes.db,
     )
     traced_client = TracedMongoClient(client)
-    ddtrace.Pin.new(service=service, tracer=tracer).onto(traced_client)
+    ddtrace.Pin(service=service, tracer=tracer).onto(traced_client)
     return traced_client
 
 
@@ -46,7 +46,7 @@ def __init__(self, client=None, *args, **kwargs):
 
         super(TracedMongoClient, self).__init__(client)
         # Default Pin
-        ddtrace.Pin.new(service=mongox.TYPE).onto(self)
+        ddtrace.Pin(service=mongox.TYPE).onto(self)
         # NOTE[matt] the TracedMongoClient attempts to trace all of the network
         # calls in the trace library. This is good because it measures the
         # actual network time. It's bad because it uses a private API which
diff --git a/ddtrace/contrib/redis/__init__.py b/ddtrace/contrib/redis/__init__.py
@@ -14,7 +14,7 @@
     client.get("my-key")
 
     # Use a pin to specify metadata related to this client
-    Pin(service='redis-queue').onto(client)
+    Pin.override(client, service='redis-queue')
 """
 
 from ..util import require_modules
diff --git a/ddtrace/contrib/redis/patch.py b/ddtrace/contrib/redis/patch.py
@@ -25,7 +25,7 @@ def patch():
     _w('redis', 'Redis.pipeline', traced_pipeline)
     _w('redis.client', 'BasePipeline.execute', traced_execute_pipeline)
     _w('redis.client', 'BasePipeline.immediate_execute_command', traced_execute_command)
-    Pin.new(service="redis", app="redis", app_type="db").onto(redis.StrictRedis)
+    Pin(service="redis", app="redis", app_type="db").onto(redis.StrictRedis)
 
 def unpatch():
     if getattr(redis, '_datadog_patch', False):
diff --git a/ddtrace/contrib/sqlite3/__init__.py b/ddtrace/contrib/sqlite3/__init__.py
@@ -15,7 +15,7 @@
     cursor.execute("select * from users where id = 1")
 
     # Use a pin to specify metadata related to this connection
-    Pin.get_from(db).service = 'sqlite-users'
+    Pin.override(db, service='sqlite-users')
 """
 from .connection import connection_factory
 from .patch import patch
diff --git a/ddtrace/contrib/sqlite3/patch.py b/ddtrace/contrib/sqlite3/patch.py
@@ -26,7 +26,7 @@ def traced_connect(func, _, args, kwargs):
 
 def patch_conn(conn):
     wrapped = TracedSQLite(conn)
-    Pin.new(service="sqlite", app="sqlite").onto(wrapped)
+    Pin(service="sqlite", app="sqlite").onto(wrapped)
     return wrapped
 
 class TracedSQLite(TracedConnection):
diff --git a/ddtrace/pin.py b/ddtrace/pin.py
@@ -1,47 +1,41 @@
-
-from collections import namedtuple
 import logging
 
 import ddtrace
 
-
 log = logging.getLogger(__name__)
 
 
-_pin = namedtuple('_pin', [
-    'service',
-    'app',
-    'app_type',
-    'tags',
-    'tracer'])
-
-
-class Pin(_pin):
+class Pin(object):
     """ Pin (a.k.a Patch INfo) is a small class which is used to
         set tracing metadata on a particular traced connection.
         This is useful if you wanted to, say, trace two different
         database clusters clusters.
 
         >>> conn = sqlite.connect("/tmp/user.db")
-        >>> pin = Pin.get_from(conn)
-        >>> if pin:
-                pin.copy(service="user-db").onto(conn)
+        >>> # Override a pin for a specific connection
+        >>> pin = Pin.override(conn, service="user-db")
         >>> conn = sqlite.connect("/tmp/image.db")
-        >>> pin = Pin.get_from(conn)
-        >>> if pin:
-                pin.copy(service="image-db").onto(conn)
     """
 
-    @staticmethod
-    def new(service, app=None, app_type=None, tags=None, tracer=None):
-        """ Return a new pin. Convience funtion with sane defaults. """
+    __slots__ = ['app', 'app_type', 'service', 'tags', 'tracer', '_initialized']
+
+    def __init__(self, service, app=None, app_type=None, tags=None, tracer=None):
         tracer = tracer or ddtrace.tracer
-        return Pin(
-            service=service,
-            app=app,
-            app_type=app_type,
-            tags=tags,
-            tracer=tracer)
+        self.service = service
+        self.app = app
+        self.app_type = app_type
+        self.tags = tags
+        self.tracer = tracer
+        self._initialized = True
+
+    def __setattr__(self, name, value):
+        if hasattr(self, '_initialized'):
+            raise AttributeError("can't mutate a pin, use override() or clone() instead")
+        super(Pin, self).__setattr__(name, value)
+
+    def __repr__(self):
+        return "Pin(service=%s, app=%s, app_type=%s, tags=%s, tracer=%s)" % (
+            self.service, self.app, self.app_type, self.tags, self.tracer)
 
     @staticmethod
     def get_from(obj):
@@ -55,12 +49,21 @@ def get_from(obj):
 
     @classmethod
     def override(cls, obj, service=None, app=None, app_type=None, tags=None, tracer=None):
+        """Override an object with the given attributes.
+
+        That's the recommended way to customize an already instrumented client, without
+        losing existing attributes.
+
+        >>> conn = sqlite.connect("/tmp/user.db")
+        >>> # Override a pin for a specific connection
+        >>> pin = Pin.override(conn, service="user-db")
+        """
         if not obj:
             return
 
         pin = cls.get_from(obj)
         if not pin:
-            pin = Pin.new(service)
+            pin = Pin(service)
 
         pin.clone(
             service=service,
diff --git a/tests/contrib/cassandra/test.py b/tests/contrib/cassandra/test.py
@@ -146,7 +146,7 @@ def setUp(self):
     def _traced_session(self):
         tracer = get_dummy_tracer()
         # pin the global Cluster to test if they will conflict
-        Pin.new(service=self.TEST_SERVICE, tracer=tracer).onto(Cluster)
+        Pin(service=self.TEST_SERVICE, tracer=tracer).onto(Cluster)
         cluster = Cluster(port=CASSANDRA_CONFIG['port'])
 
         return cluster.connect(self.TEST_KEYSPACE), tracer.writer
@@ -168,10 +168,10 @@ def setUp(self):
     def _traced_session(self):
         tracer = get_dummy_tracer()
         # pin the global Cluster to test if they will conflict
-        Pin.new(service='not-%s' % self.TEST_SERVICE).onto(Cluster)
+        Pin(service='not-%s' % self.TEST_SERVICE).onto(Cluster)
         cluster = Cluster(port=CASSANDRA_CONFIG['port'])
 
-        Pin.new(service=self.TEST_SERVICE, tracer=tracer).onto(cluster)
+        Pin(service=self.TEST_SERVICE, tracer=tracer).onto(cluster)
         return cluster.connect(self.TEST_KEYSPACE), tracer.writer
 
     def test_patch_unpatch(self):
diff --git a/tests/contrib/elasticsearch/test.py b/tests/contrib/elasticsearch/test.py
@@ -138,7 +138,7 @@ def test_elasticsearch(self):
 
         tracer = get_dummy_tracer()
         writer = tracer.writer
-        pin = Pin.new(service=self.TEST_SERVICE, tracer=tracer)
+        pin = Pin(service=self.TEST_SERVICE, tracer=tracer)
         pin.onto(es)
 
         # Test index creation
@@ -215,7 +215,7 @@ def test_patch_unpatch(self):
         patch()
 
         es = elasticsearch.Elasticsearch(port=ELASTICSEARCH_CONFIG['port'])
-        Pin.new(service=self.TEST_SERVICE, tracer=tracer).onto(es)
+        Pin(service=self.TEST_SERVICE, tracer=tracer).onto(es)
 
         # Test index creation
         es.indices.create(index=self.ES_INDEX, ignore=400)
@@ -239,7 +239,7 @@ def test_patch_unpatch(self):
         patch()
 
         es = elasticsearch.Elasticsearch(port=ELASTICSEARCH_CONFIG['port'])
-        Pin.new(service=self.TEST_SERVICE, tracer=tracer).onto(es)
+        Pin(service=self.TEST_SERVICE, tracer=tracer).onto(es)
 
         # Test index creation
         es.indices.create(index=self.ES_INDEX, ignore=400)
diff --git a/tests/contrib/mongoengine/test.py b/tests/contrib/mongoengine/test.py
@@ -150,7 +150,7 @@ class TestMongoEnginePatchConnect(TestMongoEnginePatchConnectDefault):
 
     def get_tracer_and_connect(self):
         tracer = TestMongoEnginePatchConnectDefault.get_tracer_and_connect(self)
-        Pin.new(service=self.TEST_SERVICE, tracer=tracer).onto(mongoengine.connect)
+        Pin(service=self.TEST_SERVICE, tracer=tracer).onto(mongoengine.connect)
         mongoengine.connect(port=MONGO_CONFIG['port'])
 
         return tracer
@@ -184,9 +184,9 @@ class TestMongoEnginePatchClient(TestMongoEnginePatchClientDefault):
     def get_tracer_and_connect(self):
         tracer = get_dummy_tracer()
         # Set a connect-level service, to check that we properly override it
-        Pin.new(service='not-%s' % self.TEST_SERVICE).onto(mongoengine.connect)
+        Pin(service='not-%s' % self.TEST_SERVICE).onto(mongoengine.connect)
         client = mongoengine.connect(port=MONGO_CONFIG['port'])
-        Pin.new(service=self.TEST_SERVICE, tracer=tracer).onto(client)
+        Pin(service=self.TEST_SERVICE, tracer=tracer).onto(client)
 
         return tracer
 
diff --git a/tests/contrib/pylibmc/test.py b/tests/contrib/pylibmc/test.py
@@ -242,7 +242,7 @@ def test_patch_unpatch(self):
         patch()
 
         client = pylibmc.Client([url])
-        Pin.new(service=self.TEST_SERVICE, tracer=tracer).onto(client)
+        Pin(service=self.TEST_SERVICE, tracer=tracer).onto(client)
         client.set("a", 1)
 
         spans = writer.pop()
diff --git a/tests/contrib/pymongo/test.py b/tests/contrib/pymongo/test.py
@@ -263,7 +263,7 @@ def tearDown(self):
     def get_tracer_and_client(self):
         tracer = get_dummy_tracer()
         client = pymongo.MongoClient(port=MONGO_CONFIG['port'])
-        Pin.new(service=self.TEST_SERVICE, tracer=tracer).onto(client)
+        Pin(service=self.TEST_SERVICE, tracer=tracer).onto(client)
         return tracer, client
 
     def test_patch_unpatch(self):
diff --git a/tests/test_pin.py b/tests/test_pin.py
