feat(integrations): Add integration for clickhouse-driver (#2167)

Adds an integration that automatically facilitates tracing/recording of all queries, their parameters, data, and results.

.github/workflows/test-integration-clickhouse_driver.yml

@@ -0,0 +1,85 @@
+name: Test clickhouse_driver
+
+on:
+  push:
+    branches:
+      - master
+      - release/**
+
+  pull_request:
+
+# Cancel in progress workflows on pull_requests.
+# https://docs.github.com/en/actions/using-jobs/using-concurrency#example-using-a-fallback-value
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}
+  cancel-in-progress: true
+
+permissions:
+  contents: read
+
+env:
+  BUILD_CACHE_KEY: ${{ github.sha }}
+  CACHED_BUILD_PATHS: |
+    ${{ github.workspace }}/dist-serverless
+
+jobs:
+  test:
+    name: clickhouse_driver, python ${{ matrix.python-version }}, ${{ matrix.os }}
+    runs-on: ${{ matrix.os }}
+    timeout-minutes: 30
+
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: ["3.8","3.9","3.10","3.11"]
+        # python3.6 reached EOL and is no longer being supported on
+        # new versions of hosted runners on Github Actions
+        # ubuntu-20.04 is the last version that supported python3.6
+        # see https://github.com/actions/setup-python/issues/544#issuecomment-1332535877
+        os: [ubuntu-20.04]
+
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v4
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - uses: getsentry/action-clickhouse-in-ci@v1
+
+      - name: Setup Test Env
+        run: |
+          pip install coverage "tox>=3,<4"
+
+      - name: Test clickhouse_driver
+        uses: nick-fields/retry@v2
+        with:
+          timeout_minutes: 15
+          max_attempts: 2
+          retry_wait_seconds: 5
+          shell: bash
+          command: |
+            set -x # print commands that are executed
+            coverage erase
+
+            # Run tests
+            ./scripts/runtox.sh "py${{ matrix.python-version }}-clickhouse_driver" --cov=tests --cov=sentry_sdk --cov-report= --cov-branch &&
+            coverage combine .coverage* &&
+            coverage xml -i
+
+      - uses: codecov/codecov-action@v3
+        with:
+          token: ${{ secrets.CODECOV_TOKEN }}
+          files: coverage.xml
+
+
+  check_required_tests:
+    name: All clickhouse_driver tests passed or skipped
+    needs: test
+    # Always run this, even if a dependent job failed
+    if: always()
+    runs-on: ubuntu-20.04
+    steps:
+      - name: Check for failures
+        if: contains(needs.test.result, 'failure')
+        run: |
+          echo "One of the dependent jobs has failed. You may need to re-run it." && exit 1

scripts/split-tox-gh-actions/ci-yaml-test-snippet.txt

@@ -10,6 +10,7 @@
       - uses: actions/setup-python@v4
         with:
           python-version: ${{ matrix.python-version }}
+{{ additional_uses }}
 
       - name: Setup Test Env
         run: |

scripts/split-tox-gh-actions/split-tox-gh-actions.py

@@ -36,6 +36,10 @@
     "asyncpg",
 ]
 
+FRAMEWORKS_NEEDING_CLICKHOUSE = [
+    "clickhouse_driver",
+]
+
 MATRIX_DEFINITION = """
     strategy:
       fail-fast: false
@@ -48,6 +52,11 @@
         os: [ubuntu-20.04]
 """
 
+ADDITIONAL_USES_CLICKHOUSE = """\
+
+      - uses: getsentry/action-clickhouse-in-ci@v1
+"""
+
 CHECK_NEEDS = """\
     needs: test
 """
@@ -119,6 +128,10 @@ def write_yaml_file(
                 f = open(TEMPLATE_FILE_SETUP_DB, "r")
                 out += "".join(f.readlines())
 
+        elif template_line.strip() == "{{ additional_uses }}":
+            if current_framework in FRAMEWORKS_NEEDING_CLICKHOUSE:
+                out += ADDITIONAL_USES_CLICKHOUSE
+
         elif template_line.strip() == "{{ check_needs }}":
             if py27_supported:
                 out += CHECK_NEEDS_PY27

sentry_sdk/integrations/clickhouse_driver.py

@@ -0,0 +1,150 @@
+from sentry_sdk import Hub
+from sentry_sdk.consts import OP, SPANDATA
+from sentry_sdk.hub import _should_send_default_pii
+from sentry_sdk.integrations import Integration, DidNotEnable
+from sentry_sdk.tracing import Span
+from sentry_sdk._types import TYPE_CHECKING
+from sentry_sdk.utils import capture_internal_exceptions
+
+from typing import TypeVar
+
+# Hack to get new Python features working in older versions
+# without introducing a hard dependency on `typing_extensions`
+# from: https://stackoverflow.com/a/71944042/300572
+if TYPE_CHECKING:
+    from typing import ParamSpec, Callable
+else:
+    # Fake ParamSpec
+    class ParamSpec:
+        def __init__(self, _):
+            self.args = None
+            self.kwargs = None
+
+    # Callable[anything] will return None
+    class _Callable:
+        def __getitem__(self, _):
+            return None
+
+    # Make instances
+    Callable = _Callable()
+
+
+try:
+    import clickhouse_driver  # type: ignore[import]
+
+except ImportError:
+    raise DidNotEnable("clickhouse-driver not installed.")
+
+if clickhouse_driver.VERSION < (0, 2, 0):
+    raise DidNotEnable("clickhouse-driver >= 0.2.0 required")
+
+
+class ClickhouseDriverIntegration(Integration):
+    identifier = "clickhouse_driver"
+
+    @staticmethod
+    def setup_once() -> None:
+        # Every query is done using the Connection's `send_query` function
+        clickhouse_driver.connection.Connection.send_query = _wrap_start(
+            clickhouse_driver.connection.Connection.send_query
+        )
+
+        # If the query contains parameters then the send_data function is used to send those parameters to clickhouse
+        clickhouse_driver.client.Client.send_data = _wrap_send_data(
+            clickhouse_driver.client.Client.send_data
+        )
+
+        # Every query ends either with the Client's `receive_end_of_query` (no result expected)
+        # or its `receive_result` (result expected)
+        clickhouse_driver.client.Client.receive_end_of_query = _wrap_end(
+            clickhouse_driver.client.Client.receive_end_of_query
+        )
+        clickhouse_driver.client.Client.receive_result = _wrap_end(
+            clickhouse_driver.client.Client.receive_result
+        )
+
+
+P = ParamSpec("P")
+T = TypeVar("T")
+
+
+def _wrap_start(f: Callable[P, T]) -> Callable[P, T]:
+    def _inner(*args: P.args, **kwargs: P.kwargs) -> T:
+        hub = Hub.current
+        if hub.get_integration(ClickhouseDriverIntegration) is None:
+            return f(*args, **kwargs)
+        connection = args[0]
+        query = args[1]
+        query_id = args[2] if len(args) > 2 else kwargs.get("query_id")
+        params = args[3] if len(args) > 3 else kwargs.get("params")
+
+        span = hub.start_span(op=OP.DB, description=query)
+
+        connection._sentry_span = span  # type: ignore[attr-defined]
+
+        _set_db_data(span, connection)
+
+        span.set_data("query", query)
+
+        if query_id:
+            span.set_data("db.query_id", query_id)
+
+        if params and _should_send_default_pii():
+            span.set_data("db.params", params)
+
+        # run the original code
+        ret = f(*args, **kwargs)
+
+        return ret
+
+    return _inner
+
+
+def _wrap_end(f: Callable[P, T]) -> Callable[P, T]:
+    def _inner_end(*args: P.args, **kwargs: P.kwargs) -> T:
+        res = f(*args, **kwargs)
+        instance = args[0]
+        span = instance.connection._sentry_span  # type: ignore[attr-defined]
+
+        if span is not None:
+            if res is not None and _should_send_default_pii():
+                span.set_data("db.result", res)
+
+            with capture_internal_exceptions():
+                span.hub.add_breadcrumb(
+                    message=span._data.pop("query"), category="query", data=span._data
+                )
+
+            span.finish()
+
+        return res
+
+    return _inner_end
+
+
+def _wrap_send_data(f: Callable[P, T]) -> Callable[P, T]:
+    def _inner_send_data(*args: P.args, **kwargs: P.kwargs) -> T:
+        instance = args[0]  # type: clickhouse_driver.client.Client
+        data = args[2]
+        span = instance.connection._sentry_span
+
+        _set_db_data(span, instance.connection)
+
+        if _should_send_default_pii():
+            db_params = span._data.get("db.params", [])
+            db_params.extend(data)
+            span.set_data("db.params", db_params)
+
+        return f(*args, **kwargs)
+
+    return _inner_send_data
+
+
+def _set_db_data(
+    span: Span, connection: clickhouse_driver.connection.Connection
+) -> None:
+    span.set_data(SPANDATA.DB_SYSTEM, "clickhouse")
+    span.set_data(SPANDATA.SERVER_ADDRESS, connection.host)
+    span.set_data(SPANDATA.SERVER_PORT, connection.port)
+    span.set_data(SPANDATA.DB_NAME, connection.database)
+    span.set_data(SPANDATA.DB_USER, connection.user)

setup.py

@@ -51,6 +51,7 @@ def get_file_text(file_name):
         "bottle": ["bottle>=0.12.13"],
         "celery": ["celery>=3"],
         "chalice": ["chalice>=1.16.0"],
+        "clickhouse-driver": ["clickhouse-driver>=0.2.0"],
         "django": ["django>=1.8"],
         "falcon": ["falcon>=1.4"],
         "fastapi": ["fastapi>=0.79.0"],

tests/integrations/clickhouse_driver/__init__.py

@@ -0,0 +1,3 @@
+import pytest
+
+pytest.importorskip("clickhouse_driver")