From 10bcca6964ab2af6df8b6abeae8153002e4cc76d Mon Sep 17 00:00:00 2001 From: Peter Zhu Date: Tue, 11 Feb 2025 14:29:58 -0800 Subject: [PATCH 1/2] feat: Add new cookie format --- src/amplitude_experiment/cookie.py | 46 +++++++++++++++++++++++++++--- tests/cookie_test.py | 27 ++++++++++++++++++ 2 files changed, 69 insertions(+), 4 deletions(-) diff --git a/src/amplitude_experiment/cookie.py b/src/amplitude_experiment/cookie.py index 31cd0f8..351b58e 100644 --- a/src/amplitude_experiment/cookie.py +++ b/src/amplitude_experiment/cookie.py @@ -1,3 +1,7 @@ +import json +import logging +import urllib.parse + from .user import User import base64 @@ -6,29 +10,52 @@ class AmplitudeCookie: """This class provides utility functions for parsing and handling identity from Amplitude cookies.""" @staticmethod - def cookie_name(api_key: str) -> str: + def cookie_name(api_key: str, new_format: bool = False) -> str: """ Get the cookie name that Amplitude sets for the provided Parameters: api_key (str): The Amplitude API Key + new_format (bool): True if cookie is in Browser SDK 2.0 format Returns: The cookie name that Amplitude sets for the provided Amplitude API Key """ if not api_key: raise ValueError("Invalid Amplitude API Key") + + if new_format: + if len(api_key) < 10: + raise ValueError("Invalid Amplitude API Key") + return f"AMP_{api_key[0:10]}" + + if len(api_key) < 6: + raise ValueError("Invalid Amplitude API Key") return f"amp_{api_key[0:6]}" @staticmethod - def parse(amplitude_cookie: str) -> User: + def parse(amplitude_cookie: str, new_format: bool = False) -> User: """ Parse a cookie string and returns user Parameters: amplitude_cookie (str): A string from the amplitude cookie + new_format: True if cookie is in Browser SDK 2.0 format Returns: Experiment User context containing a device_id and user_id (if available) """ + + if new_format: + decoding = base64.b64decode(amplitude_cookie).decode("utf-8") + try: + user_session = json.loads(urllib.parse.unquote_plus(decoding)) + if "userId" not in user_session: + return User(user_id=None, device_id=user_session["deviceId"]) + return User(user_id=user_session["userId"], device_id=user_session["deviceId"]) + except Exception as e: + logger = logging.getLogger("Amplitude") + logger.error("Error parsing the Amplitude cookie: " + str(e)) + return User() + values = amplitude_cookie.split('.') user_id = None if values[1]: @@ -39,13 +66,24 @@ def parse(amplitude_cookie: str) -> User: return User(user_id=user_id, device_id=values[0]) @staticmethod - def generate(device_id: str) -> str: + def generate(device_id: str, new_format: bool = False) -> str: """ Generates a cookie string to set for the Amplitude Javascript SDK Parameters: device_id (str): A device id to set + new_format: True if cookie is in Browser SDK 2.0 format Returns: A cookie string to set for the Amplitude Javascript SDK to read """ - return f"{device_id}.........." + if not new_format: + return f"{device_id}.........." + + user_session_hash = { + "deviceId": device_id + } + json_data = json.dumps(user_session_hash) + encoded_json = urllib.parse.quote(json_data) + base64_encoded = base64.b64encode(bytearray(encoded_json, "utf-8")).decode("utf-8") + + return base64_encoded diff --git a/tests/cookie_test.py b/tests/cookie_test.py index 9964a2e..fda0639 100644 --- a/tests/cookie_test.py +++ b/tests/cookie_test.py @@ -30,6 +30,33 @@ def test_parse_cookie_with_device_id_and_utf_user_id(self): def test_generate(self): self.assertEqual(AmplitudeCookie.generate('deviceId'), 'deviceId..........') + def test_new_format_valid_api_key_return_cookie_name(self): + self.assertEqual(AmplitudeCookie.cookie_name('1234567'), 'amp_123456') + + def test_new_format_invalid_api_key_raise_error(self): + self.assertRaises(ValueError, AmplitudeCookie.cookie_name, '') + + def test_new_format_parse_cookie_with_device_id_only(self): + user = AmplitudeCookie.parse('JTdCJTIyZGV2aWNlSWQlMjIlM0ElMjJzb21lRGV2aWNlSWQlMjIlN0Q=', new_format=True) + self.assertIsNotNone(user) + self.assertEqual(user.device_id, 'someDeviceId') + self.assertIsNone(user.user_id) + + def test_new_format_parse_cookie_with_device_id_and_user_id(self): + user = AmplitudeCookie.parse('JTdCJTIyZGV2aWNlSWQlMjIlM0ElMjJzb21lRGV2aWNlSWQlMjIlMkMlMjJ1c2VySWQlMjIlM0ElMjJleGFtcGxlJTQwYW1wbGl0dWRlLmNvbSUyMiU3RA==', new_format=True) + self.assertIsNotNone(user) + self.assertEqual(user.device_id, 'someDeviceId') + self.assertEqual(user.user_id, 'example@amplitude.com') + + def test_new_format_parse_cookie_with_device_id_and_utf_user_id(self): + user = AmplitudeCookie.parse('JTdCJTIyZGV2aWNlSWQlMjIlM0ElMjJzb21lRGV2aWNlSWQlMjIlMkMlMjJ1c2VySWQlMjIlM0ElMjJjJUMzJUI3JTNFJTIyJTdE', new_format=True) + self.assertIsNotNone(user) + self.assertEqual(user.device_id, 'someDeviceId') + self.assertEqual(user.user_id, 'c÷>') + + def test_new_format_generate(self): + self.assertEqual(AmplitudeCookie.generate('someDeviceId', new_format=True), 'JTdCJTIyZGV2aWNlSWQlMjIlM0ElMjAlMjJzb21lRGV2aWNlSWQlMjIlN0Q=') + if __name__ == '__main__': unittest.main() From f0726737c0d4a571d2ffca07b2f8cdf0661db2ac Mon Sep 17 00:00:00 2001 From: Peter Zhu Date: Wed, 12 Feb 2025 13:59:53 -0800 Subject: [PATCH 2/2] Fix new format test --- tests/cookie_test.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/cookie_test.py b/tests/cookie_test.py index fda0639..7c90c7d 100644 --- a/tests/cookie_test.py +++ b/tests/cookie_test.py @@ -31,10 +31,10 @@ def test_generate(self): self.assertEqual(AmplitudeCookie.generate('deviceId'), 'deviceId..........') def test_new_format_valid_api_key_return_cookie_name(self): - self.assertEqual(AmplitudeCookie.cookie_name('1234567'), 'amp_123456') + self.assertEqual(AmplitudeCookie.cookie_name('12345678901', new_format=True), 'AMP_1234567890') def test_new_format_invalid_api_key_raise_error(self): - self.assertRaises(ValueError, AmplitudeCookie.cookie_name, '') + self.assertRaises(ValueError, AmplitudeCookie.cookie_name, '12345678', new_format=True) def test_new_format_parse_cookie_with_device_id_only(self): user = AmplitudeCookie.parse('JTdCJTIyZGV2aWNlSWQlMjIlM0ElMjJzb21lRGV2aWNlSWQlMjIlN0Q=', new_format=True)