Add tracking headers

Add Archivist-User-Agent and Archivist-Partner-ID headers. Fixes AB#9572
datatrails · Jun 7, 2024 · a2006c1 · a2006c1
1 parent 683f62a
commit a2006c1
Show file tree

Hide file tree

Showing 44 changed files with 475 additions and 49 deletions.
diff --git a/archivist/archivist.py b/archivist/archivist.py
@@ -108,11 +108,13 @@ def __init__(
         fixtures: "dict[str,dict[Any,Any]]|None" = None,
         verify: bool = True,
         max_time: float = MAX_TIME,
+        partner_id: str = "",
     ):
         super().__init__(
             fixtures=fixtures,
             verify=verify,
             max_time=max_time,
+            partner_id=partner_id,
         )
 
         if isinstance(auth, tuple):
@@ -214,10 +216,11 @@ def __copy__(self) -> "Archivist":
             fixtures=deepcopy(self._fixtures),
             verify=self._verify,
             max_time=self._max_time,
+            partner_id=self._partner_id,
         )
 
     def _add_headers(self, headers: "dict[str,str]|None") -> "dict[str,Any]":
-        newheaders = {**headers} if isinstance(headers, dict) else {}
+        newheaders = super()._add_headers(headers)
 
         auth = self.auth  # this may trigger a refetch so only do it once here
         # for appidp endpoint there may not be an authtoken

diff --git a/archivist/archivistpublic.py b/archivist/archivistpublic.py
@@ -33,12 +33,15 @@
     from requests.models import Response
 
 
+from .about import __version__ as VERSION
 from .assetattachments import _AssetAttachmentsClient
 from .assets import _AssetsPublic
 from .confirmer import MAX_TIME
 from .constants import (
     HEADERS_REQUEST_TOTAL_COUNT,
     HEADERS_TOTAL_COUNT,
+    PARTNER_ID,
+    USER_AGENT,
 )
 from .dictmerge import _deepmerge, _dotdict
 from .errors import (
@@ -82,12 +85,14 @@ def __init__(
         fixtures: "dict[str, Any]|None" = None,
         verify: bool = True,
         max_time: float = MAX_TIME,
+        partner_id: str = "",
     ):
         self._verify = verify
         self._response_ring_buffer = deque(maxlen=self.RING_BUFFER_MAX_LEN)
         self._session = None
         self._max_time = max_time
         self._fixtures = fixtures or {}
+        self._partner_id = partner_id
 
         # Type hints for IDE autocomplete, keep in sync with CLIENTS map above
         self.assets: _AssetsPublic
@@ -151,6 +156,16 @@ def max_time(self) -> float:
         """bool: Returns maximum time in seconds to wait for confirmation"""
         return self._max_time
 
+    @property
+    def version(self) -> str:
+        """str: Returns version of the archivist package"""
+        return VERSION
+
+    @property
+    def partner_id(self) -> str:
+        """str: Returns partner id of the archivist package"""
+        return self._partner_id
+
     @property
     def fixtures(self) -> "dict[str, Any]":
         """dict: Contains predefined attributes for each endpoint"""
@@ -170,6 +185,10 @@ def __copy__(self):
 
     def _add_headers(self, headers: "dict[str, str]|None") -> "dict[str, str]":
         newheaders = {**headers} if headers is not None else {}
+        newheaders[USER_AGENT] = f"datatrails/pysdk/{self.version}"
+        p = self.partner_id
+        if p:
+            newheaders[PARTNER_ID] = p
 
         return newheaders
 

diff --git a/archivist/constants.py b/archivist/constants.py
@@ -13,6 +13,10 @@
     "RecordEvidence",
 ]
 
+USER_AGENT = "Archivist-User-Agent"
+USER_AGENT_PREFIX = "datatrails/pysdk/"
+PARTNER_ID = "Archivist-Partner-ID"
+
 # define in MIME canonical form
 HEADERS_REQUEST_TOTAL_COUNT = "X-Request-Total-Count"
 HEADERS_TOTAL_COUNT = "X-Total-Count"

diff --git a/archivist/parser.py b/archivist/parser.py
@@ -129,6 +129,15 @@ def common_parser(description: str):
         default=None,
         help="namespace of item population",
     )
