Add google firestore instrumentation (#893)

* Add instrumentation for Google Firestore documents and collections (#876) * Initial GCP firestore instrumentation commit. * Add testing for documents and collections + test generators Co-authored-by: Tim Pansino <[email protected]> Co-authored-by: Lalleh Rafeei <[email protected]> Co-authored-by: Hannah Stepanek <[email protected]> * Add co-authors. Co-authored-by: Tim Pansino <[email protected]> Co-authored-by: Lalleh Rafeei <[email protected]> Co-authored-by: Hannah Stepanek <[email protected]> * Add co-authors. Co-authored-by: Tim Pansino <[email protected]> Co-authored-by: Lalleh Rafeei <[email protected]> Co-authored-by: Hannah Stepanek <[email protected]> * Trim whitespace --------- Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Tim Pansino <[email protected]> Co-authored-by: Lalleh Rafeei <[email protected]> Co-authored-by: Hannah Stepanek <[email protected]> Co-authored-by: Timothy Pansino <[email protected]> * Firestore CI (#877) * Add firestore CI runner * Correct hook file name * Setup emulator credentials * Swap dependency to firestore alone * Hacky setup for firestore * Fix firestore hostname * Ensure firestore connection * Fix CI issues * Refactor Firestore Hooks (#879) * Remove unnecessary instrumentation * Simplify existing instrumentation * Remove unnecessary settings lookups * Firestore Sync Client Instrumentation (#880) * Remove unnecessary instrumentation * Simplify existing instrumentation * Remove unnecessary settings lookups * Client instrumentation * Add query and aggregation query instrumentation * Fix deprecation warning * Simplify collection lookup * Combine query test files * Rename methods for clarity * Instrument Firestore batching * Add transaction instrumentation * Consumer iterators on <=Py38 * Allow better parallelization in firestore tests * Clean out unnecessary code * [Mega-Linter] Apply linters fixes * Better parallelization safeguards * Add collection group instrumentation * [Mega-Linter] Apply linters fixes * Change imports to native APIs * Swap target functions to lambdas * Convert exercise functions to fixtures --------- Co-authored-by: TimPansino <[email protected]> * Update datastore_trace wrapper to take instance info (#883) * Update datastore trace wrapper to take instance info. * [Mega-Linter] Apply linters fixes * Make instance info args optional. * [Mega-Linter] Apply linters fixes * Add datastore trace testing. * Add background task decorator. * [Mega-Linter] Apply linters fixes * Fix typo in validator. --------- Co-authored-by: umaannamalai <[email protected]> * Async Generator Wrapper (#884) * Add async generator wrapper * Add no harm test * Remove anext calls * Add graphql traces to decorator testing * Remove pypy generator gc logic --------- Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> * Trace Async Wrapper Argument (#886) * Add async_wrapper to datastore_trace api * Add async wrapper argument to all trace APIs * Add testing for automatic and manual asyncwrappers * Firstore Async Instrumentation (#882) * Remove unnecessary instrumentation * Simplify existing instrumentation * Remove unnecessary settings lookups * Client instrumentation * Add query and aggregation query instrumentation * Fix deprecation warning * Simplify collection lookup * Combine query test files * Rename methods for clarity * Instrument Firestore batching * Add transaction instrumentation * Consumer iterators on <=Py38 * Add async generator wrapper * Allow better parallelization in firestore tests * Fix issue in async generator wrapper * Add async client instrumentation * Squashed commit of the following: commit 9d411e0 Author: Tim Pansino <[email protected]> Date: Wed Jul 26 15:57:39 2023 -0700 Clean out unnecessary code commit cb550ba Author: Tim Pansino <[email protected]> Date: Wed Jul 26 14:27:01 2023 -0700 Allow better parallelization in firestore tests * Add async collection instrumentation * Add async document instrumentation * Async Query instrumentation * Add async batch instrumentation * Add instrumentation for AsyncTransaction * Squashed commit of the following: commit c836f8f Author: TimPansino <[email protected]> Date: Thu Jul 27 19:54:35 2023 +0000 [Mega-Linter] Apply linters fixes commit 02a55a1 Author: Tim Pansino <[email protected]> Date: Thu Jul 27 12:46:46 2023 -0700 Add collection group instrumentation commit ab1f4ff Author: Tim Pansino <[email protected]> Date: Thu Jul 27 12:00:33 2023 -0700 Better parallelization safeguards commit fa5f39a Author: TimPansino <[email protected]> Date: Wed Jul 26 22:59:11 2023 +0000 [Mega-Linter] Apply linters fixes commit 9d411e0 Author: Tim Pansino <[email protected]> Date: Wed Jul 26 15:57:39 2023 -0700 Clean out unnecessary code commit cb550ba Author: Tim Pansino <[email protected]> Date: Wed Jul 26 14:27:01 2023 -0700 Allow better parallelization in firestore tests * Remove reset_firestore * Re-merge of test_query * Use public API imports * Add async collection group instrumentation * Refactor exercise functions to fixtures * Squashed commit of the following: commit 09c5e11 Author: Tim Pansino <[email protected]> Date: Wed Aug 2 14:33:24 2023 -0700 Add testing for automatic and manual asyncwrappers commit fc3ef6b Author: Tim Pansino <[email protected]> Date: Wed Aug 2 14:33:05 2023 -0700 Add async wrapper argument to all trace APIs commit 479f9e2 Merge: faf3ccc edd1f94 Author: Tim Pansino <[email protected]> Date: Wed Aug 2 13:44:24 2023 -0700 Merge remote-tracking branch 'origin/develop-google-firestore-instrumentation' into feature-async-wrapper-argument commit edd1f94 Author: Timothy Pansino <[email protected]> Date: Wed Aug 2 13:40:51 2023 -0700 Async Generator Wrapper (#884) * Add async generator wrapper * Add no harm test * Remove anext calls * Add graphql traces to decorator testing * Remove pypy generator gc logic --------- Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> commit faf3ccc Author: Tim Pansino <[email protected]> Date: Mon Jul 31 15:10:56 2023 -0700 Add async_wrapper to datastore_trace api * Remove custom wrapper code from firestore * Undo wrapper edits --------- Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> * Firestore Instance Info (#887) * Add instance info testing to query * Instance info for query.stream * Squashed commit of the following: commit 1c426c8 Author: umaannamalai <[email protected]> Date: Mon Jul 31 23:01:49 2023 +0000 [Mega-Linter] Apply linters fixes commit 7687c06 Author: Uma Annamalai <[email protected]> Date: Mon Jul 31 15:47:09 2023 -0700 Make instance info args optional. commit 53f8400 Author: umaannamalai <[email protected]> Date: Mon Jul 31 22:23:20 2023 +0000 [Mega-Linter] Apply linters fixes commit d95d477 Author: Uma Annamalai <[email protected]> Date: Mon Jul 31 15:20:41 2023 -0700 Update datastore trace wrapper to take instance info. * Add instance info testing to all apis * Separate transaction instance info tests * Implement all instance info getters * Squashed commit of the following: commit db3561e Merge: 844e556 edd1f94 Author: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Date: Wed Aug 2 22:10:32 2023 +0000 Merge branch 'develop-google-firestore-instrumentation' into feature-firstore-async-instrumentation commit 844e556 Author: Tim Pansino <[email protected]> Date: Wed Aug 2 15:09:49 2023 -0700 Remove custom wrapper code from firestore commit ad2999f Author: Tim Pansino <[email protected]> Date: Wed Aug 2 14:58:38 2023 -0700 Squashed commit of the following: commit 09c5e11 Author: Tim Pansino <[email protected]> Date: Wed Aug 2 14:33:24 2023 -0700 Add testing for automatic and manual asyncwrappers commit fc3ef6b Author: Tim Pansino <[email protected]> Date: Wed Aug 2 14:33:05 2023 -0700 Add async wrapper argument to all trace APIs commit 479f9e2 Merge: faf3ccc edd1f94 Author: Tim Pansino <[email protected]> Date: Wed Aug 2 13:44:24 2023 -0700 Merge remote-tracking branch 'origin/develop-google-firestore-instrumentation' into feature-async-wrapper-argument commit edd1f94 Author: Timothy Pansino <[email protected]> Date: Wed Aug 2 13:40:51 2023 -0700 Async Generator Wrapper (#884) * Add async generator wrapper * Add no harm test * Remove anext calls * Add graphql traces to decorator testing * Remove pypy generator gc logic --------- Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> commit faf3ccc Author: Tim Pansino <[email protected]> Date: Mon Jul 31 15:10:56 2023 -0700 Add async_wrapper to datastore_trace api commit edd1f94 Author: Timothy Pansino <[email protected]> Date: Wed Aug 2 13:40:51 2023 -0700 Async Generator Wrapper (#884) * Add async generator wrapper * Add no harm test * Remove anext calls * Add graphql traces to decorator testing * Remove pypy generator gc logic --------- Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> commit 29579fc Merge: 4a8a3fe 7596fb4 Author: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Date: Wed Aug 2 19:54:09 2023 +0000 Merge branch 'develop-google-firestore-instrumentation' into feature-firstore-async-instrumentation commit 7596fb4 Author: Uma Annamalai <[email protected]> Date: Wed Aug 2 12:53:29 2023 -0700 Update datastore_trace wrapper to take instance info (#883) * Update datastore trace wrapper to take instance info. * [Mega-Linter] Apply linters fixes * Make instance info args optional. * [Mega-Linter] Apply linters fixes * Add datastore trace testing. * Add background task decorator. * [Mega-Linter] Apply linters fixes * Fix typo in validator. --------- Co-authored-by: umaannamalai <[email protected]> commit 4a8a3fe Merge: 7bf6f49 dcc92a9 Author: Tim Pansino <[email protected]> Date: Mon Jul 31 14:51:20 2023 -0700 Merge remote-tracking branch 'origin/develop-google-firestore-instrumentation' into feature-firstore-async-instrumentation commit 7bf6f49 Author: Tim Pansino <[email protected]> Date: Mon Jul 31 14:34:26 2023 -0700 Refactor exercise functions to fixtures commit d3e4732 Author: Tim Pansino <[email protected]> Date: Thu Jul 27 13:20:37 2023 -0700 Add async collection group instrumentation commit 5902515 Author: Tim Pansino <[email protected]> Date: Thu Jul 27 13:09:13 2023 -0700 Use public API imports commit 9266924 Author: Tim Pansino <[email protected]> Date: Thu Jul 27 13:04:19 2023 -0700 Re-merge of test_query commit b6bc9a4 Author: Tim Pansino <[email protected]> Date: Thu Jul 27 13:01:27 2023 -0700 Remove reset_firestore commit 87fbe62 Author: Tim Pansino <[email protected]> Date: Thu Jul 27 13:00:37 2023 -0700 Squashed commit of the following: commit c836f8f Author: TimPansino <[email protected]> Date: Thu Jul 27 19:54:35 2023 +0000 [Mega-Linter] Apply linters fixes commit 02a55a1 Author: Tim Pansino <[email protected]> Date: Thu Jul 27 12:46:46 2023 -0700 Add collection group instrumentation commit ab1f4ff Author: Tim Pansino <[email protected]> Date: Thu Jul 27 12:00:33 2023 -0700 Better parallelization safeguards commit fa5f39a Author: TimPansino <[email protected]> Date: Wed Jul 26 22:59:11 2023 +0000 [Mega-Linter] Apply linters fixes commit 9d411e0 Author: Tim Pansino <[email protected]> Date: Wed Jul 26 15:57:39 2023 -0700 Clean out unnecessary code commit cb550ba Author: Tim Pansino <[email protected]> Date: Wed Jul 26 14:27:01 2023 -0700 Allow better parallelization in firestore tests commit e04ec6f Author: Tim Pansino <[email protected]> Date: Thu Jul 27 11:55:44 2023 -0700 Add instrumentation for AsyncTransaction commit 6b7fc79 Author: Tim Pansino <[email protected]> Date: Wed Jul 26 16:56:04 2023 -0700 Add async batch instrumentation commit c392e78 Author: Tim Pansino <[email protected]> Date: Wed Jul 26 16:36:03 2023 -0700 Async Query instrumentation commit aab244b Author: Tim Pansino <[email protected]> Date: Wed Jul 26 16:20:58 2023 -0700 Add async document instrumentation commit 3fb6a6c Author: Tim Pansino <[email protected]> Date: Wed Jul 26 16:11:17 2023 -0700 Add async collection instrumentation commit 7851baf Author: Tim Pansino <[email protected]> Date: Wed Jul 26 15:58:12 2023 -0700 Squashed commit of the following: commit 9d411e0 Author: Tim Pansino <[email protected]> Date: Wed Jul 26 15:57:39 2023 -0700 Clean out unnecessary code commit cb550ba Author: Tim Pansino <[email protected]> Date: Wed Jul 26 14:27:01 2023 -0700 Allow better parallelization in firestore tests commit c49a1cf Author: Tim Pansino <[email protected]> Date: Wed Jul 26 15:54:13 2023 -0700 Add async client instrumentation commit c857358 Author: Tim Pansino <[email protected]> Date: Wed Jul 26 15:53:21 2023 -0700 Fix issue in async generator wrapper commit 5693dd2 Author: Tim Pansino <[email protected]> Date: Wed Jul 26 14:27:01 2023 -0700 Allow better parallelization in firestore tests commit fbe40ea Author: Tim Pansino <[email protected]> Date: Wed Jul 26 14:22:53 2023 -0700 Add async generator wrapper commit b9a91e5 Author: Tim Pansino <[email protected]> Date: Wed Jul 26 12:21:25 2023 -0700 Consumer iterators on <=Py38 commit ef06df5 Author: Tim Pansino <[email protected]> Date: Wed Jul 26 12:01:25 2023 -0700 Add transaction instrumentation commit 2ce45c8 Author: Tim Pansino <[email protected]> Date: Tue Jul 25 15:55:50 2023 -0700 Instrument Firestore batching commit d17b62f Author: Tim Pansino <[email protected]> Date: Tue Jul 25 15:31:48 2023 -0700 Rename methods for clarity commit 6214f0b Author: Tim Pansino <[email protected]> Date: Tue Jul 25 15:30:23 2023 -0700 Combine query test files commit b4e8700 Author: Tim Pansino <[email protected]> Date: Tue Jul 25 15:23:03 2023 -0700 Simplify collection lookup commit a0c78a2 Author: Tim Pansino <[email protected]> Date: Tue Jul 25 15:18:51 2023 -0700 Fix deprecation warning commit 44598cc Author: Tim Pansino <[email protected]> Date: Tue Jul 25 15:15:13 2023 -0700 Add query and aggregation query instrumentation commit b9eaa5b Author: Tim Pansino <[email protected]> Date: Tue Jul 25 13:33:42 2023 -0700 Client instrumentation commit 19f5a48 Author: Tim Pansino <[email protected]> Date: Mon Jul 24 15:55:52 2023 -0700 Remove unnecessary settings lookups commit ba7850a Author: Tim Pansino <[email protected]> Date: Mon Jul 24 15:44:54 2023 -0700 Simplify existing instrumentation commit e07ffc3 Author: Tim Pansino <[email protected]> Date: Mon Jul 24 15:44:10 2023 -0700 Remove unnecessary instrumentation * Add instance info to async client * Simplify lookup logic for instance info * Precompute closures for memory usage * Undo wrapper edits * Fix typo * Change port from int ot str * Fix Generator Wrappers (#890) * Fix async wrapper implementations * Add regression testing --------- Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --------- Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Tim Pansino <[email protected]> Co-authored-by: Lalleh Rafeei <[email protected]> Co-authored-by: Hannah Stepanek <[email protected]> Co-authored-by: Timothy Pansino <[email protected]> Co-authored-by: TimPansino <[email protected]> Co-authored-by: umaannamalai <[email protected]>
newrelic · Aug 10, 2023 · 8ebe9a3 · 8ebe9a3
1 parent 17f8937
commit 8ebe9a3
Show file tree

Hide file tree

Showing 34 changed files with 3,281 additions and 76 deletions.
diff --git a/.github/workflows/tests.yml b/.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
diff --git a/newrelic/api/database_trace.py b/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))
diff --git a/newrelic/api/datastore_trace.py b/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)
+    )
diff --git a/newrelic/api/external_trace.py b/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))
diff --git a/newrelic/api/function_trace.py b/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))