From c50c5994905ca8c16d45de47c712c0cd6c15c7dc Mon Sep 17 00:00:00 2001
From: SageGJ <gianna.jordan@sagebase.org>
Date: Tue, 28 Jan 2025 12:35:23 -0800
Subject: [PATCH] [Schematic-90] Use synapseclient user-agent string to track
 schematic library and CLI usage (#1569)

set synapseclient user strings to track schematic cli and library usage
---
 schematic/__init__.py                | 12 +++++++++-
 schematic/__main__.py                |  7 +++++-
 tests/unit/test_metric_collection.py | 36 ++++++++++++++++++++++++++++
 3 files changed, 53 insertions(+), 2 deletions(-)
 create mode 100644 tests/unit/test_metric_collection.py

diff --git a/schematic/__init__.py b/schematic/__init__.py
index 18b1216c0..6cc29fb5c 100644
--- a/schematic/__init__.py
+++ b/schematic/__init__.py
@@ -21,7 +21,7 @@
 from opentelemetry.trace import Span, SpanContext, get_current_span
 from opentelemetry.sdk.trace.export import BatchSpanProcessor, Span
 from opentelemetry.sdk.trace.sampling import ALWAYS_OFF
-from synapseclient import Synapse
+from synapseclient import Synapse, USER_AGENT
 from werkzeug import Request
 
 from schematic.configuration.configuration import CONFIG
@@ -35,6 +35,16 @@
 # Ensure environment variables are loaded
 load_dotenv()
 
+USER_AGENT_LIBRARY = {
+    "User-Agent": USER_AGENT["User-Agent"] + f" schematic/{__version__}"
+}
+
+USER_AGENT_COMMAND_LINE = {
+    "User-Agent": USER_AGENT["User-Agent"] + f" schematiccommandline/{__version__}"
+}
+
+USER_AGENT |= USER_AGENT_LIBRARY
+
 
 class AttributePropagatingSpanProcessor(SpanProcessor):
     """A custom span processor that propagates specific attributes from the parent span
diff --git a/schematic/__main__.py b/schematic/__main__.py
index ab7a3fc76..973ede7c2 100644
--- a/schematic/__main__.py
+++ b/schematic/__main__.py
@@ -14,7 +14,8 @@
 from schematic.visualization.commands import (
     viz as viz_cli,
 )  # viz generation commands
-from schematic import __version__
+from schematic import __version__, USER_AGENT_COMMAND_LINE
+
 
 logger = logging.getLogger()
 click_log.basic_config(logger)
@@ -31,6 +32,10 @@ def main():
     """
     Command line interface to the `schematic` backend services.
     """
+    from synapseclient import USER_AGENT
+
+    USER_AGENT |= USER_AGENT_COMMAND_LINE
+
     logger.info("Starting schematic...")
     logger.debug("Existing sub-commands need to be used with schematic.")
 
diff --git a/tests/unit/test_metric_collection.py b/tests/unit/test_metric_collection.py
new file mode 100644
index 000000000..e995dba7f
--- /dev/null
+++ b/tests/unit/test_metric_collection.py
@@ -0,0 +1,36 @@
+from re import search
+from click.testing import CliRunner
+import pytest
+
+
+@pytest.fixture
+def command_line_user_agent_pattern():
+    yield "schematiccommandline" + "/(\\S+)"
+
+
+@pytest.fixture
+def library_user_agent_pattern():
+    yield "schematic" + "/(\\S+)"
+
+
+class TestUserAgentString:
+    def test_user_agent_string(
+        self,
+        library_user_agent_pattern,
+        command_line_user_agent_pattern,
+    ):
+        # GIVEN the default USER_AGENT string from the synapse client
+        from synapseclient import USER_AGENT
+
+        # WHEN schematic is imported to be used as a library
+        from schematic.__main__ import main
+
+        # THEN the User-Agent string should be updated to include the schematic library client string
+        assert search(library_user_agent_pattern, USER_AGENT["User-Agent"])
+
+        # AND when the command line is used to execute commands
+        runner = CliRunner()
+        result = runner.invoke(main)
+
+        # THEN the User-Agent string should be updated to include the schematic command line client string
+        assert search(command_line_user_agent_pattern, USER_AGENT["User-Agent"])