def __init__(self,
api_key: str,
on_error: Callable[[str], None],
flush_queue_size: int,
flush_interval: timedelta,
request_timeout: timedelta,
min_id_length: Optional[int],
events_endpoint: Optional[str],
identification_endpoint: Optional[str]) -> None:
self._api_key = api_key
self._request_timeout = request_timeout
self._min_id_length = min_id_length
self._on_error = on_error
self._queue: queue.Queue = AsyncConsumer.create_queue()
self._endpoints = {
"events": Endpoint(url=events_endpoint or "https://api.amplitude.com/2/httpapi", is_json=True),
"identification": Endpoint(url=identification_endpoint or "https://api.amplitude.com/identify",
is_json=False),
}
self._session = Session()
self._consumer = AsyncConsumer(message_queue=self._queue,
do_upload=self._upload_batch,
flush_queue_size=flush_queue_size,
flush_interval=flush_interval)
atexit.register(self.shutdown)
self._consumer.start()
def __init__(self, api_endpoint: str, api_key: str, flush_queue_size: int,
flush_interval: timedelta, request_timeout: timedelta,
omit_values: bool, retry_options: IterativelyRetryOptions,
on_error: Callable[[str], None]) -> None:
self._api_endpoint = api_endpoint
self._api_key = api_key
self._request_timeout = request_timeout
self._omit_values = omit_values
self._retry_options = retry_options
self._on_error = on_error
self._queue: queue.Queue = AsyncConsumer.create_queue()
self._session = Session()
self._consumer = AsyncConsumer(self._queue,
do_upload=self._upload_batch,
flush_queue_size=flush_queue_size,
flush_interval=flush_interval)
atexit.register(self.shutdown)
self._consumer.start()
def __init__(
self,
base_url: str,
api_key: str,
flush_queue_size: int,
flush_interval: timedelta,
request_timeout: timedelta,
logger: Logger,
) -> None:
self._api_key = api_key
self._request_timeout = request_timeout
self._queue: queue.Queue = AsyncConsumer.create_queue()
base_url = base_url.rstrip("/")
self._user_track_url = f'{base_url}/users/track'
self._session = Session()
self._logger = logger
self._consumer = AsyncConsumer(message_queue=self._queue,
do_upload=self._upload_batch,
flush_queue_size=flush_queue_size,
flush_interval=flush_interval)
atexit.register(self.shutdown)
self._consumer.start()
class IterativelyClient:
def __init__(self, api_endpoint: str, api_key: str, flush_queue_size: int,
flush_interval: timedelta, request_timeout: timedelta,
omit_values: bool, retry_options: IterativelyRetryOptions,
on_error: Callable[[str], None]) -> None:
self._api_endpoint = api_endpoint
self._api_key = api_key
self._request_timeout = request_timeout
self._omit_values = omit_values
self._retry_options = retry_options
self._on_error = on_error
self._queue: queue.Queue = AsyncConsumer.create_queue()
self._session = Session()
self._consumer = AsyncConsumer(self._queue,
do_upload=self._upload_batch,
flush_queue_size=flush_queue_size,
flush_interval=flush_interval)
atexit.register(self.shutdown)
self._consumer.start()
def track(self,
track_type: TrackType,
event: Optional[Event] = None,
properties: Optional[Properties] = None,
validation: Optional[ValidationResponse] = None) -> None:
date_sent = datetime.utcnow().strftime(
'%Y-%m-%dT%H:%M:%S.%f')[:-3] + "Z"
model = {
"type": track_type.value,
"dateSent": date_sent,
"properties": {},
"valid": True,
"validation": {
"details": "",
},
}
if event is not None:
model['eventName'] = event.name
if event.id is not None:
model['eventId'] = event.id
if event.version is not None:
model['eventSchemaVersion'] = event.version
if properties is not None:
if self._omit_values:
model['properties'] = {k: None for k in properties.to_json()}
else:
model['properties'] = properties.to_json()
if validation is not None:
model['valid'] = validation.valid
if not self._omit_values and validation.message is not None:
model['validation']['details'] = validation.message
self._enqueue(AsyncConsumerMessage("events", model))
def _upload_batch(self, batch: List[AsyncConsumerMessage],
stop_event: threading.Event) -> None:
data = {
'objects': [message.data for message in batch],
}
try:
self._send_request(data, stop_event)
except Exception as e:
self._on_error(str(e))
def _send_request(self, data: Any, stop_event: threading.Event) -> None:
need_retry = self._post_request(data)
if not need_retry:
return
for delay in backoff(
start=self._retry_options.delay_initial.total_seconds(),
stop=self._retry_options.delay_maximum.total_seconds(),
count=self._retry_options.max_retries - 1,
factor=2.0,
jitter=1.0):
if stop_event.wait(delay):
return
need_retry = self._post_request(data)
if not need_retry:
return
raise Exception("Failed to upload events. Maximum attempts exceeded.")
def _post_request(self, data: Any) -> bool:
try:
response = self._session.post(
self._api_endpoint,
json=data,
headers={'Authorization': 'Bearer ' + self._api_key},
timeout=self._request_timeout.total_seconds())
except requests.ConnectionError:
return True
except requests.Timeout:
return True
except Exception as e:
raise Exception(f"A unhandled exception occurred. ({e}).")
if 200 <= response.status_code < 300:
return False
if 500 <= response.status_code < 600:
return True
if response.status_code == 429:
return True
raise Exception(
f"Upload failed due to unhandled HTTP error ({response.status_code})."
)
def shutdown(self) -> None:
self._consumer.shutdown()
def _enqueue(self, message: AsyncConsumerMessage) -> None:
try:
self._queue.put(message)
except queue.Full:
self._on_error("async queue is full")
def flush(self) -> None:
self._consumer.flush()
class AmplitudeClient:
def __init__(self,
api_key: str,
on_error: Callable[[str], None],
flush_queue_size: int,
flush_interval: timedelta,
request_timeout: timedelta,
min_id_length: Optional[int],
events_endpoint: Optional[str],
identification_endpoint: Optional[str]) -> None:
self._api_key = api_key
self._request_timeout = request_timeout
self._min_id_length = min_id_length
self._on_error = on_error
self._queue: queue.Queue = AsyncConsumer.create_queue()
self._endpoints = {
"events": Endpoint(url=events_endpoint or "https://api.amplitude.com/2/httpapi", is_json=True),
"identification": Endpoint(url=identification_endpoint or "https://api.amplitude.com/identify",
is_json=False),
}
self._session = Session()
self._consumer = AsyncConsumer(message_queue=self._queue,
do_upload=self._upload_batch,
flush_queue_size=flush_queue_size,
flush_interval=flush_interval)
atexit.register(self.shutdown)
self._consumer.start()
def track(self, user_id: str, event_name: str, properties: Optional[Dict[str, Any]], metadata: Optional[AmplitudeMetadata]) -> None:
data = {k: v for (k, v) in vars(metadata).items() if v is not None} if metadata is not None else {}
data["user_id"] = user_id
data["event_type"] = event_name
data["event_properties"] = properties if properties is not None else {}
if "time" not in data:
data["time"] = int(time.time() * 1000)
self._enqueue(AsyncConsumerMessage("events", data))
def identify(self, user_id: str, properties: Optional[Dict[str, Any]], metadata: Optional[AmplitudeMetadata]) -> None:
data = {k: v for (k, v) in vars(metadata).items() if v is not None} if metadata is not None else {}
data["user_id"] = user_id
data["user_properties"] = properties if properties is not None else {}
self._enqueue(AsyncConsumerMessage("identification", data))
def _upload_batch(self, batch: List[AsyncConsumerMessage], stop_event: Event) -> None:
message_type = batch[0].message_type
endpoint_url, is_json = self._endpoints[message_type]
try:
if message_type == "events":
data = {
"events": [message.data for message in batch],
"api_key": self._api_key,
}
if self._min_id_length is not None:
data["options"] = {"min_id_length": self._min_id_length}
else:
data = {
"identification": json.dumps([message.data for message in batch]),
"api_key": self._api_key
}
self._send_request(Request(url=endpoint_url, is_json=is_json, data=data))
except Exception as e:
self._on_error(str(e))
def _send_request(self, request: Request) -> None:
if request.is_json:
response = self._session.post(request.url, json=request.data, timeout=self._request_timeout.total_seconds())
else:
response = self._session.post(request.url, data=request.data, timeout=self._request_timeout.total_seconds())
if response.status_code >= 300:
self._on_error(f'Unexpected status code for {request.url}: {response.status_code}')
def shutdown(self) -> None:
self._consumer.shutdown()
def _enqueue(self, message: AsyncConsumerMessage) -> None:
try:
self._queue.put(message)
except queue.Full:
self._on_error("async queue is full")
def flush(self) -> None:
self._consumer.flush()
def test_consumer():
batches = []
q = AsyncConsumer.create_queue()
consumer = AsyncConsumer(
message_queue=q,
do_upload=lambda batch, event: batches.append([msg.data for msg in batch]),
flush_queue_size=3,
flush_interval=timedelta(seconds=1)
)
try:
consumer.start()
q.put(AsyncConsumerMessage(message_type='data', data='1'))
q.put(AsyncConsumerMessage(message_type='data', data='2'))
time.sleep(0.1)
assert batches == []
q.put(AsyncConsumerMessage(message_type='data', data='3'))
time.sleep(0.1)
assert batches == [["1", "2", "3"]]
q.put(AsyncConsumerMessage(message_type='data', data='4'))
time.sleep(0.1)
assert batches == [["1", "2", "3"]]
consumer.flush()
time.sleep(0.1)
assert batches == [["1", "2", "3"], ["4"]]
consumer.flush()
consumer.flush()
q.put(AsyncConsumerMessage(message_type='data', data='5'))
q.put(AsyncConsumerMessage(message_type='data', data='6'))
time.sleep(0.1)
assert batches == [["1", "2", "3"], ["4"]]
time.sleep(1)
assert batches == [["1", "2", "3"], ["4"], ["5", "6"]]
q.put(AsyncConsumerMessage(message_type='message', data='7'))
q.put(AsyncConsumerMessage(message_type='data', data='8'))
time.sleep(0.1)
assert batches == [["1", "2", "3"], ["4"], ["5", "6"], ["7"]]
consumer.flush()
time.sleep(0.1)
assert batches == [["1", "2", "3"], ["4"], ["5", "6"], ["7"], ["8"]]
q.put(AsyncConsumerMessage(message_type='data', data='9'))
q.put(AsyncConsumerMessage(message_type='data', data='10'))
finally:
consumer.shutdown()
time.sleep(0.1)
assert batches == [["1", "2", "3"], ["4"], ["5", "6"], ["7"], ["8"], ["9", "10"]]
class BrazeClient:
def __init__(
self,
base_url: str,
api_key: str,
flush_queue_size: int,
flush_interval: timedelta,
request_timeout: timedelta,
logger: Logger,
) -> None:
self._api_key = api_key
self._request_timeout = request_timeout
self._queue: queue.Queue = AsyncConsumer.create_queue()
base_url = base_url.rstrip("/")
self._user_track_url = f'{base_url}/users/track'
self._session = Session()
self._logger = logger
self._consumer = AsyncConsumer(message_queue=self._queue,
do_upload=self._upload_batch,
flush_queue_size=flush_queue_size,
flush_interval=flush_interval)
atexit.register(self.shutdown)
self._consumer.start()
def identify(self, user_id: str, properties: Optional[Dict[str,
Any]]) -> None:
data = self._to_braze_properties(properties)
data["external_id"] = user_id
self._enqueue(AsyncConsumerMessage("", {"attributes": data}))
def track(self, user_id: str, event_name: str,
properties: Optional[Dict[str, Any]]) -> None:
data = {
"external_id": user_id,
"name": event_name,
"time": datetime.now().isoformat(),
"properties": self._to_braze_properties(properties),
}
self._enqueue(AsyncConsumerMessage("", {"events": data}))
def _upload_batch(self, batch: List[AsyncConsumerMessage],
stop_event: Event) -> None:
body = {}
count = 0
for event in batch:
for key, value in event.data.items():
if key not in body:
body[key] = []
count += 1
body[key].append(value)
self._logger.info(f"uploading {count} items")
try:
response = self._session.post(
self._user_track_url,
headers={'Authorization': f'Bearer {self._api_key}'},
json=body,
timeout=self._request_timeout.total_seconds(),
)
if response.status_code >= 300:
self._logger.error(
f'unexpected response status: {response.status_code}')
else:
self._logger.info(f'response status: {response.status_code}')
except Exception as e:
self._logger.error(str(e))
def flush(self) -> None:
self._consumer.flush()
def shutdown(self) -> None:
self._consumer.shutdown()
def _enqueue(self, message: AsyncConsumerMessage) -> None:
try:
self._queue.put(message)
except queue.Full:
self._logger.error("async queue is full")
@staticmethod
def _to_braze_properties(
properties: Optional[Dict[str, Any]]) -> Optional[Dict[str, Any]]:
if properties is None:
return None
return {
key: BrazeClient._value_for_api(value)
for key, value in properties.items()
}
@staticmethod
def _value_for_api(value: Any) -> Any:
# https://www.braze.com/docs/api/objects_filters/user_attributes_object/
# https://www.braze.com/docs/api/objects_filters/event_object/
# API doesn't support objects and (non-string) arrays.
return json.dumps(value) if isinstance(value, (list, dict)) else value