Add snowflake integration (#2901)

* scaffold snowflake integration

* setup snowflake in CircleCI

* scaffold tests

* work on snowflake tests

* get mocked up authentication flow

* fix tests

* fix circleci

* fix contirb_job mapping

* remove unused import

* use machine executor

* update riotfile for snowflake 2.3.0

* add more tests

* add parallelism to circleci job

* add to monkey.py

* add release note

* Add to docs

* fix link

* Update ddtrace/contrib/snowflake/patch.py

Co-authored-by: Kyle Verhoog <[email protected]>

* Update snowflake docs

* remove todo comment

* remove docstring

* fix passing pin and config to TracedConnection

* add new json fixtures

* fix service env fixture

* fix service from env test

* remove unused import

* Do not assert on err

Co-authored-by: Kyle Verhoog <[email protected]>

Author: brettlangdon

.circleci/config.yml

@@ -748,6 +748,14 @@ jobs:
       - run_tox_scenario:
           pattern: '^sanic_contrib-'
 
+  snowflake:
+    <<: *machine_executor
+    parallelism: 4
+    steps:
+      - run_test:
+          pattern: "snowflake"
+          snapshot: true
+
   starlette:
     <<: *machine_executor
     steps:
@@ -991,6 +999,7 @@ requires_tests: &requires_tests
     - requests
     - rq
     - sanic
+    - snowflake
     - sqlalchemy
     - sqlite3
     - starlette
@@ -1074,6 +1083,7 @@ workflows:
       - requests: *requires_base_venvs
       - rq: *requires_base_venvs
       - sanic: *requires_base_venvs
+      - snowflake: *requires_base_venvs
       - starlette: *requires_base_venvs
       - sqlalchemy: *requires_base_venvs
       - sqlite3: *requires_base_venvs

ddtrace/contrib/snowflake/__init__.py

@@ -0,0 +1,66 @@
+"""
+The snowflake integration instruments the ``snowflake-connector-python`` library to trace Snowflake queries.
+
+Enabling
+~~~~~~~~
+
+The integration is enabled automatically when using
+:ref:`ddtrace-run<ddtracerun>` or :ref:`patch_all()<patch_all>`.
+
+Or use :ref:`patch()<patch>` to manually enable the integration::
+
+    from ddtrace import patch
+    patch(snowflake=True)
+
+
+Global Configuration
+~~~~~~~~~~~~~~~~~~~~
+
+.. py:data:: ddtrace.config.snowflake["service"]
+
+   The service name reported by default for snowflake spans.
+
+   This option can also be set with the ``DD_SNOWFLAKE_SERVICE`` environment
+   variable.
+
+   Default: ``"snowflake"``
+
+.. py:data:: ddtrace.config.snowflake["trace_fetch_methods"]
+
+   Whether or not to trace fetch methods.
+
+   Can also configured via the ``DD_SNOWFLAKE_TRACE_FETCH_METHODS`` environment variable.
+
+   Default: ``False``
+
+
+Instance Configuration
+~~~~~~~~~~~~~~~~~~~~~~
+
+To configure the integration on an per-connection basis use the
+``Pin`` API::
+
+    from ddtrace import Pin
+    from snowflake.connector import connect
+
+    # This will report a span with the default settings
+    conn = connect(user="alice", password="b0b", account="dev")
+
+    # Use a pin to override the service name for this connection.
+    Pin.override(conn, service="snowflake-dev")
+
+
+    cursor = conn.cursor()
+    cursor.execute("SELECT current_version()")
+"""
+from ...utils.importlib import require_modules
+
+
+required_modules = ["snowflake.connector"]
+
+with require_modules(required_modules) as missing_modules:
+    if not missing_modules:
+        from .patch import patch
+        from .patch import unpatch
+
+        __all__ = ["patch", "unpatch"]

ddtrace/contrib/snowflake/patch.py

@@ -0,0 +1,62 @@
+from ddtrace import Pin
+from ddtrace import config
+from ddtrace.vendor import wrapt
+
+from ...ext import db
+from ...ext import net
+from ...utils.formats import asbool
+from ...utils.formats import get_env
+from ...utils.wrappers import unwrap
+from ..dbapi import TracedConnection
+
+
+config._add(
+    "snowflake",
+    dict(
+        _default_service="snowflake",
+        trace_fetch_methods=asbool(get_env("snowflake", "trace_fetch_methods", default=False)),
+    ),
+)
+
+
+def patch():
+    import snowflake.connector
+
+    if getattr(snowflake.connector, "_datadog_patch", False):
+        return
+    setattr(snowflake.connector, "_datadog_patch", True)
+
+    wrapt.wrap_function_wrapper(snowflake.connector, "Connect", patched_connect)
+    wrapt.wrap_function_wrapper(snowflake.connector, "connect", patched_connect)
+
+
+def unpatch():
+    import snowflake.connector
+
+    if getattr(snowflake.connector, "_datadog_patch", False):
+        setattr(snowflake.connector, "_datadog_patch", False)
+
+        unwrap(snowflake.connector, "Connect")
+        unwrap(snowflake.connector, "connect")
+
+
+def patched_connect(connect_func, _, args, kwargs):
+    conn = connect_func(*args, **kwargs)
+    if isinstance(conn, TracedConnection):
+        return conn
+
+    # Add default tags to each query
+    tags = {
+        net.TARGET_HOST: conn.host,
+        net.TARGET_PORT: conn.port,
+        db.NAME: conn.database,
+        db.USER: conn.user,
+        "db.application": conn.application,
+        "db.schema": conn.schema,
+        "db.warehouse": conn.warehouse,
+    }
+
+    pin = Pin(tags=tags)
+    traced_conn = TracedConnection(conn, pin=pin, cfg=config.snowflake)
+    pin.onto(traced_conn)
+    return traced_conn

ddtrace/monkey.py

@@ -53,6 +53,7 @@
     "requests": True,
     "rq": True,
     "sanic": True,
+    "snowflake": True,
     "sqlalchemy": False,  # Prefer DB client instrumentation
     "sqlite3": True,
     "aiohttp": True,  # requires asyncio (Python 3.4+)

docs/index.rst

@@ -136,6 +136,8 @@ contacting support.
 +--------------------------------------------------+---------------+----------------+
 | :ref:`sanic`                                     | >= 19.6.0     | Yes [4]_       |
 +--------------------------------------------------+---------------+----------------+
+| :ref:`snowflake`                                 | >= 2.0.0      | Yes            |
++--------------------------------------------------+---------------+----------------+
 | :ref:`sqlalchemy`                                | >= 1.0        | No             |
 +--------------------------------------------------+---------------+----------------+
 | :ref:`sqlite`                                    | \*            | Yes            |

docs/integrations.rst

@@ -345,6 +345,13 @@ Sanic
 .. automodule:: ddtrace.contrib.sanic
 
 
+.. _snowflake:
+
+Snowflake
+^^^^^^^^^
+.. automodule:: ddtrace.contrib.snowflake
+
+
 .. _starlette:
 
 Starlette

releasenotes/notes/add-snowflake-9fdedcd5058de2f7.yaml

@@ -0,0 +1,4 @@
+---
+features:
+  - |
+    Add support for `snowflake-connector-python <https://pypi.org/project/snowflake-connector-python/>`_ >= 2.0.0.

riotfile.py

@@ -1204,5 +1204,53 @@ def select_pys(min_version=MIN_PYTHON_VERSION, max_version=MAX_PYTHON_VERSION):
                 ],
             },
         ),
