feat: add PostgreSQL synchronization service and management command for labels and roles

AnujChhikara · AnujChhikara · commit 017f5dfdd45a · 2025-08-25T13:17:14.000+05:30
diff --git a/todo/management/commands/sync_postgres_tables.py b/todo/management/commands/sync_postgres_tables.py
@@ -0,0 +1,31 @@
+from django.core.management.base import BaseCommand
+from todo.services.postgres_sync_service import PostgresSyncService
+
+
+class Command(BaseCommand):
+    help = "Synchronize labels and roles PostgreSQL tables with MongoDB data"
+
+    def add_arguments(self, parser):
+        parser.add_argument(
+            "--force",
+            action="store_true",
+            help="Force sync even if tables already have data",
+        )
+
+    def handle(self, *args, **options):
+        self.stdout.write(self.style.SUCCESS("Starting PostgreSQL table synchronization for labels and roles..."))
+
+        try:
+            postgres_sync_service = PostgresSyncService()
+
+            if options["force"]:
+                self.stdout.write("Force sync enabled - will sync all tables regardless of existing data")
+
+            success = postgres_sync_service.sync_all_tables()
+
+            if success:
+                self.stdout.write(self.style.SUCCESS("PostgreSQL table synchronization completed successfully!"))
+            else:
+                self.stdout.write(self.style.ERROR("Some PostgreSQL table synchronizations failed!"))
+        except Exception as e:
+            self.stdout.write(self.style.ERROR(f"PostgreSQL table synchronization failed: {str(e)}"))
diff --git a/todo/repositories/team_creation_invite_code_repository.py b/todo/repositories/team_creation_invite_code_repository.py
@@ -33,7 +33,7 @@ def validate_and_consume_code(cls, code: str, used_by: str) -> Optional[dict]:
                 {"$set": {"is_used": True, "used_by": used_by, "used_at": current_time.isoformat()}},
                 return_document=True,
             )
-            
+
             if result:
                 # Sync the update to PostgreSQL
                 dual_write_service = EnhancedDualWriteService()
@@ -48,16 +48,15 @@ def validate_and_consume_code(cls, code: str, used_by: str) -> Optional[dict]:
                 }
 
                 dual_write_success = dual_write_service.update_document(
-                    collection_name="team_creation_invite_codes", 
-                    data=invite_code_data, 
-                    mongo_id=str(result["_id"])
+                    collection_name="team_creation_invite_codes", data=invite_code_data, mongo_id=str(result["_id"])
                 )
 
                 if not dual_write_success:
                     import logging
+
                     logger = logging.getLogger(__name__)
                     logger.warning(f"Failed to sync team creation invite code update {result['_id']} to Postgres")
-            
+
             return result
         except Exception as e:
             raise Exception(f"Error validating and consuming code: {e}")
diff --git a/todo/services/dual_write_service.py b/todo/services/dual_write_service.py
@@ -294,8 +294,8 @@ def _transform_label_data(self, data: Dict[str, Any]) -> Dict[str, Any]:
             "name": data.get("name"),
             "color": data.get("color", "#000000"),
             "description": data.get("description"),
-            "created_at": data.get("created_at"),
-            "updated_at": data.get("updated_at"),
+            "created_at": data.get("createdAt"),
+            "updated_at": data.get("updatedAt"),
         }
 
     def _transform_role_data(self, data: Dict[str, Any]) -> Dict[str, Any]:
diff --git a/todo/services/postgres_sync_service.py b/todo/services/postgres_sync_service.py
@@ -0,0 +1,228 @@
+import logging
+from django.db import connection
+from django.conf import settings
+
+from todo_project.db.config import DatabaseManager
+from todo.services.dual_write_service import DualWriteService
+
+logger = logging.getLogger(__name__)
+
+
+class PostgresSyncService:
+    """
+    Service to synchronize PostgreSQL tables with MongoDB data.
+    Checks if tables exist and copies data from MongoDB if needed.
+    Currently handles labels and roles tables only.
+    """
+
+    def __init__(self):
+        self.db_manager = DatabaseManager()
+        self.dual_write_service = DualWriteService()
+        self.enabled = getattr(settings, "POSTGRES_SYNC_ENABLED", True)
+
+    def sync_all_tables(self) -> bool:
+        """
+        Synchronize labels and roles PostgreSQL tables with MongoDB data.
+
+        Returns:
+            bool: True if all syncs completed successfully, False otherwise
+        """
+        if not self.enabled:
+            logger.info("PostgreSQL sync is disabled, skipping")
+            return True
+
+        logger.info("Starting PostgreSQL table synchronization for labels and roles")
+        logger.info(f"PostgreSQL sync enabled: {self.enabled}")
+
+        sync_operations = [
+            ("labels", self._sync_labels_table),
+            ("roles", self._sync_roles_table),
+        ]
+
+        success_count = 0
+        total_operations = len(sync_operations)
+
+        for table_name, sync_func in sync_operations:
+            try:
+                logger.info(f"Syncing table: {table_name}")
+                if sync_func():
+                    logger.info(f"Successfully synced table: {table_name}")
+                    success_count += 1
+                else:
+                    logger.error(f"Failed to sync table: {table_name}")
+            except Exception as e:
+                logger.error(f"Error syncing table {table_name}: {str(e)}")
+
+        logger.info(f"PostgreSQL sync completed - {success_count}/{total_operations} tables synced successfully")
+        return success_count == total_operations
+
+    def _check_table_exists(self, table_name: str) -> bool:
+        """
+        Check if a PostgreSQL table exists.
+
+        Args:
+            table_name: Name of the table to check
+
+        Returns:
+            bool: True if table exists, False otherwise
+        """
+        try:
+            with connection.cursor() as cursor:
+                cursor.execute(
+                    """
+                    SELECT EXISTS (
+                        SELECT FROM information_schema.tables 
+                        WHERE table_schema = 'public' 
+                        AND table_name = %s
+                    );
+                """,
+                    [table_name],
+                )
+                return cursor.fetchone()[0]
+        except Exception as e:
+            logger.error(f"Error checking if table {table_name} exists: {str(e)}")
+            return False
+
+    def _get_mongo_collection_count(self, collection_name: str) -> int:
+        """
+        Get the count of documents in a MongoDB collection.
+
+        Args:
+            collection_name: Name of the MongoDB collection
+
+        Returns:
+            int: Number of documents in the collection
+        """
+        try:
+            collection = self.db_manager.get_collection(collection_name)
+
+            # Labels use isDeleted field for soft deletes
+            if collection_name == "labels":
+                return collection.count_documents({"isDeleted": {"$ne": True}})
+            else:
+                # For roles and other collections without soft delete, count all documents
+                return collection.count_documents({})
+
+        except Exception as e:
+            logger.error(f"Error getting count for collection {collection_name}: {str(e)}")
+            return 0
+
+    def _get_postgres_table_count(self, table_name: str) -> int:
+        """
+        Get the count of records in a PostgreSQL table.
+
+        Args:
+            table_name: Name of the PostgreSQL table
+
+        Returns:
+            int: Number of records in the table
+        """
+        try:
+            with connection.cursor() as cursor:
+                cursor.execute(f"SELECT COUNT(*) FROM {table_name};")
+                return cursor.fetchone()[0]
+        except Exception as e:
+            logger.error(f"Error getting count for table {table_name}: {str(e)}")
+            return 0
+
+    def _sync_labels_table(self) -> bool:
+        """Synchronize the labels table."""
+        table_name = "postgres_labels"
+
+        if not self._check_table_exists(table_name):
+            logger.warning(f"Table {table_name} does not exist, skipping sync")
+            return True
+
+        mongo_count = self._get_mongo_collection_count("labels")
+        postgres_count = self._get_postgres_table_count(table_name)
+
+        if postgres_count >= mongo_count:
+            logger.info(f"Labels table already has {postgres_count} records, MongoDB has {mongo_count}. Skipping sync.")
+            return True
+
+        logger.info(f"Syncing labels: MongoDB has {mongo_count} records, PostgreSQL has {postgres_count} records")
+        logger.info(f"Will sync {mongo_count - postgres_count} labels to PostgreSQL")
+
+        try:
+            collection = self.db_manager.get_collection("labels")
+            labels = collection.find({"isDeleted": {"$ne": True}})
+
+            synced_count = 0
+            for label in labels:
+                try:
+                    # Check if label already exists in PostgreSQL
+                    from todo.models.postgres.label import PostgresLabel
+
+                    existing = PostgresLabel.objects.filter(mongo_id=str(label["_id"])).first()
+                    if existing:
+                        continue
+
+                    # Transform data for PostgreSQL
+                    postgres_data = self.dual_write_service._transform_label_data(label)
+                    postgres_data["mongo_id"] = str(label["_id"])
+                    postgres_data["sync_status"] = "SYNCED"
+
+                    logger.debug(f"Creating label in PostgreSQL: {postgres_data}")
+
+                    # Create in PostgreSQL
+                    PostgresLabel.objects.create(**postgres_data)
+                    synced_count += 1
+
+                except Exception as e:
+                    logger.error(f"Error syncing label {label.get('_id')}: {str(e)}")
+                    continue
+
+            logger.info(f"Successfully synced {synced_count} labels to PostgreSQL")
+            return True
+
+        except Exception as e:
+            logger.error(f"Error syncing labels table: {str(e)}")
+            return False
+
+    def _sync_roles_table(self) -> bool:
+        """Synchronize the roles table."""
+        table_name = "postgres_roles"
+
+        if not self._check_table_exists(table_name):
+            logger.warning(f"Table {table_name} does not exist, skipping sync")
+            return True
+
+        mongo_count = self._get_mongo_collection_count("roles")
+        postgres_count = self._get_postgres_table_count(table_name)
+
+        if postgres_count >= mongo_count:
+            logger.info(f"Roles table already has {postgres_count} records, MongoDB has {mongo_count}. Skipping sync.")
+            return True
+
+        logger.info(f"Syncing roles: MongoDB has {mongo_count} records, PostgreSQL has {postgres_count} records")
+
+        try:
+            collection = self.db_manager.get_collection("roles")
+            roles = collection.find({})
+
+            synced_count = 0
+            for role in roles:
+                try:
+                    from todo.models.postgres.role import PostgresRole
+
+                    existing = PostgresRole.objects.filter(mongo_id=str(role["_id"])).first()
+                    if existing:
+                        continue
+
+                    postgres_data = self.dual_write_service._transform_role_data(role)
+                    postgres_data["mongo_id"] = str(role["_id"])
+                    postgres_data["sync_status"] = "SYNCED"
+
+                    PostgresRole.objects.create(**postgres_data)
+                    synced_count += 1
+
+                except Exception as e:
+                    logger.error(f"Error syncing role {role.get('_id')}: {str(e)}")
+                    continue
+
+            logger.info(f"Successfully synced {synced_count} roles to PostgreSQL")
+            return True
+
+        except Exception as e:
+            logger.error(f"Error syncing roles table: {str(e)}")
+            return False
diff --git a/todo_project/db/init.py b/todo_project/db/init.py
@@ -2,6 +2,7 @@
 import time
 from todo_project.db.config import DatabaseManager
 from todo_project.db.migrations import run_all_migrations
