source-posthog: make persons an incremental stream

Author: nicolaslazo

source-posthog/acmeCo/Persons.schema.yaml

@@ -15,6 +15,13 @@ $defs:
         description: Row ID of the Document, counting up from zero, or -1 if not known
         title: Row Id
         type: integer
+      project_id:
+        anyOf:
+        - type: integer
+        - type: 'null'
+        default: null
+        description: The PostHog project this document belongs to
+        title: Project Id
     title: Meta
     type: object
 additionalProperties: true
@@ -30,6 +37,13 @@ properties:
     format: date-time
     title: Created At
     type: string
+  last_seen_at:
+    anyOf:
+    - format: date-time
+      type: string
+    - type: 'null'
+    default: null
+    title: Last Seen At
 required:
 - id
 - created_at

source-posthog/acmeCo/flow.yaml

@@ -24,7 +24,8 @@ collections:
   acmeCo/Persons:
     schema: Persons.schema.yaml
     key:
-    - /_meta/row_id
+    - /_meta/project_id
+    - /id
   acmeCo/Projects:
     schema: Projects.schema.yaml
     key:

source-posthog/source_posthog/api.py

@@ -2,7 +2,6 @@
 PostHog API client functions.
 """
 
-import functools
 from collections.abc import AsyncGenerator
 from datetime import UTC, datetime, timedelta
 from logging import Logger
@@ -30,6 +29,7 @@
 )
 
 HOGQL_PAGE_SIZE = 50_000
+BACKFILL_TIMEOUT_PERIOD = timedelta(minutes=5)
 
 
 # Cache for project IDs per organization (avoids re-fetching on retry).
@@ -160,9 +160,21 @@ async def _query_hogql[T: HogQLEntity[str] | HogQLEntity[int]](
     http: HTTPSession,
     log: Logger,
 ) -> AsyncGenerator[T]:
+    log.debug(
+        "Querying HogQL",
+        {
+            "table": model.table_name,
+            "project_id": project_id,
+            "start": start_date,
+            "end": end_date,
+        },
+    )
+
     url = model.get_api_endpoint_url(base_url, project_id)
     column_names = await _get_hogql_columns(model, base_url, project_id, http, log)
 
+    coalesced_cursor_fields = f'COALESCE({",".join(model.cursor_columns)})'
+
     serialized_start_date = start_date.astimezone(UTC).replace(tzinfo=None).isoformat()
     serialized_end_date = (
         end_date.astimezone(UTC).replace(tzinfo=None).isoformat()
@@ -171,11 +183,11 @@ async def _query_hogql[T: HogQLEntity[str] | HogQLEntity[int]](
     )
 
     start_date_clause = (
-        f"WHERE {model.cursor_column} > "
+        f"WHERE {coalesced_cursor_fields} > "
         + f"toDateTime64('{serialized_start_date}', 6, 'UTC') "
     )
     end_date_clause = (
-        f"AND {model.cursor_column} <= toDateTime64('{serialized_end_date}', 6, 'UTC') "
+        f"AND {coalesced_cursor_fields} <= toDateTime64('{serialized_end_date}', 6, 'UTC') "
         if end_date is not None
         else ""
     )
@@ -187,7 +199,7 @@ async def _query_hogql[T: HogQLEntity[str] | HogQLEntity[int]](
             + f"FROM {model.table_name} "
             + start_date_clause
             + end_date_clause
-            + f"ORDER BY {model.cursor_column} DESC "
+            + f"ORDER BY {coalesced_cursor_fields} ASC "
             + f"LIMIT {HOGQL_PAGE_SIZE}",
         },
     }
@@ -299,7 +311,6 @@ async def backfill_project_events(
         return
 
     base_url = config.advanced.base_url
-    ctx = ProjectIdValidationContext(project_id=project_id)
     new_cursor = start_date
     doc_count = 0
 
@@ -343,7 +354,6 @@ async def fetch_project_events(
     assert isinstance(cursor, datetime)
 
     base_url = config.advanced.base_url
-    ctx = ProjectIdValidationContext(project_id=project_id)
     now = datetime.now(tz=UTC)
     upper_bound = now - horizon if horizon else None
 
@@ -380,34 +390,98 @@ async def fetch_project_events(
         yield new_cursor
 
 
-async def snapshot_persons(
+async def backfill_persons(
     http: HTTPSession,
     config: EndpointConfig,
+    project_id: int,
     log: Logger,
-) -> AsyncGenerator[Person, None]:
+    page: PageCursor | None,
+    cutoff: LogCursor,
+) -> AsyncGenerator[Person | PageCursor, None]:
+    assert isinstance(page, str | None)
+    assert isinstance(cutoff, datetime)
+
+    start_date = datetime.fromisoformat(page) if page is not None else config.start_date
+
+    if start_date >= cutoff:
+        return
+
     base_url = config.advanced.base_url
-    project_ids = await fetch_project_ids(http, config, log)
+    new_cursor = start_date
+    doc_count = 0
+    backfill_timeout = datetime.now(tz=UTC) + BACKFILL_TIMEOUT_PERIOD
 
-    total_doc_count = 0
+    while True:
+        batch_count = 0
 
-    for project_id in project_ids:
-        doc_count = 0
-        project_cursor = datetime.min.replace(tzinfo=UTC)
+        async for item in _query_hogql(
+            Person,
+            new_cursor,
+            cutoff,
+            base_url,
+            project_id,
+            http,
+            log,
+        ):
+            batch_count += 1
+            doc_count += 1
+            item_cursor = item.get_cursor()
+            new_cursor = max(new_cursor, item_cursor)
 
-        while True:
-            async for item in _query_hogql(
-                Person, project_cursor, None, base_url, project_id, http, log
-            ):
-                doc_count += 1
-                project_cursor = item.get_cursor()
+            yield item
 
-                yield item
+        if batch_count < HOGQL_PAGE_SIZE:
+            break
+
+        if datetime.now(tz=UTC) > backfill_timeout:
+            log.info(
+                f"{BACKFILL_TIMEOUT_PERIOD.total_seconds() / 60} "
+                + "minutes have elapsed, emitting a checkpoint"
+            )
+            break
+
+    log.info(f"Backfilled {doc_count} persons from project {project_id}")
+
+    if new_cursor > start_date:
+        yield new_cursor.isoformat()
+
+
+async def fetch_persons(
+    http: HTTPSession,
+    config: EndpointConfig,
+    project_id: int,
+    log: Logger,
+    cursor: LogCursor,
+) -> AsyncGenerator[Person | LogCursor, None]:
+    assert isinstance(cursor, datetime)
+
+    base_url = config.advanced.base_url
+    new_cursor = cursor
+    doc_count = 0
+
+    while True:
+        batch_count = 0
 
-            # If we got fewer rows than page_size, we've reached the end
-            if doc_count < HOGQL_PAGE_SIZE:
-                break
+        async for item in _query_hogql(
+            Person,
+            new_cursor,
+            None,
+            base_url,
+            project_id,
+            http,
+            log,
+        ):
+            batch_count += 1
+            doc_count += 1
+            item_cursor = item.get_cursor()
+            new_cursor = max(new_cursor, item_cursor)
 
-        log.info(f"Fetched {doc_count} persons from project {project_id}")
-        total_doc_count += doc_count
+            yield item
 
-    log.info(f"Fetched {total_doc_count} total persons across all projects")
+        if batch_count < HOGQL_PAGE_SIZE:
+            break
+
+    log.info(f"Fetched {doc_count} persons from project {project_id}")
+
+    if new_cursor > cursor:
+        yield new_cursor

source-posthog/source_posthog/models.py

@@ -164,7 +164,7 @@ class HogQLEntity[T: PostHogPrimaryKey](BasePostHogEntity[T], metaclass=ABCMeta)
     """Entity fetched via HogQL Query API."""
 
     table_name: ClassVar[str]
-    cursor_column: ClassVar[str]
+    cursor_columns: ClassVar[list[str]]
 
     @classmethod
     def get_api_endpoint_url(cls, base_url: str, project_id: int) -> str:
@@ -288,7 +288,7 @@ class Event(HogQLEntity[str]):
 
     resource_name: ClassVar[str] = "Events"
     table_name: ClassVar[str] = "events"
-    cursor_column: ClassVar[str] = "timestamp"
+    cursor_columns: ClassVar[list[str]] = ["timestamp"]
 
     class Meta(BaseDocument.Meta):
         model_config = ConfigDict(validate_assignment=True)
@@ -327,13 +327,38 @@ class Person(HogQLEntity[str]):
 
     resource_name: ClassVar[str] = "Persons"
     table_name: ClassVar[str] = "persons"
-    cursor_column: ClassVar[str] = "created_at"
+    cursor_columns: ClassVar[list[str]] = ["last_seen_at", "created_at"]
+
+    class Meta(BaseDocument.Meta):
+        model_config = ConfigDict(validate_assignment=True)
+        project_id: int | None = Field(
+            default=None,
+            description="The PostHog project this document belongs to",
+        )
+
+    meta_: Meta = Field(  # pyright: ignore[reportIncompatibleVariableOverride]
+        default_factory=lambda: Person.Meta(op="u"),
+        alias="_meta",
+        description="Document metadata",
+    )
 
     created_at: AwareDatetime
+    last_seen_at: AwareDatetime | None = None
+
+    @model_validator(mode="before")
+    @classmethod
+    def _inject_project_id_from_context(cls, data: Any, info: ValidationInfo) -> Any:
+        if not isinstance(data, dict):
+            return data
+        if info.context and isinstance(info.context, ProjectIdValidationContext):
+            meta = data.get("_meta") or {}
+            meta["project_id"] = info.context.project_id
+            data["_meta"] = meta
+        return data
 
     @override
     def get_cursor(self) -> AwareDatetime:
-        return self.created_at
+        return self.last_seen_at or self.created_at
 
 
 # =============================================================================