Add google firestore instrumentation

.github/workflows/tests.yml

@@ -1024,3 +1024,68 @@ jobs:
           name: coverage-${{ github.job }}-${{ strategy.job-index }}
           path: ./**/.coverage.*
           retention-days: 1
+
+  firestore:
+    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:
+      firestore:
+        # Image set here MUST be repeated down below in options. See comment below.
+        image: gcr.io/google.com/cloudsdktool/google-cloud-cli:437.0.1-emulators
+        ports:
+          - 8080:8080
+        # Set health checks to wait 5 seconds in lieu of an actual healthcheck
+        options: >-
+          --health-cmd "echo success"
+          --health-interval 10s
+          --health-timeout 5s
+          --health-retries 5
+          --health-start-period 5s
+          gcr.io/google.com/cloudsdktool/google-cloud-cli:437.0.1-emulators /bin/bash -c "gcloud emulators firestore start --host-port=0.0.0.0:8080" ||
+        # This is a very hacky solution. GitHub Actions doesn't provide APIs for setting commands on services, but allows adding arbitrary options.
+        # --entrypoint won't work as it only accepts an executable and not the [] syntax.
+        # Instead, we specify the image again the command afterwards like a call to docker create. The result is a few environment variables
+        # and the original command being appended to our hijacked docker create command. We can avoid any issues by adding || to prevent that
+        # from every being executed as bash commands.
+
+    steps:
+      - uses: actions/checkout@v3
+
+      - name: Fetch git tags
+        run: |
+          git config --global --add safe.directory "$GITHUB_WORKSPACE"
+          git fetch --tags origin
+
+      - 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@v3
+        with:
+          name: coverage-${{ github.job }}-${{ strategy.job-index }}
+          path: ./**/.coverage.*
+          retention-days: 1

newrelic/api/database_trace.py

@@ -16,7 +16,7 @@
 import logging
 
 from newrelic.api.time_trace import TimeTrace, current_trace
-from newrelic.common.async_wrapper import async_wrapper
+from newrelic.common.async_wrapper import async_wrapper as get_async_wrapper
 from newrelic.common.object_wrapper import FunctionWrapper, wrap_object
 from newrelic.core.database_node import DatabaseNode
 from newrelic.core.stack_trace import current_stack
@@ -244,9 +244,9 @@ def create_node(self):
         )
 
 
-def DatabaseTraceWrapper(wrapped, sql, dbapi2_module=None):
+def DatabaseTraceWrapper(wrapped, sql, dbapi2_module=None, async_wrapper=None):
     def _nr_database_trace_wrapper_(wrapped, instance, args, kwargs):
-        wrapper = async_wrapper(wrapped)
+        wrapper = async_wrapper if async_wrapper is not None else get_async_wrapper(wrapped)
         if not wrapper:
             parent = current_trace()
             if not parent:
@@ -273,9 +273,9 @@ def _nr_database_trace_wrapper_(wrapped, instance, args, kwargs):
     return FunctionWrapper(wrapped, _nr_database_trace_wrapper_)
 
 
-def database_trace(sql, dbapi2_module=None):
-    return functools.partial(DatabaseTraceWrapper, sql=sql, dbapi2_module=dbapi2_module)
+def database_trace(sql, dbapi2_module=None, async_wrapper=None):
+    return functools.partial(DatabaseTraceWrapper, sql=sql, dbapi2_module=dbapi2_module, async_wrapper=async_wrapper)
 
 
-def wrap_database_trace(module, object_path, sql, dbapi2_module=None):
-    wrap_object(module, object_path, DatabaseTraceWrapper, (sql, dbapi2_module))
+def wrap_database_trace(module, object_path, sql, dbapi2_module=None, async_wrapper=None):
+    wrap_object(module, object_path, DatabaseTraceWrapper, (sql, dbapi2_module, async_wrapper))

newrelic/api/datastore_trace.py

@@ -15,7 +15,7 @@
 import functools
 
 from newrelic.api.time_trace import TimeTrace, current_trace
