Function migrate_documents

migrate_to_postgres.py:74–137

Migrate all documents from NDB to PostgreSQL

(logger)

def migrate_documents(logger) -> int:
    """Migrate all documents from NDB to PostgreSQL"""
    from questions.db_models import Document as NDBDocument
    from questions.db_models import client
    from questions.db_models_postgres import Document as PGDocument
    from questions.db_models_postgres import SessionLocal

    logger.info("Starting document migration...")

    db = SessionLocal()
    migrated_count = 0
    error_count = 0

    try:
        with client.context():
            # Fetch all documents from NDB
            ndb_documents = NDBDocument.query().fetch()
            logger.info(f"Found {len(ndb_documents)} documents in NDB")

            for ndb_doc in ndb_documents:
                try:
                    # Check if document already exists in PostgreSQL
                    # Since NDB uses auto-generated keys, we'll check by user_id, title, and created time
                    existing_doc = (
                        db.query(PGDocument)
                        .filter(
                            PGDocument.user_id == ndb_doc.user_id,
                            PGDocument.title == (ndb_doc.title or "Untitled Document"),
                            PGDocument.created == ndb_doc.created,
                        )
                        .first()
                    )

                    if existing_doc:
                        logger.debug(f"Document for user {ndb_doc.user_id} already exists, skipping...")
                        continue

                    # Create new PostgreSQL document
                    pg_doc = PGDocument(
                        user_id=ndb_doc.user_id,
                        title=ndb_doc.title or "Untitled Document",
                        content=ndb_doc.content,
                        created=ndb_doc.created,
                        updated=ndb_doc.updated,
                    )

                    db.add(pg_doc)
                    db.commit()
                    migrated_count += 1

                    if migrated_count % 100 == 0:
                        logger.info(f"Migrated {migrated_count} documents...")

                except Exception as e:
                    error_count += 1
                    logger.error(f"Error migrating document for user {ndb_doc.user_id}: {e}")
                    db.rollback()
                    continue