192 lines
7.3 KiB
Python
192 lines
7.3 KiB
Python
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()
|