-from newrelic.common.async_wrapper import async_wrapper
+from newrelic.common.async_wrapper import async_wrapper as get_async_wrapper
 from newrelic.common.object_wrapper import FunctionWrapper, wrap_object
 from newrelic.core.datastore_node import DatastoreNode
 
@@ -82,6 +82,9 @@ def __enter__(self):
             self.product = transaction._intern_string(self.product)
             self.target = transaction._intern_string(self.target)
             self.operation = transaction._intern_string(self.operation)
+            self.host = transaction._intern_string(self.host)
+            self.port_path_or_id = transaction._intern_string(self.port_path_or_id)
+            self.database_name = transaction._intern_string(self.database_name)
 
             datastore_tracer_settings = transaction.settings.datastore_tracer
             self.instance_reporting_enabled = datastore_tracer_settings.instance_reporting.enabled
@@ -92,7 +95,14 @@ def __repr__(self):
         return "<%s object at 0x%x %s>" % (
             self.__class__.__name__,
             id(self),
-            dict(product=self.product, target=self.target, operation=self.operation),
+            dict(
+                product=self.product,
+                target=self.target,
+                operation=self.operation,
+                host=self.host,
+                port_path_or_id=self.port_path_or_id,
+                database_name=self.database_name,
+            ),
         )
 
     def finalize_data(self, transaction, exc=None, value=None, tb=None):
@@ -125,7 +135,7 @@ def create_node(self):
         )
 
 
-def DatastoreTraceWrapper(wrapped, product, target, operation):
+def DatastoreTraceWrapper(wrapped, product, target, operation, host=None, port_path_or_id=None, database_name=None, async_wrapper=None):
     """Wraps a method to time datastore queries.
 
     :param wrapped: The function to apply the trace to.
@@ -140,6 +150,16 @@ def DatastoreTraceWrapper(wrapped, product, target, operation):
                       or the name of any API function/method in the client
                       library.
     :type operation: str or callable
+    :param host: The name of the server hosting the actual datastore.
+    :type host: str
+    :param port_path_or_id: The value passed in can represent either the port,
+                            path, or id of the datastore being connected to.
+    :type port_path_or_id: str
+    :param database_name: The name of database where the current query is being
+                          executed.
+    :type database_name: str
+    :param async_wrapper: An async trace wrapper from newrelic.common.async_wrapper.
+    :type async_wrapper: callable or None
     :rtype: :class:`newrelic.common.object_wrapper.FunctionWrapper`
 
     This is typically used to wrap datastore queries such as calls to Redis or
@@ -155,7 +175,7 @@ def DatastoreTraceWrapper(wrapped, product, target, operation):
     """
 
     def _nr_datastore_trace_wrapper_(wrapped, instance, args, kwargs):
-        wrapper = async_wrapper(wrapped)
+        wrapper = async_wrapper if async_wrapper is not None else get_async_wrapper(wrapped)
         if not wrapper:
             parent = current_trace()
             if not parent:
@@ -187,7 +207,33 @@ def _nr_datastore_trace_wrapper_(wrapped, instance, args, kwargs):
         else:
             _operation = operation
 
-        trace = DatastoreTrace(_product, _target, _operation, parent=parent, source=wrapped)
+        if callable(host):
+            if instance is not None:
+                _host = host(instance, *args, **kwargs)
+            else:
+                _host = host(*args, **kwargs)
+        else:
+            _host = host
+
+        if callable(port_path_or_id):
+            if instance is not None:
+                _port_path_or_id = port_path_or_id(instance, *args, **kwargs)
+            else:
+                _port_path_or_id = port_path_or_id(*args, **kwargs)
+        else:
+            _port_path_or_id = port_path_or_id
+
+        if callable(database_name):
+            if instance is not None:
+                _database_name = database_name(instance, *args, **kwargs)
+            else:
+                _database_name = database_name(*args, **kwargs)
+        else:
+            _database_name = database_name
+
+        trace = DatastoreTrace(
+            _product, _target, _operation, _host, _port_path_or_id, _database_name, parent=parent, source=wrapped
+        )
 
         if wrapper:  # pylint: disable=W0125,W0126
             return wrapper(wrapped, trace)(*args, **kwargs)
