works

anaxi/sync.py

@@ -0,0 +1,159 @@
+import datetime
+import json
+
+from psycopg2.extras import execute_batch
+
+from anaxi.checkpoint import CheckpointManager
+from anaxi.config import (
+    get_cosmos_database_config_from_file,
+    get_postgres_database_config_from_file,
+    get_stream_config_from_file,
+)
+from anaxi.constants import (
+    ASSUMED_PREHISTORICAL_DATETIME,
+    DEFAULT_BATCH_SIZE,
+    DESTINATION_COLUMN_NAME,
+    PATHS,
+)
+from anaxi.cosmos_tools import get_container_client
+from anaxi.logging import get_anaxi_logger
+from anaxi.postgres_tools import create_postgres_connection_from_config
+
+logger = get_anaxi_logger(__name__)
+
+
+class SyncJob:
+
+    def __init__(
+        self,
+        stream_id: str,
+        batch_size: int = DEFAULT_BATCH_SIZE,
+    ) -> None:
+
+        self.stream_config = get_stream_config_from_file(stream_id)
+        self.cosmos_db_config = get_cosmos_database_config_from_file(
+            self.stream_config.cosmos_database_id
+        )
+        self.postgres_db_config = get_postgres_database_config_from_file(
+            self.stream_config.postgres_database
+        )
+        self._batch_size = batch_size
+
+        self.checkpoint_manager = self._create_checkpoint_manager()
+        self.cosmos_container_client = get_container_client(
+            config=self.cosmos_db_config,
+            container_name=self.stream_config.cosmos_container_name,
+        )
+        self.postgres_connection = create_postgres_connection_from_config(
+            config=self.postgres_db_config
+        )
+
+    @property
+    def current_checkpoint(self):
+        return self.checkpoint_manager.highest_synced_timestamp
+
+    def run_sync(self):
+        logger.info("Starting sync.")
+
+        changed_feed_response, cursor = self._prepare_source_and_sink()
+
+        docs_batch = []
+        for doc in changed_feed_response:
+            docs_batch.append(doc)
+            if len(docs_batch) < self._batch_size:
+                continue
+
+            logger.info("Reached batch size, trying to commit to database...")
+            self.commit_batch(
+                docs_batch,
+                cursor,
+            )
+            logger.info("Committed batch.")
+            logger.info(
+                f"Checkpoint is a now at: {self.checkpoint_manager.highest_synced_timestamp}"
+            )
+
+            docs_batch = []
+
+        if docs_batch:
+            logger.info(f"Committing dangling documents ({len(docs_batch)})")
+            self.commit_batch(
+                docs_batch,
+                cursor,
+            )
+            logger.info("Committed final batch.")
+
+        logger.info("Finished sync job.")
+
+    def commit_batch(self, docs_batch, cursor):
+        self.checkpoint_manager.set_new_goal(
+            goal_timestamp=datetime.datetime.fromtimestamp(docs_batch[-1]["_ts"])
+        )
+
+        json_data = [(json.dumps(d),) for d in docs_batch]
+
+        insert_query = (
+            f"INSERT INTO {self.stream_config.postgres_schema_name}.{self.stream_config.cosmos_container_name.replace('-', '_')}"
+            + f" ({DESTINATION_COLUMN_NAME})"
+            + " VALUES (%s)"
+        )
+        self.postgres_connection.rollback()
+
+        execute_batch(cursor, insert_query, json_data)
+
+        try:
+            self.postgres_connection.commit()
+        except Exception as e:
+            logger.error(
+                "An error has happened while trying to commit a batch. Rolling back."
+            )
+            self.postgres_connection.rollback()
+            raise e
+        else:
+            self.checkpoint_manager.commit_goal()
+
+    def _prepare_source_and_sink(self):
+        changed_feed_response = self.cosmos_container_client.query_items_change_feed(
+            start_time=self.checkpoint_manager.highest_synced_timestamp
+        )
+
+        cursor = self.postgres_connection.cursor()
+        cursor.execute(
+            f"""
+            CREATE TABLE IF NOT EXISTS {self.stream_config.postgres_schema_name}.{self.stream_config.cosmos_container_name.replace('-', '_')} (
+                {DESTINATION_COLUMN_NAME} JSONB
+            );
+            """
+        )
+        self.postgres_connection.commit()
+        return changed_feed_response, cursor
+
+    def _create_checkpoint_manager(self):
+        checkpoint_file_path = PATHS.checkpoints_folder_path / (
+            self.stream_config.stream_name + ".yml"
+        )
+        try:
+            checkpoint_manager = CheckpointManager.load_from_yml(
+                file_path=checkpoint_file_path
+            )
+        except FileNotFoundError:
+            checkpoint_manager = CheckpointManager(
+                file_path=checkpoint_file_path, highest_synced_timestamp=None
+            )
+
+        if checkpoint_manager.highest_synced_timestamp is None:
+            logger.info(
+                "Couldn't find an existing checkpoint. Will run from the stream's configured cutoff."
+            )
+            if self.stream_config.cutoff_timestamp is None:
+                logger.info(
+                    "No cutoff has ben set. Loading since the start of the container's history."
+                )
+                checkpoint_manager.highest_synced_timestamp = (
+                    ASSUMED_PREHISTORICAL_DATETIME
+                )
+            else:
+                checkpoint_manager.highest_synced_timestamp = (
+                    self.stream_config.cutoff_timestamp
+                )
+        return checkpoint_manager