feat(cdc) Actually load the csv file and import snapshot

snuba/cli/bulk_load.py

@@ -5,6 +5,7 @@
 from snuba import settings
 from snuba.datasets.factory import get_dataset, DATASET_NAMES
 from snuba.snapshots.postgres_snapshot import PostgresSnapshot
+from snuba.writer import BufferedWriterWrapper
 
 
 @click.command()
@@ -32,5 +33,9 @@ def bulk_load(dataset, dest_table, source, log_level):
     )
 
     loader = dataset.get_bulk_loader(snapshot_source, dest_table)
-
-    loader.load()
+    # TODO: see whether we need to pass options to the writer
+    writer = BufferedWriterWrapper(
+        dataset.get_writer({}, dest_table),
+        settings.BULK_CLICKHOUSE_BUFFER,
+    )
+    loader.load(writer)

snuba/datasets/__init__.py

@@ -27,7 +27,7 @@ def get_schema(self):
     def get_processor(self):
         return self.__processor
 
-    def get_writer(self, options=None):
+    def get_writer(self, options=None, table_name=None):
         from snuba import settings
         from snuba.writer import HTTPBatchWriter
 
@@ -36,6 +36,7 @@ def get_writer(self, options=None):
             settings.CLICKHOUSE_HOST,
             settings.CLICKHOUSE_HTTP_PORT,
             options,
+            table_name,
         )
 
     def default_conditions(self):

snuba/datasets/cdc/cdcprocessors.py

@@ -1,26 +1,92 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Any, Mapping, Optional, Sequence, Type
+
 from snuba.processor import MessageProcessor
+from snuba.writer import WriterTableRow
 
 KAFKA_ONLY_PARTITION = 0  # CDC only works with single partition topics. So partition must be 0
 
 
+class CdcMessageRow(ABC):
+    """
+    Takes care of the data transformation from WAL to clickhouse and from
+    bulk load to clickhouse. The goal is to keep all these transformation
+    function in the same place because they ultimately have to be consistent
+    with the Clickhouse schema.
+    """
+
+    @classmethod
+    def from_wal(
+        cls,
+        offset: int,
+        columnnames: Sequence[str],
+        columnvalues: Sequence[Any],
+    ) -> CdcMessageRow:
+        raise NotImplementedError
+
+    @classmethod
+    def from_bulk(
+        cls,
+        row: Mapping[str, Any],
+    ) -> CdcMessageRow:
+        raise NotImplementedError
+
+    @abstractmethod
+    def to_clickhouse(self) -> WriterTableRow:
+        raise NotImplementedError
+
+
 class CdcProcessor(MessageProcessor):
 
-    def __init__(self, pg_table):
+    def __init__(self, pg_table: str, message_row_class: Type[CdcMessageRow]):
         self.pg_table = pg_table
+        self._message_row_class = message_row_class
 
-    def _process_begin(self, offset):
+    def _process_begin(self, offset: int):
         pass
 
-    def _process_commit(self, offset):
+    def _process_commit(self, offset: int):
         pass
 
-    def _process_insert(self, offset, columnnames, columnvalues):
-        pass
+    def _process_insert(self,
+        offset: int,
+        columnnames: Sequence[str],
+        columnvalues: Sequence[Any],
+    ) -> Optional[WriterTableRow]:
+        return self._message_row_class.from_wal(
+            offset,
+            columnnames,
+            columnvalues
+        ).to_clickhouse()
 
-    def _process_update(self, offset, key, columnnames, columnvalues):
-        pass
+    def _process_update(self,
+        offset: int,
+        key: Mapping[str, Any],
+        columnnames: Sequence[str],
+        columnvalues: Sequence[Any],
+    ) -> Optional[WriterTableRow]:
+        old_key = dict(zip(key['keynames'], key['keyvalues']))
+        new_key = {
+            key: columnvalues[columnnames.index(key)]
+            for key
+            in key['keynames']
+        }
+        # We cannot support a change in the identity of the record
+        # clickhouse will use the identity column to find rows to merge.
+        # if we change it, merging won't work.
+        assert old_key == new_key, 'Changing Primary Key is not supported.'
+        return self._message_row_class.from_wal(
+            offset,
+            columnnames,
+            columnvalues
+        ).to_clickhouse()
 
-    def _process_delete(self, offset, key):
+    def _process_delete(self,
+        offset: int,
+        key: Mapping[str, Any],
+    ) -> Optional[WriterTableRow]:
         pass
 
     def process_message(self, value, metadata):

snuba/datasets/cdc/groupedmessage.py

@@ -1,6 +1,6 @@
 from snuba.clickhouse import ColumnSet, DateTime, Nullable, UInt
 from snuba.datasets import Dataset
-from snuba.datasets.cdc.groupedmessage_processor import GroupedMessageProcessor
+from snuba.datasets.cdc.groupedmessage_processor import GroupedMessageProcessor, GroupedMessageRow
 from snuba.datasets.schema import ReplacingMergeTreeSchema
 from snuba.snapshots.bulk_load import SingleTableBulkLoader
 
@@ -55,4 +55,5 @@ def get_bulk_loader(self, source, dest_table):
             source=source,
             source_table=self.POSTGRES_TABLE,
             dest_table=dest_table,
+            row_processor=lambda row: GroupedMessageRow.from_bulk(row).to_clickhouse(),
         )

snuba/datasets/cdc/groupedmessage_processor.py

@@ -1,54 +1,99 @@
+from __future__ import annotations
+
+from datetime import datetime
+from typing import Any, Mapping, Optional, Sequence
+
+from dataclasses import dataclass
 from dateutil.parser import parse as dateutil_parse
-from snuba.datasets.cdc.cdcprocessors import CdcProcessor
+from snuba.datasets.cdc.cdcprocessors import CdcProcessor, CdcMessageRow
+from snuba.writer import WriterTableRow
 
 
-class GroupedMessageProcessor(CdcProcessor):
+@dataclass(frozen=True)
+class GroupMessageRecord:
+    status: int
+    last_seen: datetime
+    first_seen: datetime
+    active_at: Optional[datetime] = None
+    first_release_id: Optional[int] = None
 
-    def __init__(self, postgres_table):
-        super(GroupedMessageProcessor, self).__init__(
-            pg_table=postgres_table,
-        )
 
-    def __build_record(self, offset, columnnames, columnvalues):
-        raw_data = dict(zip(columnnames, columnvalues))
-        output = {
-            'offset': offset,
-            'id': raw_data['id'],
-            'record_deleted': 0,
-            'status': raw_data['status'],
-            'last_seen': dateutil_parse(raw_data['last_seen']),
-            'first_seen': dateutil_parse(raw_data['first_seen']),
-            'active_at': dateutil_parse(raw_data['active_at']),
-            'first_release_id': raw_data['first_release_id'],
-        }
+@dataclass(frozen=True)
+class GroupedMessageRow(CdcMessageRow):
+    offset: Optional[int]
+    id: int
+    record_deleted: bool
+    record_content: Optional[GroupMessageRecord]
 
-        return output
+    @classmethod
+    def from_wal(cls,
+        offset: int,
+        columnnames: Sequence[str],
+        columnvalues: Sequence[Any],
+    ) -> GroupedMessageRow:
+        raw_data = dict(zip(columnnames, columnvalues))
+        return GroupedMessageRow(
+            offset=offset,
+            id=raw_data['id'],
+            record_deleted=False,
+            record_content=GroupMessageRecord(
+                status=raw_data['status'],
+                last_seen=dateutil_parse(raw_data['last_seen']),
+                first_seen=dateutil_parse(raw_data['first_seen']),
+                active_at=dateutil_parse(raw_data['active_at']),
+                first_release_id=raw_data['first_release_id'],
+            )
+        )
 
-    def _process_insert(self, offset, columnnames, columnvalues):
-        return self.__build_record(offset, columnnames, columnvalues)
+    @classmethod
+    def from_bulk(cls,
+        row: Mapping[str, Any],
+    ) -> GroupedMessageRow:
+        return GroupedMessageRow(
+            offset=None,
+            id=int(row['id']),
+            record_deleted=False,
+            record_content=GroupMessageRecord(
+                status=int(row['status']),
+                last_seen=dateutil_parse(row['last_seen']),
+                first_seen=dateutil_parse(row['first_seen']),
+                active_at=dateutil_parse(row['active_at']),
+                first_release_id=int(row['first_release_id']) if row['first_release_id'] else None,
+            )
+        )
 
-    def _process_delete(self, offset, key):
-        key_names = key['keynames']
-        key_values = key['keyvalues']
-        id = key_values[key_names.index('id')]
+    def to_clickhouse(self) -> WriterTableRow:
+        record = self.record_content
         return {
-            'offset': offset,
-            'id': id,
-            'record_deleted': 1,
-            'status': None,
-            'last_seen': None,
-            'first_seen': None,
-            'active_at': None,
-            'first_release_id': None,
+            'offset': self.offset if self.offset is not None else 0,
+            'id': self.id,
+            'record_deleted': 1 if self.record_deleted else 0,
+            'status': None if not record else record.status,
+            'last_seen': None if not record else record.last_seen,
+            'first_seen': None if not record else record.first_seen,
+            'active_at': None if not record else record.active_at,
+            'first_release_id': None if not record else record.first_release_id,
         }
 
-    def _process_update(self, offset, key, columnnames, columnvalues):
-        new_id = columnvalues[columnnames.index('id')]
+
+class GroupedMessageProcessor(CdcProcessor):
+
+    def __init__(self, postgres_table):
+        super(GroupedMessageProcessor, self).__init__(
+            pg_table=postgres_table,
+            message_row_class=GroupedMessageRow,
+        )
+
+    def _process_delete(self,
+        offset: int,
+        key: Mapping[str, Any],
+    ) -> Optional[WriterTableRow]:
         key_names = key['keynames']
         key_values = key['keyvalues']
-        old_id = key_values[key_names.index('id')]
-        # We cannot support a change in the identity of the record
-        # clickhouse will use the identity column to find rows to merge.
-        # if we change it, merging won't work.
-        assert old_id == new_id, 'Changing Primary Key is not supported.'
-        return self.__build_record(offset, columnnames, columnvalues)
+        id = key_values[key_names.index('id')]
+        return GroupedMessageRow(
+            offset=offset,
+            id=id,
+            record_deleted=True,
+            record_content=None
+        ).to_clickhouse()

snuba/settings_base.py

@@ -43,6 +43,7 @@
 # Snuba Options
 
 SNAPSHOT_LOAD_PRODUCT = 'snuba'
+BULK_CLICKHOUSE_BUFFER = 1000
 
 # Processor/Writer Options
 DEFAULT_BROKERS = ['localhost:9092']

snuba/snapshots/__init__.py

@@ -1,12 +1,14 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
+from collections.abc import Iterator
 
 from contextlib import contextmanager
-from typing import NewType, Generator, IO, Optional, Sequence
+from typing import Any, Mapping, NewType, Generator, IO, Iterable, Optional, Sequence
 from dataclasses import dataclass
 
 SnapshotId = NewType("SnapshotId", str)
+SnapshotTableRow = Mapping[str, Any]
 
 
 @dataclass(frozen=True)
@@ -26,6 +28,12 @@ class SnapshotDescriptor:
     id: SnapshotId
     tables: Sequence[TableConfig]
 
+    def get_table(self, table_name: str):
+        for t in self.tables:
+            if t.table == table_name:
+                return t
+        raise ValueError(f"Table {table_name} does not exists in the snapshot")
+
 
 class BulkLoadSource(ABC):
     """
@@ -41,5 +49,5 @@ def get_descriptor(self) -> SnapshotDescriptor:
 
     @abstractmethod
     @contextmanager
-    def get_table_file(self, table: str) -> Generator[IO[bytes], None, None]:
+    def get_table_file(self, table: str) -> Generator[Iterable[SnapshotTableRow], None, None]:
         raise NotImplementedError

snuba/snapshots/bulk_load.py

@@ -1,8 +1,12 @@
 from abc import ABC, abstractmethod
+from typing import Any, Callable, Mapping
+
 import logging
 
 from snuba.clickhouse import ClickhousePool
 from snuba.snapshots import BulkLoadSource
+from snuba.writer import BufferedWriterWrapper, WriterTableRow
+from snuba.snapshots import SnapshotTableRow
 
 
 class BulkLoader(ABC):
@@ -14,7 +18,7 @@ class BulkLoader(ABC):
     the bulk load operation.
     """
     @abstractmethod
-    def load(self) -> None:
+    def load(self, writer: BufferedWriterWrapper) -> None:
         raise NotImplementedError
 
 
@@ -24,15 +28,17 @@ class SingleTableBulkLoader(BulkLoader):
     """
 
     def __init__(self,
-        source: BulkLoadSource,
-        dest_table: str,
-        dataset_table: str,
-    ):
+                source: BulkLoadSource,
+                dest_table: str,
+                source_table: str,
+                row_processor: Callable[[SnapshotTableRow], WriterTableRow],
+            ):
         self.__source = source
         self.__dest_table = dest_table
         self.__source_table = source_table
+        self.__row_processor = row_processor
 
-    def load(self) -> None:
+    def load(self, writer: BufferedWriterWrapper) -> None:
         logger = logging.getLogger('snuba.bulk-loader')
 
         clickhouse_ro = ClickhousePool(client_settings={
@@ -50,6 +56,11 @@ def load(self) -> None:
         logger.info("Loading snapshot %s", descriptor.id)
 
         with self.__source.get_table_file(self.__source_table) as table:
-            logger.info("Loading table from file %s", table.name)
-            # TODO: Do something with the table file
-            raise NotImplementedError
+            logger.info("Loading table %s from file", self.__source_table)
+            row_count = 0
+            with writer as buffer_writer:
+                for row in table:
+                    clickhouse_data = self.__row_processor(row)
+                    buffer_writer.write(clickhouse_data)
+                    row_count += 1
+            logger.info("Load complete %d records loaded", row_count)