+        Venv(
+            name="snowflake",
+            command="pytest {cmdargs} tests/contrib/snowflake",
+            pkgs={
+                "responses": latest,
+            },
+            venvs=[
+                Venv(
+                    # 2.2.0 dropped 2.7 support
+                    pys=select_pys(),
+                    pkgs={
+                        "snowflake-connector-python": [
+                            "~=2.0.0",
+                            "~=2.1.0",
+                        ],
+                    },
+                ),
+                Venv(
+                    # 2.3.7 dropped 3.5 support
+                    pys=select_pys(min_version="3.5"),
+                    pkgs={
+                        "snowflake-connector-python": [
+                            "~=2.2.0",
+                        ],
+                    },
+                ),
+                Venv(
+                    # 2.3.x needs pyarrow >=0.17,<0.18 which does not install on Python 3.9
+                    pys=select_pys(min_version="3.6", max_version="3.8"),
+                    pkgs={
+                        "snowflake-connector-python": [
+                            "~=2.3.0",
+                        ],
+                    },
+                ),
+                Venv(
+                    pys=select_pys(min_version="3.6"),
+                    pkgs={
+                        "snowflake-connector-python": [
+                            "~=2.4.0",
+                            "~=2.5.0",
+                            "~=2.6.0",
+                            latest,
+                        ],
+                    },
+                ),
+            ],
+        ),
     ],
 )