posthog/ee/clickhouse/queries/actor_base_query.py

import uuid
from typing import (
    Any,
    Dict,
    List,
    Literal,
    Optional,
    Tuple,
    TypedDict,
    Union,
)

from django.db.models.query import QuerySet

from ee.clickhouse.client import sync_execute
from posthog.models import Entity, Filter, Team
from posthog.models.filters.mixins.utils import cached_property
from posthog.models.group import Group
from posthog.models.person import Person


class SerializedPerson(TypedDict):
    type: Literal["person"]
    id: uuid.UUID
    created_at: Optional[str]
    properties: Dict[str, Any]
    is_identified: Optional[bool]
    name: str
    distinct_ids: List[str]


class SerializedGroup(TypedDict):
    type: Literal["group"]
    group_key: str
    created_at: Optional[str]
    properties: Dict[str, Any]


SerializedActor = Union[SerializedPerson, SerializedGroup]
Actor = Union[Person, Group]


class ActorBaseQuery:
    aggregating_by_groups = False

    def __init__(self, team: Team, filter: Filter, entity: Optional[Entity] = None):
        self.team = team
        self.entity = entity
        self.filter = filter

    def actor_query(self) -> Tuple[str, Dict]:
        """ Implemented by subclasses. Must return list of uuids. Can be group uuids (group_key) or person uuids """
        raise NotImplementedError()

    @cached_property
    def is_aggregating_by_groups(self) -> bool:
        if self.entity and self.entity.math == "unique_group":
            return True
        else:
            return False

    def get_actors(self) -> Tuple[QuerySet[Actor], List[SerializedActor]]:
        """ Get actors in data model and dict formats. Builds query and executes """
        query, params = self.actor_query()
        raw_result = sync_execute(query, params)
        actors: QuerySet[Actor]
        serialized_actors: List[SerializedActor] = []
        if self.is_aggregating_by_groups:
            actors, serialized_actors = self._get_groups(raw_result)
        else:
            actors, serialized_actors = self._get_people(raw_result)
        return actors, serialized_actors

    def _get_groups(self, results) -> Tuple[QuerySet[Actor], List[SerializedActor]]:
        """ Get groups from raw SQL results in data model and dict formats """
        groups: QuerySet[Group] = Group.objects.filter(team_id=self.team.pk, group_key__in=[val[0] for val in results])
        return groups, self._serialize_groups(groups)

    def _get_people(self, results) -> Tuple[QuerySet[Actor], List[SerializedActor]]:
        """ Get people from raw SQL results in data model and dict formats """
        persons: QuerySet[Person] = Person.objects.filter(team_id=self.team.pk, uuid__in=[val[0] for val in results])
        return persons, self._serialize_people(persons)

    def _serialize_people(self, data: QuerySet[Person]) -> List[SerializedActor]:
        from posthog.api.person import get_person_name

        return [
            SerializedPerson(
                type="person",
                id=person.uuid,
                created_at=person.created_at,
                properties=person.properties,
                is_identified=person.is_identified,
                name=get_person_name(person),
                distinct_ids=person.distinct_ids,
            )
            for person in data
        ]

    def _serialize_groups(self, data: QuerySet[Group]) -> List[SerializedActor]:
        return [
            SerializedGroup(
                type="group", group_key=group.group_key, created_at=group.created_at, properties=group.group_properties
            )
            for group in data
        ]