+from todo.services.postgres_sync_service import PostgresSyncService
 
 logger = logging.getLogger(__name__)
 
@@ -50,6 +51,16 @@ def initialize_database(max_retries=5, retry_delay=2):
         if not migrations_success:
             logger.warning("Some database migrations failed, but continuing with initialization")
 
+        try:
+            postgres_sync_service = PostgresSyncService()
+            postgres_sync_success = postgres_sync_service.sync_all_tables()
+            if not postgres_sync_success:
+                logger.warning("Some PostgreSQL table synchronizations failed, but continuing with initialization")
+            else:
+                logger.info("PostgreSQL table synchronization completed successfully")
+        except Exception as e:
+            logger.warning(f"PostgreSQL table synchronization failed: {str(e)}, but continuing with initialization")
+
         logger.info("Database initialization completed successfully")
         return True
     except Exception as e:
diff --git a/todo_project/settings/test.py b/todo_project/settings/test.py
@@ -0,0 +1,10 @@
+from .base import *
+
+DUAL_WRITE_ENABLED = False
+
+# Remove PostgreSQL database configuration for tests
+# This prevents Django from trying to connect to PostgreSQL
+DATABASES = {}
+
+# Use MongoDB only for tests
+# The tests will use testcontainers to spin up their own MongoDB instance
