from datetime import datetime import logging import os import psycopg2 import sqlite3 import sys # Configure logging logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') logger = logging.getLogger(__name__) # Define the order of table migration to respect foreign key constraints TABLE_MIGRATION_ORDER = [ "users", "servers", "tickets", "email_verification_tokens", "password_reset_tokens", "demo_servers_used", "server_status_history", "sessions", "credit_transactions", "database_configs", "ticket_messages", "ticket_attachments", "backups", "user_domains", "server_billing" ] def get_sqlite_schema(cursor, table_name): """Fetches the schema for a given table from SQLite.""" cursor.execute(f"PRAGMA table_info('{table_name}')") return {row['name']: row['type'].upper() for row in cursor.fetchall()} def transform_row(row, schema): """Transforms a row for PostgreSQL insertion, handling data types.""" transformed = list(row) columns = row.keys() for i, col_name in enumerate(columns): col_type = schema.get(col_name) # Handle booleans if col_type == 'BOOLEAN' and transformed[i] in (0, 1): transformed[i] = bool(transformed[i]) # Handle timestamps elif col_type == 'TIMESTAMP' and isinstance(transformed[i], str) and ' ' in transformed[i]: try: dt = datetime.strptime(transformed[i].split('.')[0], '%Y-%m-%d %H:%M:%S') transformed[i] = dt.isoformat() + 'Z' # Add Z for UTC except (ValueError, TypeError): pass return tuple(transformed) def ensure_server_billing_table(pg_conn): """Checks if server_billing table exists in PostgreSQL and creates it if not.""" with pg_conn.cursor() as cursor: cursor.execute("SELECT to_regclass('public.server_billing');") if cursor.fetchone()[0] is None: logger.info("Table 'server_billing' not found in PostgreSQL. Creating it now.") create_table_sql = """ CREATE TABLE server_billing ( id TEXT PRIMARY KEY, server_id TEXT NOT NULL UNIQUE, creation_date DATE NOT NULL, next_billing_date DATE NOT NULL, last_billing_date DATE, created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP, updated_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP, FOREIGN KEY (server_id) REFERENCES servers (id) ON DELETE CASCADE ); """ cursor.execute(create_table_sql) pg_conn.commit() logger.info("Table 'server_billing' created successfully.") def clear_postgres_tables(pg_conn): """Clears all data from the tables in PostgreSQL by checking existence first.""" logger.warning("Preparing to clear existing data from PostgreSQL tables.") with pg_conn.cursor() as cursor: # Get a list of all existing tables in the public schema cursor.execute("SELECT tablename FROM pg_tables WHERE schemaname = 'public';") existing_tables = {row[0] for row in cursor.fetchall()} # Truncate tables in reverse order to respect foreign keys tables_to_clear = reversed(TABLE_MIGRATION_ORDER) for table in tables_to_clear: if table in existing_tables: try: logger.info(f"Clearing table: {table}") cursor.execute(f'TRUNCATE TABLE "{table}" RESTART IDENTITY CASCADE;') except Exception as e: logger.error(f"Error clearing table '{table}': {e}") pg_conn.rollback() raise else: logger.warning(f"Table '{table}' does not exist in PostgreSQL, skipping clear.") pg_conn.commit() logger.info("All existing tables cleared successfully.") def migrate_data(): sqlite_path = os.environ.get('SQLITE_DB_PATH', 'anchorhost.db') pg_url = os.environ.get('DATABASE_URL') if not pg_url or 'postgres' not in pg_url: logger.error("DATABASE_URL environment variable not set or not a PostgreSQL URL.") return logger.info(f"Starting migration from SQLite DB at '{sqlite_path}' to PostgreSQL.") if not os.path.exists(sqlite_path): logger.error(f"SQLite database not found at '{sqlite_path}'.") return # Use a single connection for the whole process pg_conn = None sqlite_conn = None try: pg_conn = psycopg2.connect(pg_url) # Prompt user for confirmation before clearing data try: if '--no-clear' not in sys.argv: confirm = input("This script will clear all existing data from the target tables in PostgreSQL. Are you sure you want to continue? (y/n): ") if confirm.lower() != 'y': logger.info("Migration aborted by user.") return clear_postgres_tables(pg_conn) except (IOError, EOFError): logger.warning("Could not read user confirmation. Assuming non-interactive mode and proceeding WITH CLEARING DATA.") if '--no-clear' not in sys.argv: clear_postgres_tables(pg_conn) sqlite_conn = sqlite3.connect(sqlite_path) sqlite_conn.row_factory = sqlite3.Row sqlite_cursor = sqlite_conn.cursor() ensure_server_billing_table(pg_conn) for table in TABLE_MIGRATION_ORDER: logger.info(f"Migrating table: {table}") try: schema = get_sqlite_schema(sqlite_cursor, table) sqlite_cursor.execute(f'SELECT * FROM "{table}"') rows = sqlite_cursor.fetchall() except sqlite3.OperationalError: logger.warning(f"Table '{table}' does not exist in SQLite, skipping.") continue if not rows: logger.info(f"Table '{table}' is empty, skipping.") continue columns = rows[0].keys() columns_str = ', '.join(f'"{col}"' for col in columns) placeholders = ', '.join(['%s'] * len(columns)) data_to_insert = [transform_row(row, schema) for row in rows] try: with pg_conn.cursor() as pg_cursor: query = f'INSERT INTO "{table}" ({columns_str}) VALUES ({placeholders})' pg_cursor.executemany(query, data_to_insert) logger.info(f"Successfully inserted {pg_cursor.rowcount} rows into table '{table}'.") except Exception as e: logger.error(f"Error inserting data for table '{table}': {e}") pg_conn.rollback() # Rollback this specific table's transaction continue # Continue with the next table pg_conn.commit() logger.info("Migration completed successfully.") except Exception as e: logger.error(f"An unexpected error occurred during migration: {e}") if pg_conn: pg_conn.rollback() finally: if sqlite_conn: sqlite_conn.close() if pg_conn: pg_conn.close() logger.info("Migration process finished.") if __name__ == "__main__": migrate_data()