Psycopg3 Instrumentation (#1155)

* Sort tox file * Update pytest * Add pscyopg3 test suite Remove psycopg 3.0 tests Fix tox for psycopg-binary * Add new error message to slow sql validator * Add DBAPI2 async instrumentation * Add psycopg3 instrumentation * Upgrade postgres in CI * Add separate postgres 16 and 9 testing versions * Format and lint * Add more __enter__ comments * Don't needlessly append rollup metrics * [Mega-Linter] Apply linters fixes * Bump tests * Rename dsn to conninfo --------- Co-authored-by: TimPansino <[email protected]> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
newrelic · Jun 10, 2024 · 73f3197 · 73f3197
1 parent 60c5549
commit 73f3197
Show file tree

Hide file tree

Showing 22 changed files with 2,831 additions and 63 deletions.
diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
@@ -43,7 +43,8 @@ jobs:
       - mongodb
       - mssql
       - mysql
-      - postgres
+      - postgres16
+      - postgres9
       - rabbitmq
       - redis
       - rediscluster
@@ -207,7 +208,7 @@ jobs:
           path: ./**/.coverage.*
           retention-days: 1
 
-  postgres:
+  postgres16:
     env:
       TOTAL_GROUPS: 2
 
@@ -223,7 +224,71 @@ jobs:
         --add-host=host.docker.internal:host-gateway
     timeout-minutes: 30
     services:
-      postgres:
+      postgres16:
+        image: postgres:16
+        env:
+          POSTGRES_PASSWORD: postgres
+        ports:
+          - 8080:5432
+          - 8081:5432
+        # Set health checks to wait until postgres has started
+        options: >-
+          --health-cmd pg_isready
+          --health-interval 10s
+          --health-timeout 5s
+          --health-retries 5
+
+    steps:
+      - uses: actions/checkout@b4ffde65f46336ab88eb53be808477a3936bae11 # 4.1.1
+
+      - name: Fetch git tags
+        run: |
+          git config --global --add safe.directory "$GITHUB_WORKSPACE"
+          git fetch --tags origin
+
+      - name: Configure pip cache
+        run: |
+          mkdir -p /github/home/.cache/pip
+          chown -R $(whoami) /github/home/.cache/pip
+
+      - name: Get Environments
+        id: get-envs
+        run: |
+          echo "envs=$(tox -l | grep '^${{ github.job }}\-' | ./.github/workflows/get-envs.py)" >> $GITHUB_OUTPUT
+        env:
+          GROUP_NUMBER: ${{ matrix.group-number }}
+
+      - name: Test
+        run: |
+          tox -vv -e ${{ steps.get-envs.outputs.envs }} -p auto
+        env:
+          TOX_PARALLEL_NO_SPINNER: 1
+          PY_COLORS: 0
+
+      - name: Upload Coverage Artifacts
+        uses: actions/upload-artifact@5d5d22a31266ced268874388b861e4b58bb5c2f3 # 4.3.1
+        with:
+          name: coverage-${{ github.job }}-${{ strategy.job-index }}
+          path: ./**/.coverage.*
+          retention-days: 1
+
+  postgres9:
+    env:
+      TOTAL_GROUPS: 1
+
+    strategy:
+      fail-fast: false
+      matrix:
+        group-number: [1]
+
+    runs-on: ubuntu-20.04
+    container:
+      image: ghcr.io/newrelic/newrelic-python-agent-ci:latest
+      options: >-
+        --add-host=host.docker.internal:host-gateway
+    timeout-minutes: 30
+    services:
+      postgres9:
         image: postgres:9
         env:
           POSTGRES_PASSWORD: postgres

diff --git a/newrelic/config.py b/newrelic/config.py
@@ -3100,6 +3100,9 @@ def _process_module_builtin_defaults():
 
     _process_module_definition("pymssql", "newrelic.hooks.database_pymssql", "instrument_pymssql")
 
+    _process_module_definition("psycopg", "newrelic.hooks.database_psycopg", "instrument_psycopg")
+    _process_module_definition("psycopg.sql", "newrelic.hooks.database_psycopg", "instrument_psycopg_sql")
+
     _process_module_definition("psycopg2", "newrelic.hooks.database_psycopg2", "instrument_psycopg2")
     _process_module_definition(
         "psycopg2._psycopg2",

diff --git a/newrelic/hooks/database_dbapi2.py b/newrelic/hooks/database_dbapi2.py
@@ -14,13 +14,12 @@
 
 from newrelic.api.database_trace import DatabaseTrace, register_database_client
 from newrelic.api.function_trace import FunctionTrace
-from newrelic.api.transaction import current_transaction
 from newrelic.common.object_names import callable_name
-from newrelic.common.object_wrapper import wrap_object, ObjectProxy
-from newrelic.core.config import global_settings
+from newrelic.common.object_wrapper import ObjectProxy, wrap_object
 
 DEFAULT = object()
 
+
 class CursorWrapper(ObjectProxy):
 
     def __init__(self, cursor, dbapi2_module, connect_params, cursor_params):
@@ -31,15 +30,25 @@ def __init__(self, cursor, dbapi2_module, connect_params, cursor_params):
 
     def execute(self, sql, parameters=DEFAULT, *args, **kwargs):
         if parameters is not DEFAULT:
-            with DatabaseTrace(sql, self._nr_dbapi2_module,
-                    self._nr_connect_params, self._nr_cursor_params,
-                    parameters, (args, kwargs), source=self.__wrapped__.execute):
-                return self.__wrapped__.execute(sql, parameters,
-                        *args, **kwargs)
+            with DatabaseTrace(
+                sql=sql,
+                dbapi2_module=self._nr_dbapi2_module,
+                connect_params=self._nr_connect_params,
+                cursor_params=self._nr_cursor_params,
+                sql_parameters=parameters,
+                execute_params=(args, kwargs),
+                source=self.__wrapped__.execute,
+            ):
+                return self.__wrapped__.execute(sql, parameters, *args, **kwargs)
         else:
-            with DatabaseTrace(sql, self._nr_dbapi2_module,
-                    self._nr_connect_params, self._nr_cursor_params,
-                    None, (args, kwargs), source=self.__wrapped__.execute):
+            with DatabaseTrace(
+                sql=sql,
+                dbapi2_module=self._nr_dbapi2_module,
+                connect_params=self._nr_connect_params,
+                cursor_params=self._nr_cursor_params,
+                execute_params=(args, kwargs),
+                source=self.__wrapped__.execute,
+            ):
                 return self.__wrapped__.execute(sql, **kwargs)
 
     def executemany(self, sql, seq_of_parameters):
@@ -49,23 +58,38 @@ def executemany(self, sql, seq_of_parameters):
         except (TypeError, IndexError):
             parameters = DEFAULT
         if parameters is not DEFAULT:
-            with DatabaseTrace(sql, self._nr_dbapi2_module,
-                    self._nr_connect_params, self._nr_cursor_params,
-                    parameters, source=self.__wrapped__.executemany):
+            with DatabaseTrace(
+                sql=sql,
+                dbapi2_module=self._nr_dbapi2_module,
+                connect_params=self._nr_connect_params,
+                cursor_params=self._nr_cursor_params,
+                sql_parameters=parameters,
+                source=self.__wrapped__.executemany,
+            ):
                 return self.__wrapped__.executemany(sql, seq_of_parameters)
         else:
-            with DatabaseTrace(sql, self._nr_dbapi2_module,
-                    self._nr_connect_params, self._nr_cursor_params, source=self.__wrapped__.executemany):
+            with DatabaseTrace(
+                sql=sql,
+                dbapi2_module=self._nr_dbapi2_module,
+                connect_params=self._nr_connect_params,
+                cursor_params=self._nr_cursor_params,
+                source=self.__wrapped__.executemany,
+            ):
                 return self.__wrapped__.executemany(sql, seq_of_parameters)
 
     def callproc(self, procname, parameters=DEFAULT):
-        with DatabaseTrace('CALL %s' % procname,
-                self._nr_dbapi2_module, self._nr_connect_params, source=self.__wrapped__.callproc):
+        with DatabaseTrace(
+            sql="CALL %s" % procname,
+            dbapi2_module=self._nr_dbapi2_module,
+            connect_params=self._nr_connect_params,
+            source=self.__wrapped__.callproc,
+        ):
             if parameters is not DEFAULT:
                 return self.__wrapped__.callproc(procname, parameters)
             else:
                 return self.__wrapped__.callproc(procname)
 
+
 class ConnectionWrapper(ObjectProxy):
 
     __cursor_wrapper__ = CursorWrapper
@@ -76,20 +100,29 @@ def __init__(self, connection, dbapi2_module, connect_params):
         self._nr_connect_params = connect_params
 
     def cursor(self, *args, **kwargs):
-        return self.__cursor_wrapper__(self.__wrapped__.cursor(
-                *args, **kwargs), self._nr_dbapi2_module,
-                self._nr_connect_params, (args, kwargs))
+        return self.__cursor_wrapper__(
+            self.__wrapped__.cursor(*args, **kwargs), self._nr_dbapi2_module, self._nr_connect_params, (args, kwargs)
+        )
 
     def commit(self):
-        with DatabaseTrace('COMMIT', self._nr_dbapi2_module,
-                self._nr_connect_params, source=self.__wrapped__.commit):
+        with DatabaseTrace(
+            sql="COMMIT",
+            dbapi2_module=self._nr_dbapi2_module,
+            connect_params=self._nr_connect_params,
+            source=self.__wrapped__.commit,
+        ):
             return self.__wrapped__.commit()
 
     def rollback(self):
-        with DatabaseTrace('ROLLBACK', self._nr_dbapi2_module,
-                self._nr_connect_params, source=self.__wrapped__.rollback):
+        with DatabaseTrace(
+            sql="ROLLBACK",
+            dbapi2_module=self._nr_dbapi2_module,
+            connect_params=self._nr_connect_params,
+            source=self.__wrapped__.rollback,
+        ):
             return self.__wrapped__.rollback()
 
+
 class ConnectionFactory(ObjectProxy):
 
     __connection_wrapper__ = ConnectionWrapper
@@ -99,17 +132,15 @@ def __init__(self, connect, dbapi2_module):
         self._nr_dbapi2_module = dbapi2_module
 
     def __call__(self, *args, **kwargs):
-        rollup = []
-        rollup.append('Datastore/all')
-        rollup.append('Datastore/%s/all' %
-                self._nr_dbapi2_module._nr_database_product)
+        rollup = ["Datastore/all", "Datastore/%s/all" % self._nr_dbapi2_module._nr_database_product]
+
+        with FunctionTrace(name=callable_name(self.__wrapped__), terminal=True, rollup=rollup, source=self.__wrapped__):
+            return self.__connection_wrapper__(
+                self.__wrapped__(*args, **kwargs), self._nr_dbapi2_module, (args, kwargs)
+            )
 
-        with FunctionTrace(callable_name(self.__wrapped__),
-                terminal=True, rollup=rollup, source=self.__wrapped__):
-            return self.__connection_wrapper__(self.__wrapped__(
-                    *args, **kwargs), self._nr_dbapi2_module, (args, kwargs))
 
 def instrument(module):
-    register_database_client(module, 'DBAPI2', 'single')
+    register_database_client(module, "DBAPI2", "single")
 
-    wrap_object(module, 'connect', ConnectionFactory, (module,))
+    wrap_object(module, "connect", ConnectionFactory, (module,))