@@ -198,7 +244,7 @@ def _nr_datastore_trace_wrapper_(wrapped, instance, args, kwargs):
     return FunctionWrapper(wrapped, _nr_datastore_trace_wrapper_)
 
 
-def datastore_trace(product, target, operation):
+def datastore_trace(product, target, operation, host=None, port_path_or_id=None, database_name=None, async_wrapper=None):
     """Decorator allows datastore query to be timed.
 
     :param product: The name of the vendor.
@@ -211,6 +257,16 @@ def datastore_trace(product, target, operation):
                       or the name of any API function/method in the client
                       library.
     :type operation: str
+    :param host: The name of the server hosting the actual datastore.
+    :type host: str
+    :param port_path_or_id: The value passed in can represent either the port,
+                            path, or id of the datastore being connected to.
+    :type port_path_or_id: str
+    :param database_name: The name of database where the current query is being
+                          executed.
+    :type database_name: str
+    :param async_wrapper: An async trace wrapper from newrelic.common.async_wrapper.
+    :type async_wrapper: callable or None
 
     This is typically used to decorate datastore queries such as calls to Redis
     or ElasticSearch.
@@ -224,10 +280,21 @@ def datastore_trace(product, target, operation):
         ...     time.sleep(*args, **kwargs)
 
     """
-    return functools.partial(DatastoreTraceWrapper, product=product, target=target, operation=operation)
-
-
-def wrap_datastore_trace(module, object_path, product, target, operation):
+    return functools.partial(
+        DatastoreTraceWrapper,
+        product=product,
+        target=target,
+        operation=operation,
+        host=host,
+        port_path_or_id=port_path_or_id,
+        database_name=database_name,
+        async_wrapper=async_wrapper,
+    )
+
+
+def wrap_datastore_trace(
+    module, object_path, product, target, operation, host=None, port_path_or_id=None, database_name=None, async_wrapper=None
+):
     """Method applies custom timing to datastore query.
 
     :param module: Module containing the method to be instrumented.
@@ -244,6 +311,16 @@ def wrap_datastore_trace(module, object_path, product, target, operation):
                       or the name of any API function/method in the client
                       library.
     :type operation: str
+    :param host: The name of the server hosting the actual datastore.
+    :type host: str
+    :param port_path_or_id: The value passed in can represent either the port,
+                            path, or id of the datastore being connected to.
+    :type port_path_or_id: str
+    :param database_name: The name of database where the current query is being
+                          executed.
+    :type database_name: str
+    :param async_wrapper: An async trace wrapper from newrelic.common.async_wrapper.
+    :type async_wrapper: callable or None
 
     This is typically used to time database query method calls such as Redis
     GET.
@@ -256,4 +333,6 @@ def wrap_datastore_trace(module, object_path, product, target, operation):
         ...        'sleep')
 
     """
-    wrap_object(module, object_path, DatastoreTraceWrapper, (product, target, operation))
+    wrap_object(
+        module, object_path, DatastoreTraceWrapper, (product, target, operation, host, port_path_or_id, database_name, async_wrapper)
+    )

newrelic/api/external_trace.py

@@ -16,7 +16,7 @@
 
 from newrelic.api.cat_header_mixin import CatHeaderMixin
 from newrelic.api.time_trace import TimeTrace, current_trace
-from newrelic.common.async_wrapper import async_wrapper
+from newrelic.common.async_wrapper import async_wrapper as get_async_wrapper
 from newrelic.common.object_wrapper import FunctionWrapper, wrap_object
 from newrelic.core.external_node import ExternalNode
 
@@ -66,9 +66,9 @@ def create_node(self):
         )
 
 
-def ExternalTraceWrapper(wrapped, library, url, method=None):
+def ExternalTraceWrapper(wrapped, library, url, method=None, async_wrapper=None):
     def dynamic_wrapper(wrapped, instance, args, kwargs):
-        wrapper = async_wrapper(wrapped)
+        wrapper = async_wrapper if async_wrapper is not None else get_async_wrapper(wrapped)
         if not wrapper:
             parent = current_trace()
             if not parent:
@@ -103,7 +103,7 @@ def dynamic_wrapper(wrapped, instance, args, kwargs):
             return wrapped(*args, **kwargs)
 
     def literal_wrapper(wrapped, instance, args, kwargs):
-        wrapper = async_wrapper(wrapped)
+        wrapper = async_wrapper if async_wrapper is not None else get_async_wrapper(wrapped)
         if not wrapper:
             parent = current_trace()
             if not parent:
@@ -125,9 +125,9 @@ def literal_wrapper(wrapped, instance, args, kwargs):
     return FunctionWrapper(wrapped, literal_wrapper)
 
 
-def external_trace(library, url, method=None):
-    return functools.partial(ExternalTraceWrapper, library=library, url=url, method=method)
+def external_trace(library, url, method=None, async_wrapper=None):
+    return functools.partial(ExternalTraceWrapper, library=library, url=url, method=method, async_wrapper=async_wrapper)
 
 
-def wrap_external_trace(module, object_path, library, url, method=None):
-    wrap_object(module, object_path, ExternalTraceWrapper, (library, url, method))
+def wrap_external_trace(module, object_path, library, url, method=None, async_wrapper=None):
+    wrap_object(module, object_path, ExternalTraceWrapper, (library, url, method, async_wrapper))

newrelic/api/function_trace.py

@@ -15,7 +15,7 @@
 import functools
 
 from newrelic.api.time_trace import TimeTrace, current_trace
-from newrelic.common.async_wrapper import async_wrapper
+from newrelic.common.async_wrapper import async_wrapper as get_async_wrapper
 from newrelic.common.object_names import callable_name
 from newrelic.common.object_wrapper import FunctionWrapper, wrap_object
 from newrelic.core.function_node import FunctionNode
@@ -89,9 +89,9 @@ def create_node(self):
         )
 
 
-def FunctionTraceWrapper(wrapped, name=None, group=None, label=None, params=None, terminal=False, rollup=None):
+def FunctionTraceWrapper(wrapped, name=None, group=None, label=None, params=None, terminal=False, rollup=None, async_wrapper=None):
     def dynamic_wrapper(wrapped, instance, args, kwargs):
-        wrapper = async_wrapper(wrapped)
+        wrapper = async_wrapper if async_wrapper is not None else get_async_wrapper(wrapped)
         if not wrapper:
             parent = current_trace()
             if not parent:
@@ -147,7 +147,7 @@ def dynamic_wrapper(wrapped, instance, args, kwargs):
             return wrapped(*args, **kwargs)
 
     def literal_wrapper(wrapped, instance, args, kwargs):
-        wrapper = async_wrapper(wrapped)
+        wrapper = async_wrapper if async_wrapper is not None else get_async_wrapper(wrapped)
         if not wrapper:
             parent = current_trace()
             if not parent:
@@ -171,13 +171,13 @@ def literal_wrapper(wrapped, instance, args, kwargs):
     return FunctionWrapper(wrapped, literal_wrapper)
 
 
-def function_trace(name=None, group=None, label=None, params=None, terminal=False, rollup=None):
+def function_trace(name=None, group=None, label=None, params=None, terminal=False, rollup=None, async_wrapper=None):
     return functools.partial(
-        FunctionTraceWrapper, name=name, group=group, label=label, params=params, terminal=terminal, rollup=rollup
+        FunctionTraceWrapper, name=name, group=group, label=label, params=params, terminal=terminal, rollup=rollup, async_wrapper=async_wrapper
     )
 
 
 def wrap_function_trace(
-    module, object_path, name=None, group=None, label=None, params=None, terminal=False, rollup=None
+    module, object_path, name=None, group=None, label=None, params=None, terminal=False, rollup=None, async_wrapper=None
 ):
-    return wrap_object(module, object_path, FunctionTraceWrapper, (name, group, label, params, terminal, rollup))
+    return wrap_object(module, object_path, FunctionTraceWrapper, (name, group, label, params, terminal, rollup, async_wrapper))