Skip to content

Commit

Permalink
Iteration 2.5
Browse files Browse the repository at this point in the history
  • Loading branch information
spalmurray-codecov committed Jan 28, 2025
1 parent e9d093e commit 7b14aea
Showing 2 changed files with 34 additions and 9 deletions.
41 changes: 33 additions & 8 deletions shared/events/amplitude.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
from typing import Literal, TypedDict
from shared.events.base import EventPublisher, EventPublisherPropertyError
from shared.events.base import EventPublisher, EventPublisherPropertyException
from amplitude import Amplitude, BaseEvent, Config, EventOptions

AMPLITUDE_API_KEY = None
@@ -11,14 +11,25 @@
"set_orgs" # special event for setting a user's member orgs
]

# Separate type required to make user_owner_id mandatory, but all others optional
type AmplitudeEventProperty = Literal[
"user_ownerid",
"ownerid",
"org_ids",
]

# Separate type required to make user_ownerid mandatory, but all others optional
class BaseAmplitudeEventProperties(TypedDict, total=True):
user_owner_id: int # ownerid of user performing event action
user_ownerid: int # ownerid of user performing event action

class AmplitudeEventProperties(BaseAmplitudeEventProperties, total=False):
ownerid: int # ownerid of owner being acted upon
org_ids: list[str]

# user_ownerid is always required, don't need to check here.
amplitude_required_properties: dict[AmplitudeEventType, list[AmplitudeEventProperty]] = {
'App Installed': ['ownerid']
}

class AmplitudeEventPublisher(EventPublisher):
client: Amplitude

@@ -31,21 +42,35 @@ def __init__(self):
))

def publish(self, event_type: AmplitudeEventType, event_properties: AmplitudeEventProperties):
# Handle special set_orgs event
if event_type == "set_orgs":
if 'org_ids' not in event_properties:
raise EventPublisherPropertyError("event_type 'set_orgs' requires property 'org_ids'")
raise EventPublisherPropertyException("event_type 'set_orgs' requires property 'org_ids'")

self.client.set_group(
group_type="org",
group_name=[str(orgid) for orgid in event_properties['org_ids']],
event_options=EventOptions(user_id=str(event_properties['user_owner_id']))
event_options=EventOptions(user_id=str(event_properties['user_ownerid']))
)
return

# Handle normal events
structured_payload = transform_properties(event_type, event_properties)

# Track event with validated payload, we will raise an exception before
# this if bad payload.
self.client.track(
BaseEvent(event_type, user_id=str(event_properties['user_owner_id']), event_properties=dict(event_properties))
BaseEvent(event_type, user_id=str(event_properties['user_ownerid']), event_properties=structured_payload)
)
return

def transform_properties(event_type: AmplitudeEventType, event_properties: AmplitudeEventProperties) -> dict:
payload = {}

for property in amplitude_required_properties[event_type]:
if property not in event_properties:
raise EventPublisherPropertyException(f"Property {property} is required for event type {event_type}")
payload[property] = event_properties.get(property)

a = AmplitudeEventPublisher()
return payload

a.publish('set_orgs', {"user_owner_id": 2, "org_ids": ['1', '2']})
2 changes: 1 addition & 1 deletion shared/events/base.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
from abc import ABC, abstractmethod

class EventPublisherPropertyError(Exception):
class EventPublisherPropertyException(Exception):
pass

class EventPublisher[T, P](ABC):

0 comments on commit 7b14aea

Please sign in to comment.