+    parser.add_argument(
+        "-p",
+        "--partner_id",
+        type=str,
+        dest="partner_id",
+        action="store",
+        default=None,
+        help="partner id",
+    )
 
     return parser
 
@@ -171,7 +180,12 @@ def endpoint(args):
         LOGGER.error("Critical error.  Aborting.")
         sys_exit(1)
 
-    arch = Archivist(args.url, auth, fixtures=fixtures)
+    arch = Archivist(
+            args.url,
+            auth,
+            fixtures=fixtures,
+            partner_id=args.partner_id,
+    )
     if arch is None:
         LOGGER.error("Critical error.  Aborting.")
         sys_exit(1)

diff --git a/examples/create_asset.py b/examples/create_asset.py
@@ -75,7 +75,7 @@ def main():
     create an example archivist connection and create an asset.
 
     """
-    # optional call to set the logger level for all subsystems. The argumant can
+    # optional call to set the logger level for all subsystems. The argument can
     # be either "INFO" or "DEBUG". For more sophisticated logging control see the
     # documentation.
     set_logger("INFO")
@@ -97,10 +97,13 @@ def main():
     # Initialize connection to Archivist. max_time is the time to wait for confirmation
     # of an asset or event creation - the default is 300 seconds but one can optionally
     # specify a different value.
+    # The optional partner id field is allocated by Datatrails to partners - partners are then
+    # expected to specify this value when submitting any request to the archivist product.
     with Archivist(
         "https://app.datatrails.ai",
         auth,
         max_time=300,
+        partner_id="acme/f7a6beef-f01c-4b39-a494-3fa6b45d6bf4",
     ) as arch:
         # Create a new asset
         asset = create_asset(arch)

diff --git a/functests/execaccess_policies.py b/functests/execaccess_policies.py
@@ -95,6 +95,7 @@
         },
     ],
 }
+PARTNER_IDENTIFIER = "datatrails/1234567890"
 
 if getenv("DATATRAILS_LOGLEVEL") is not None:
     logger.set_logger(getenv("DATATRAILS_LOGLEVEL"))
@@ -117,7 +118,11 @@ def setUp(self):
             client_secret=getenv("DATATRAILS_APPREG_SECRET"),
             client_secret_filename=getenv("DATATRAILS_APPREG_SECRET_FILENAME"),
         )
-        self.arch = Archivist(getenv("DATATRAILS_URL"), auth)
+        self.arch = Archivist(
+            getenv("DATATRAILS_URL"),
+            auth,
+            partner_id=PARTNER_IDENTIFIER,
+        )
 
         # these are for access_policies
         self.ac_props = deepcopy(PROPS)
@@ -298,8 +303,11 @@ def setUp(self):
         super().setUp()
         with open(getenv("DATATRAILS_AUTHTOKEN_FILENAME_2"), encoding="utf-8") as fd:
             auth_2 = fd.read().strip()
-        self.arch_2 = Archivist(getenv("DATATRAILS_URL"), auth_2)
-
+        self.arch_2 = Archivist(
+            getenv("DATATRAILS_URL"),
+            auth_2,
+            partner_id=PARTNER_IDENTIFIER,
+        )
         # creates reciprocal subjects for arch 1 and arch 2.
         # subject 1 contains details of subject 2 to be shared
         self.subject_1, self.subject_2 = self.arch.subjects.share(

diff --git a/functests/execapplications.py b/functests/execapplications.py
@@ -32,6 +32,7 @@
     "arc_description": "Traffic flow control light at A603 North East",
     "some_custom_attribute": "value",
 }
+PARTNER_IDENTIFIER = "datatrails/1234567890"
 
 if getenv("DATATRAILS_LOGLEVEL") is not None:
     logger.set_logger(getenv("DATATRAILS_LOGLEVEL"))
@@ -54,7 +55,11 @@ def setUp(self):
             client_secret=getenv("DATATRAILS_APPREG_SECRET"),
             client_secret_filename=getenv("DATATRAILS_APPREG_SECRET_FILENAME"),
         )
-        self.arch = Archivist(getenv("DATATRAILS_URL"), auth)
+        self.arch = Archivist(
+            getenv("DATATRAILS_URL"),
+            auth,
+            partner_id=PARTNER_IDENTIFIER,
+        )
         self.display_name = f"{DISPLAY_NAME} {uuid4()}"
 
     def tearDown(self):

diff --git a/functests/execassets.py b/functests/execassets.py
@@ -70,6 +70,7 @@
         },
     ],
 }
+PARTNER_IDENTIFIER = "datatrails/1234567890"
 
 
 class TestAssetCreate(TestCase):
@@ -87,7 +88,12 @@ def setUp(self):
             client_secret=getenv("DATATRAILS_APPREG_SECRET"),
             client_secret_filename=getenv("DATATRAILS_APPREG_SECRET_FILENAME"),
         )
-        self.arch = Archivist(getenv("DATATRAILS_URL"), auth, max_time=30)
+        self.arch = Archivist(
+            getenv("DATATRAILS_URL"),
+            auth,
+            max_time=30,
+            partner_id=PARTNER_IDENTIFIER,
+        )
         self.attrs = deepcopy(ATTRS)
         self.traffic_light = deepcopy(ATTRS)
         self.traffic_light["arc_display_type"] = "Traffic light with violation camera"
@@ -233,7 +239,12 @@ def setUp(self):
             client_secret=getenv("DATATRAILS_APPREG_SECRET"),
             client_secret_filename=getenv("DATATRAILS_APPREG_SECRET_FILENAME"),
         )
-        self.arch = Archivist(getenv("DATATRAILS_URL"), auth, max_time=600)
+        self.arch = Archivist(
+            getenv("DATATRAILS_URL"),
+            auth,
+            max_time=30,
+            partner_id=PARTNER_IDENTIFIER,
+        )
 
     def tearDown(self):
         self.arch.close()

diff --git a/functests/execattachments.py b/functests/execattachments.py
@@ -16,6 +16,8 @@
 
 from .constants import TestCase
 
+PARTNER_IDENTIFIER = "datatrails/1234567890"
+
 if getenv("DATATRAILS_LOGLEVEL") is not None:
     logger.set_logger(getenv("DATATRAILS_LOGLEVEL"))
 
@@ -46,7 +48,11 @@ def setUp(self):
             client_secret=getenv("DATATRAILS_APPREG_SECRET"),
             client_secret_filename=getenv("DATATRAILS_APPREG_SECRET_FILENAME"),
         )
-        self.arch = Archivist(getenv("DATATRAILS_URL"), auth)
+        self.arch = Archivist(
+            getenv("DATATRAILS_URL"),
+            auth,
+            partner_id=PARTNER_IDENTIFIER,
+        )
         self.file_uuid: str = ""
 
         with suppress(FileNotFoundError):
@@ -208,7 +214,11 @@ def setUp(self):
             client_secret=getenv("DATATRAILS_APPREG_SECRET"),
             client_secret_filename=getenv("DATATRAILS_APPREG_SECRET_FILENAME"),
         )
-        self.arch = Archivist(getenv("DATATRAILS_URL"), auth)
+        self.arch = Archivist(
+            getenv("DATATRAILS_URL"),
+            auth,
+            partner_id=PARTNER_IDENTIFIER,
+        )
 
     def tearDown(self):
         self.arch.close()

diff --git a/functests/execcompliance_policies.py b/functests/execcompliance_policies.py
@@ -82,6 +82,7 @@
         ["rad<7"],
     ],
 )
+PARTNER_IDENTIFIER = "datatrails/1234567890"
 
 
 class TestCompliancePoliciesBase(TestCase):
@@ -99,7 +100,11 @@ def setUp(self):
             client_secret=getenv("DATATRAILS_APPREG_SECRET"),
             client_secret_filename=getenv("DATATRAILS_APPREG_SECRET_FILENAME"),
         )
-        self.arch = Archivist(getenv("DATATRAILS_URL"), auth)
+        self.arch = Archivist(
+            getenv("DATATRAILS_URL"),
+            auth,
+            partner_id=PARTNER_IDENTIFIER,
+        )
         self.identities = []
 
     def tearDown(self):

diff --git a/functests/execnotebooks.py b/functests/execnotebooks.py
@@ -16,6 +16,8 @@
 
 from .constants import TestCase
 
+PARTNER_IDENTIFIER = "datatrails/1234567890"
+
 if getenv("DATATRAILS_LOGLEVEL") is not None:
     logger.set_logger(getenv("DATATRAILS_LOGLEVEL"))
 
@@ -37,7 +39,11 @@ def setUp(self):
         self.client_id = getenv("DATATRAILS_APPREG_CLIENT")
         self.client_secret = getenv("DATATRAILS_APPREG_SECRET")
         self.url = getenv("DATATRAILS_URL")
-        self.arch = Archivist(self.url, (self.client_id, self.client_secret))
+        self.arch = Archivist(
+            self.url,
+            (self.client_id, self.client_secret),
+            partner_id=PARTNER_IDENTIFIER,
+        )
 
     def tearDown(self):
         self.arch.close()

diff --git a/functests/execpublicassets.py b/functests/execpublicassets.py
@@ -66,6 +66,7 @@
     ],
     "public": True,
 }
+PARTNER_IDENTIFIER = "datatrails/1234567890"
 
 
 class TestPublicAssetCreate(TestCase):
@@ -84,7 +85,12 @@ def setUp(self):
             client_secret_filename=getenv("DATATRAILS_APPREG_SECRET_FILENAME"),
         )
         self.url = getenv("DATATRAILS_URL")
-        self.arch = Archivist(self.url, auth, max_time=600)
+        self.arch = Archivist(
+            getenv("DATATRAILS_URL"),
+            auth,
+            max_time=30,
+            partner_id=PARTNER_IDENTIFIER,
+        )
         self.attrs = deepcopy(ATTRS)
         self.traffic_light = deepcopy(ATTRS)
         self.traffic_light["arc_display_type"] = "Traffic light with violation camera"

diff --git a/functests/execrunner.py b/functests/execrunner.py
@@ -18,6 +18,7 @@
 # pylint: disable=missing-docstring
 # pylint: disable=unused-variable
 
+PARTNER_IDENTIFIER = "datatrails/1234567890"
 
 if getenv("DATATRAILS_LOGLEVEL") is not None:
     logger.set_logger(getenv("DATATRAILS_LOGLEVEL"))
@@ -40,7 +41,12 @@ def setUp(self):
             client_secret=getenv("DATATRAILS_APPREG_SECRET"),
             client_secret_filename=getenv("DATATRAILS_APPREG_SECRET_FILENAME"),
         )
-        self.arch = Archivist(getenv("DATATRAILS_URL"), auth, max_time=300)
+        self.arch = Archivist(
+            getenv("DATATRAILS_URL"),
+            auth,
+            max_time=30,
+            partner_id=PARTNER_IDENTIFIER,
+        )
 
     def tearDown(self):
         self.arch.close()

diff --git a/functests/execsubjects.py b/functests/execsubjects.py
@@ -44,6 +44,7 @@
     "0aW9uX3N0YXR1cyI6ICJDT05GSVJNQVRJT05fU1RBVFVTX1VOU1BFQ0lGSU"
     "VEIn0="
 )
+PARTNER_IDENTIFIER = "datatrails/1234567890"
 
 
 class TestSubjects(TestCase):
@@ -61,7 +62,11 @@ def setUp(self):
             client_secret=getenv("DATATRAILS_APPREG_SECRET"),
             client_secret_filename=getenv("DATATRAILS_APPREG_SECRET_FILENAME"),
         )
-        self.arch = Archivist(getenv("DATATRAILS_URL"), auth)
+        self.arch = Archivist(
+            getenv("DATATRAILS_URL"),
+            auth,
+            partner_id=PARTNER_IDENTIFIER,
+        )
         self.display_name = f"{DISPLAY_NAME} {uuid4()}"
 
     def tearDown(self):

diff --git a/pyproject.toml b/pyproject.toml
@@ -259,7 +259,7 @@ valid-metaclass-classmethod-first-arg = ["cls"]
 # ignored-parents =
 
 # Maximum number of arguments for function / method.
-max-args = 6
+max-args = 7
 
 # Maximum number of attributes for a class (see R0902).
 max-attributes = 7