Function test_connections

migrate_to_postgres.py:33–71

Test both database connections

(logger)

def test_connections(logger):
    """Test both database connections"""
    logger.info("Testing database connections...")

    # Test PostgreSQL connection
    try:
        from sqlalchemy import text

        from questions.db_models_postgres import SessionLocal, create_tables

        db = SessionLocal()
        db.execute(text("SELECT 1"))
        db.close()
        logger.info("✅ PostgreSQL connection successful")

        # Ensure tables exist
        create_tables()
        logger.info("✅ PostgreSQL tables created/verified")

    except Exception as e:
        logger.error(f"❌ PostgreSQL connection failed: {e}")
        return False

    # Test Google Cloud NDB connection
    try:
        from questions.db_models import Document as NDBDocument
        from questions.db_models import client

        with client.context():
            # Try to query one document to test connection
            NDBDocument.query().fetch(1)
            logger.info("✅ Google Cloud NDB connection successful")

    except Exception as e:
        logger.error(f"❌ Google Cloud NDB connection failed: {e}")
        logger.error("Make sure GOOGLE_APPLICATION_CREDENTIALS is set correctly")
        return False

    return True


def migrate_documents(logger) -> int: