573 lines
25 KiB
Python
573 lines
25 KiB
Python
"""Database initialization and connection helpers."""
|
|
|
|
|
|
from contextlib import contextmanager
|
|
from datetime import date, datetime, timezone
|
|
from db_wrapper import get_db_connection
|
|
import json
|
|
from runtime_config import (
|
|
DATABASE_URL,
|
|
logger,
|
|
)
|
|
|
|
__all__ = ['check_and_fix_database_schema', 'database_connection', 'dict_from_row', 'get_table_columns', 'init_database']
|
|
|
|
|
|
|
|
def init_database():
|
|
"""Initialize the database with required tables"""
|
|
|
|
try:
|
|
with database_connection() as conn:
|
|
cursor = conn.cursor()
|
|
|
|
# Users table - EXTENDED
|
|
cursor.execute('''
|
|
CREATE TABLE IF NOT EXISTS users (
|
|
id TEXT PRIMARY KEY,
|
|
username TEXT UNIQUE NOT NULL,
|
|
email TEXT UNIQUE NOT NULL,
|
|
password_hash TEXT,
|
|
google_id TEXT UNIQUE,
|
|
apple_id TEXT UNIQUE,
|
|
role TEXT DEFAULT 'user',
|
|
credits INTEGER DEFAULT 100,
|
|
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
|
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
|
display_name TEXT,
|
|
first_name TEXT,
|
|
last_name TEXT,
|
|
address TEXT,
|
|
city TEXT,
|
|
country TEXT,
|
|
postal_code TEXT,
|
|
phone TEXT,
|
|
email_verified BOOLEAN DEFAULT FALSE,
|
|
email_verification_token TEXT,
|
|
password_changed_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
|
two_factor_enabled BOOLEAN DEFAULT FALSE,
|
|
two_factor_secret TEXT,
|
|
last_login TIMESTAMP,
|
|
login_count INTEGER DEFAULT 0
|
|
)
|
|
''')
|
|
|
|
# Add google_id and apple_id columns if they don't exist (migration for existing databases)
|
|
if 'postgres' in DATABASE_URL:
|
|
cursor.execute('ALTER TABLE users ADD COLUMN IF NOT EXISTS google_id TEXT UNIQUE')
|
|
cursor.execute('ALTER TABLE users ADD COLUMN IF NOT EXISTS apple_id TEXT UNIQUE')
|
|
else:
|
|
try:
|
|
cursor.execute('ALTER TABLE users ADD COLUMN google_id TEXT UNIQUE')
|
|
except:
|
|
pass
|
|
try:
|
|
cursor.execute('ALTER TABLE users ADD COLUMN apple_id TEXT UNIQUE')
|
|
except:
|
|
pass
|
|
|
|
# Make password_hash nullable if it's not already (migration for existing databases)
|
|
try:
|
|
# SQLite doesn't support ALTER COLUMN, so we'll handle NULL passwords in code
|
|
pass
|
|
except:
|
|
pass
|
|
|
|
# Email verification tokens table
|
|
cursor.execute('''
|
|
CREATE TABLE IF NOT EXISTS email_verification_tokens (
|
|
id TEXT PRIMARY KEY,
|
|
user_id TEXT NOT NULL,
|
|
token TEXT NOT NULL,
|
|
expires_at TIMESTAMP NOT NULL,
|
|
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
|
FOREIGN KEY (user_id) REFERENCES users (id) ON DELETE CASCADE
|
|
)
|
|
''')
|
|
|
|
# Password reset tokens table
|
|
cursor.execute('''
|
|
CREATE TABLE IF NOT EXISTS password_reset_tokens (
|
|
id TEXT PRIMARY KEY,
|
|
user_id TEXT NOT NULL,
|
|
token TEXT NOT NULL,
|
|
expires_at TIMESTAMP NOT NULL,
|
|
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
|
FOREIGN KEY (user_id) REFERENCES users (id) ON DELETE CASCADE
|
|
)
|
|
''')
|
|
|
|
# Servers table
|
|
cursor.execute('''
|
|
CREATE TABLE IF NOT EXISTS servers (
|
|
id TEXT PRIMARY KEY,
|
|
name TEXT NOT NULL,
|
|
type TEXT NOT NULL,
|
|
image TEXT NOT NULL,
|
|
version TEXT NOT NULL,
|
|
status TEXT DEFAULT 'Creating',
|
|
cpu TEXT NOT NULL,
|
|
ram TEXT NOT NULL,
|
|
storage TEXT NOT NULL,
|
|
ip TEXT NOT NULL,
|
|
ports TEXT NOT NULL,
|
|
owner_id TEXT NOT NULL,
|
|
namespace TEXT NOT NULL,
|
|
k8s_name TEXT NOT NULL,
|
|
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
|
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
|
is_demo BOOLEAN DEFAULT FALSE,
|
|
demo_expires_at TIMESTAMP,
|
|
java_version TEXT,
|
|
loader TEXT,
|
|
sftp_username TEXT,
|
|
sftp_password TEXT,
|
|
sftp_port INTEGER,
|
|
FOREIGN KEY (owner_id) REFERENCES users (id) ON DELETE CASCADE
|
|
)
|
|
''')
|
|
|
|
# Add demo columns if they don't exist (for existing databases)
|
|
if 'postgres' in DATABASE_URL:
|
|
cursor.execute('ALTER TABLE servers ADD COLUMN IF NOT EXISTS is_demo BOOLEAN DEFAULT FALSE')
|
|
cursor.execute('ALTER TABLE servers ADD COLUMN IF NOT EXISTS demo_expires_at TIMESTAMP')
|
|
cursor.execute('ALTER TABLE servers ADD COLUMN IF NOT EXISTS java_version TEXT')
|
|
cursor.execute('ALTER TABLE servers ADD COLUMN IF NOT EXISTS loader TEXT')
|
|
cursor.execute('ALTER TABLE servers ADD COLUMN IF NOT EXISTS sftp_username TEXT')
|
|
cursor.execute('ALTER TABLE servers ADD COLUMN IF NOT EXISTS sftp_password TEXT')
|
|
cursor.execute('ALTER TABLE servers ADD COLUMN IF NOT EXISTS sftp_port INTEGER')
|
|
else:
|
|
try:
|
|
cursor.execute('ALTER TABLE servers ADD COLUMN is_demo BOOLEAN DEFAULT FALSE')
|
|
except:
|
|
pass
|
|
try:
|
|
cursor.execute('ALTER TABLE servers ADD COLUMN demo_expires_at TIMESTAMP')
|
|
except:
|
|
pass
|
|
try:
|
|
cursor.execute('ALTER TABLE servers ADD COLUMN java_version TEXT')
|
|
except:
|
|
pass
|
|
try:
|
|
cursor.execute('ALTER TABLE servers ADD COLUMN loader TEXT')
|
|
except:
|
|
pass
|
|
try:
|
|
cursor.execute('ALTER TABLE servers ADD COLUMN sftp_username TEXT')
|
|
except:
|
|
pass
|
|
try:
|
|
cursor.execute('ALTER TABLE servers ADD COLUMN sftp_password TEXT')
|
|
except:
|
|
pass
|
|
try:
|
|
cursor.execute('ALTER TABLE servers ADD COLUMN sftp_port INTEGER')
|
|
except:
|
|
pass
|
|
|
|
# Demo servers usage tracking table
|
|
cursor.execute('''
|
|
CREATE TABLE IF NOT EXISTS demo_servers_used (
|
|
id TEXT PRIMARY KEY,
|
|
user_id TEXT NOT NULL UNIQUE,
|
|
used_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
|
FOREIGN KEY (user_id) REFERENCES users (id) ON DELETE CASCADE
|
|
)
|
|
''')
|
|
|
|
# Server status history table
|
|
cursor.execute('''
|
|
CREATE TABLE IF NOT EXISTS server_status_history (
|
|
id TEXT PRIMARY KEY,
|
|
server_id TEXT NOT NULL,
|
|
status TEXT NOT NULL,
|
|
message TEXT,
|
|
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
|
FOREIGN KEY (server_id) REFERENCES servers (id) ON DELETE CASCADE
|
|
)
|
|
''')
|
|
|
|
# Sessions table
|
|
cursor.execute('''
|
|
CREATE TABLE IF NOT EXISTS sessions (
|
|
id TEXT PRIMARY KEY,
|
|
user_id TEXT NOT NULL,
|
|
token TEXT NOT NULL,
|
|
expires_at TIMESTAMP NOT NULL,
|
|
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
|
FOREIGN KEY (user_id) REFERENCES users (id) ON DELETE CASCADE
|
|
)
|
|
''')
|
|
|
|
# Credit transactions table
|
|
cursor.execute('''
|
|
CREATE TABLE IF NOT EXISTS credit_transactions (
|
|
id TEXT PRIMARY KEY,
|
|
user_id TEXT NOT NULL,
|
|
amount INTEGER NOT NULL,
|
|
type TEXT NOT NULL,
|
|
description TEXT,
|
|
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
|
FOREIGN KEY (user_id) REFERENCES users (id) ON DELETE CASCADE
|
|
)
|
|
''')
|
|
|
|
# Database configurations table
|
|
cursor.execute('''
|
|
CREATE TABLE IF NOT EXISTS database_configs (
|
|
id TEXT PRIMARY KEY,
|
|
server_id TEXT NOT NULL,
|
|
username TEXT NOT NULL,
|
|
password TEXT NOT NULL,
|
|
database_name TEXT NOT NULL,
|
|
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
|
FOREIGN KEY (server_id) REFERENCES servers (id) ON DELETE CASCADE
|
|
)
|
|
''')
|
|
|
|
# Tickets table
|
|
cursor.execute('''
|
|
CREATE TABLE IF NOT EXISTS tickets (
|
|
id TEXT PRIMARY KEY,
|
|
user_id TEXT NOT NULL,
|
|
subject TEXT NOT NULL,
|
|
description TEXT NOT NULL,
|
|
status TEXT DEFAULT 'open',
|
|
priority TEXT DEFAULT 'medium',
|
|
category TEXT DEFAULT 'general',
|
|
assigned_admin_id TEXT,
|
|
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
|
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
|
FOREIGN KEY (user_id) REFERENCES users (id) ON DELETE CASCADE,
|
|
FOREIGN KEY (assigned_admin_id) REFERENCES users (id) ON DELETE SET NULL
|
|
)
|
|
''')
|
|
|
|
# Backups table
|
|
cursor.execute('''
|
|
CREATE TABLE IF NOT EXISTS backups (
|
|
id TEXT PRIMARY KEY,
|
|
server_id TEXT NOT NULL,
|
|
name TEXT NOT NULL,
|
|
type TEXT NOT NULL, -- 'Manual', 'Scheduled'
|
|
status TEXT DEFAULT 'Pending', -- 'Pending', 'In Progress', 'Completed', 'Failed'
|
|
size TEXT,
|
|
s3_path TEXT,
|
|
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
|
completed_at TIMESTAMP,
|
|
FOREIGN KEY (server_id) REFERENCES servers (id) ON DELETE CASCADE
|
|
)
|
|
''')
|
|
|
|
# Ticket messages table
|
|
cursor.execute('''
|
|
CREATE TABLE IF NOT EXISTS ticket_messages (
|
|
id TEXT PRIMARY KEY,
|
|
ticket_id TEXT NOT NULL,
|
|
user_id TEXT NOT NULL,
|
|
message TEXT NOT NULL,
|
|
is_admin BOOLEAN DEFAULT FALSE,
|
|
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
|
FOREIGN KEY (ticket_id) REFERENCES tickets (id) ON DELETE CASCADE,
|
|
FOREIGN KEY (user_id) REFERENCES users (id) ON DELETE CASCADE
|
|
)
|
|
''')
|
|
|
|
# Ticket attachments table
|
|
# Determine blob type based on database URL
|
|
blob_type = 'BYTEA' if 'postgres' in DATABASE_URL else 'BLOB'
|
|
|
|
cursor.execute(f'''
|
|
CREATE TABLE IF NOT EXISTS ticket_attachments (
|
|
id TEXT PRIMARY KEY,
|
|
ticket_id TEXT NOT NULL,
|
|
message_id TEXT,
|
|
filename TEXT NOT NULL,
|
|
file_data {blob_type} NOT NULL,
|
|
mime_type TEXT NOT NULL,
|
|
file_size INTEGER NOT NULL,
|
|
uploaded_by TEXT NOT NULL,
|
|
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
|
FOREIGN KEY (ticket_id) REFERENCES tickets (id) ON DELETE CASCADE,
|
|
FOREIGN KEY (message_id) REFERENCES ticket_messages (id) ON DELETE CASCADE,
|
|
FOREIGN KEY (uploaded_by) REFERENCES users (id) ON DELETE CASCADE
|
|
)
|
|
''')
|
|
|
|
# Announcements table
|
|
cursor.execute('''
|
|
CREATE TABLE IF NOT EXISTS announcements (
|
|
id TEXT PRIMARY KEY,
|
|
title TEXT NOT NULL,
|
|
content TEXT NOT NULL,
|
|
tags TEXT,
|
|
author_id TEXT NOT NULL,
|
|
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
|
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
|
FOREIGN KEY (author_id) REFERENCES users (id) ON DELETE CASCADE
|
|
)
|
|
''')
|
|
|
|
# Key-value settings (e.g. homepage showcase URLs)
|
|
cursor.execute('''
|
|
CREATE TABLE IF NOT EXISTS site_settings (
|
|
key TEXT PRIMARY KEY,
|
|
value TEXT NOT NULL,
|
|
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
|
)
|
|
''')
|
|
|
|
# Default pricing and panel settings
|
|
default_settings = [
|
|
('price_cpu', '35'),
|
|
('price_ram', '35'),
|
|
('price_storage', '5'),
|
|
('multiplier_vps', '1.0'),
|
|
('multiplier_game', '0.8'),
|
|
('multiplier_website', '1.5'),
|
|
('multiplier_database', '1.3'),
|
|
('demo_servers_enabled', 'true'),
|
|
('new_user_credits', '0')
|
|
]
|
|
for key, val in default_settings:
|
|
if 'postgres' in DATABASE_URL:
|
|
cursor.execute("INSERT INTO site_settings (key, value) VALUES (?, ?) ON CONFLICT (key) DO NOTHING", (key, val))
|
|
else:
|
|
cursor.execute("INSERT OR IGNORE INTO site_settings (key, value) VALUES (?, ?)", (key, val))
|
|
|
|
# User domains table - WITH is_custom column
|
|
cursor.execute('''
|
|
CREATE TABLE IF NOT EXISTS user_domains (
|
|
id TEXT PRIMARY KEY,
|
|
user_id TEXT NOT NULL,
|
|
subdomain TEXT,
|
|
full_domain TEXT,
|
|
domain TEXT,
|
|
is_custom BOOLEAN DEFAULT FALSE,
|
|
verified BOOLEAN DEFAULT FALSE,
|
|
record_type TEXT DEFAULT 'A',
|
|
ip_address TEXT NOT NULL,
|
|
ttl INTEGER DEFAULT 3600,
|
|
server_id TEXT,
|
|
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
|
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
|
FOREIGN KEY (user_id) REFERENCES users (id) ON DELETE CASCADE,
|
|
FOREIGN KEY (server_id) REFERENCES servers (id) ON DELETE SET NULL,
|
|
UNIQUE(user_id, full_domain),
|
|
UNIQUE(user_id, domain)
|
|
)
|
|
''')
|
|
|
|
|
|
|
|
conn.commit()
|
|
|
|
# Create indexes - moved outside transaction for safety
|
|
try:
|
|
cursor.execute('CREATE INDEX IF NOT EXISTS idx_servers_owner_id ON servers(owner_id)')
|
|
cursor.execute('CREATE INDEX IF NOT EXISTS idx_servers_status ON servers(status)')
|
|
cursor.execute('CREATE INDEX IF NOT EXISTS idx_server_status_history_server_id ON server_status_history(server_id)')
|
|
cursor.execute('CREATE INDEX IF NOT EXISTS idx_sessions_user_id ON sessions(user_id)')
|
|
cursor.execute('CREATE INDEX IF NOT EXISTS idx_sessions_token ON sessions(token)')
|
|
cursor.execute('CREATE INDEX IF NOT EXISTS idx_credit_transactions_user_id ON credit_transactions(user_id)')
|
|
cursor.execute('CREATE INDEX IF NOT EXISTS idx_database_configs_server_id ON database_configs(server_id)')
|
|
cursor.execute('CREATE INDEX IF NOT EXISTS idx_tickets_user_id ON tickets(user_id)')
|
|
cursor.execute('CREATE INDEX IF NOT EXISTS idx_tickets_status ON tickets(status)')
|
|
cursor.execute('CREATE INDEX IF NOT EXISTS idx_tickets_priority ON tickets(priority)')
|
|
cursor.execute('CREATE INDEX IF NOT EXISTS idx_ticket_messages_ticket_id ON ticket_messages(ticket_id)')
|
|
cursor.execute('CREATE INDEX IF NOT EXISTS idx_ticket_attachments_ticket_id ON ticket_attachments(ticket_id)')
|
|
cursor.execute('CREATE INDEX IF NOT EXISTS idx_user_domains_user_id ON user_domains(user_id)')
|
|
cursor.execute('CREATE INDEX IF NOT EXISTS idx_user_domains_server_id ON user_domains(server_id)')
|
|
cursor.execute('CREATE INDEX IF NOT EXISTS idx_user_domains_full_domain ON user_domains(full_domain)')
|
|
|
|
# Check if is_custom column exists before creating index
|
|
columns = get_table_columns(cursor, 'user_domains')
|
|
has_is_custom = 'is_custom' in columns
|
|
|
|
if has_is_custom:
|
|
cursor.execute('CREATE INDEX IF NOT EXISTS idx_user_domains_is_custom ON user_domains(is_custom)')
|
|
cursor.execute('CREATE INDEX IF NOT EXISTS idx_user_domains_domain ON user_domains(domain)')
|
|
else:
|
|
logger.warning("is_custom column not found in user_domains table, skipping index creation")
|
|
|
|
conn.commit()
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error creating indexes: {e}")
|
|
conn.rollback()
|
|
# Try to add missing columns
|
|
try:
|
|
columns = get_table_columns(cursor, 'user_domains')
|
|
|
|
if 'is_custom' not in columns:
|
|
logger.info("Adding missing column: is_custom to user_domains table")
|
|
cursor.execute('ALTER TABLE user_domains ADD COLUMN is_custom BOOLEAN DEFAULT FALSE')
|
|
|
|
if 'domain' not in columns:
|
|
logger.info("Adding missing column: domain to user_domains table")
|
|
cursor.execute('ALTER TABLE user_domains ADD COLUMN domain TEXT')
|
|
|
|
|
|
conn.commit()
|
|
|
|
# Now try creating indexes again
|
|
cursor.execute('CREATE INDEX IF NOT EXISTS idx_user_domains_is_custom ON user_domains(is_custom)')
|
|
cursor.execute('CREATE INDEX IF NOT EXISTS idx_user_domains_domain ON user_domains(domain)')
|
|
conn.commit()
|
|
|
|
except Exception as alter_error:
|
|
logger.error(f"Error adding missing columns: {alter_error}")
|
|
|
|
# Default rows for homepage showcase (3 configurable links)
|
|
try:
|
|
cursor.execute("SELECT 1 FROM site_settings WHERE key = ?", ('homepage_showcase',))
|
|
if not cursor.fetchone():
|
|
default_links = json.dumps([
|
|
{"label": "Site 1", "url": ""},
|
|
{"label": "Site 2", "url": ""},
|
|
{"label": "Site 3", "url": ""},
|
|
])
|
|
cursor.execute(
|
|
"INSERT INTO site_settings (key, value) VALUES (?, ?)",
|
|
("homepage_showcase", default_links),
|
|
)
|
|
conn.commit()
|
|
except Exception as seed_err:
|
|
logger.warning(f"homepage_showcase default seed: {seed_err}")
|
|
|
|
logger.info("Database initialized successfully")
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error initializing database: {e}")
|
|
raise
|
|
|
|
def check_and_fix_database_schema():
|
|
"""Check and fix any missing columns in existing tables"""
|
|
from server_provisioning import generate_sftp_credentials
|
|
|
|
try:
|
|
with database_connection() as conn:
|
|
cursor = conn.cursor()
|
|
|
|
# Check user_domains table for missing columns
|
|
columns = get_table_columns(cursor, 'user_domains')
|
|
|
|
missing_columns = []
|
|
if 'is_custom' not in columns:
|
|
missing_columns.append('is_custom')
|
|
if 'domain' not in columns:
|
|
missing_columns.append('domain')
|
|
if 'verified' not in columns:
|
|
missing_columns.append('verified')
|
|
if 'java_version' not in columns:
|
|
missing_columns.append('java_version')
|
|
if 'loader' not in columns:
|
|
missing_columns.append('loader')
|
|
|
|
if missing_columns:
|
|
logger.info(f"Adding missing columns to user_domains table: {missing_columns}")
|
|
|
|
if 'is_custom' in missing_columns:
|
|
cursor.execute('ALTER TABLE user_domains ADD COLUMN is_custom BOOLEAN DEFAULT FALSE')
|
|
|
|
if 'domain' in missing_columns:
|
|
cursor.execute('ALTER TABLE user_domains ADD COLUMN domain TEXT')
|
|
|
|
if 'verified' in missing_columns:
|
|
cursor.execute('ALTER TABLE user_domains ADD COLUMN verified BOOLEAN DEFAULT FALSE')
|
|
|
|
# Check servers table for new columns
|
|
server_columns = get_table_columns(cursor, 'servers')
|
|
if 'java_version' not in server_columns:
|
|
logger.info("Adding missing column: java_version to servers table")
|
|
cursor.execute('ALTER TABLE servers ADD COLUMN java_version TEXT')
|
|
|
|
if 'loader' not in server_columns:
|
|
logger.info("Adding missing column: loader to servers table")
|
|
cursor.execute('ALTER TABLE servers ADD COLUMN loader TEXT')
|
|
|
|
if 'sftp_username' not in server_columns:
|
|
logger.info("Adding missing column: sftp_username to servers table")
|
|
cursor.execute('ALTER TABLE servers ADD COLUMN sftp_username TEXT')
|
|
|
|
if 'sftp_password' not in server_columns:
|
|
logger.info("Adding missing column: sftp_password to servers table")
|
|
cursor.execute('ALTER TABLE servers ADD COLUMN sftp_password TEXT')
|
|
|
|
if 'sftp_port' not in server_columns:
|
|
logger.info("Adding missing column: sftp_port to servers table")
|
|
cursor.execute('ALTER TABLE servers ADD COLUMN sftp_port INTEGER')
|
|
|
|
conn.commit()
|
|
logger.info("✓ Added missing columns to user_domains and servers table")
|
|
|
|
# Create indexes that might be missing
|
|
try:
|
|
cursor.execute('CREATE INDEX IF NOT EXISTS idx_user_domains_is_custom ON user_domains(is_custom)')
|
|
cursor.execute('CREATE INDEX IF NOT EXISTS idx_user_domains_domain ON user_domains(domain)')
|
|
conn.commit()
|
|
except Exception as index_error:
|
|
logger.warning(f"Could not create indexes: {index_error}")
|
|
|
|
# Populate missing SFTP credentials for existing servers
|
|
try:
|
|
cursor.execute("SELECT id, type FROM servers WHERE sftp_username IS NULL OR sftp_username = ''")
|
|
servers_to_fix = cursor.fetchall()
|
|
if servers_to_fix:
|
|
logger.info(f"Populating SFTP credentials for {len(servers_to_fix)} servers")
|
|
for row in servers_to_fix:
|
|
srv = dict_from_row(row)
|
|
if srv['type'] != 'vps':
|
|
u, p, port = generate_sftp_credentials()
|
|
cursor.execute(
|
|
'UPDATE servers SET sftp_username = ?, sftp_password = ?, sftp_port = ? WHERE id = ?',
|
|
(u, p, port, srv['id'])
|
|
)
|
|
conn.commit()
|
|
logger.info("✓ Populated missing SFTP credentials")
|
|
except Exception as sftp_fix_error:
|
|
logger.warning(f"Could not populate SFTP credentials: {sftp_fix_error}")
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error checking database schema: {e}")
|
|
|
|
@contextmanager
|
|
def database_connection():
|
|
"""Context manager for database connections"""
|
|
conn = get_db_connection(DATABASE_URL)
|
|
# Row factory is handled in get_db_connection for SQLite, and via cursor_factory for Postgres
|
|
try:
|
|
yield conn
|
|
except Exception as e:
|
|
conn.rollback()
|
|
raise e
|
|
finally:
|
|
conn.close()
|
|
|
|
def dict_from_row(row):
|
|
"""Convert sqlite3.Row or dict to dictionary, making naive datetimes UTC-aware."""
|
|
if row is None:
|
|
return None
|
|
|
|
result = dict(row)
|
|
|
|
for key, value in result.items():
|
|
if isinstance(value, datetime):
|
|
# If datetime is naive, assume it's UTC and make it aware.
|
|
if value.tzinfo is None:
|
|
value = value.replace(tzinfo=timezone.utc)
|
|
result[key] = value.isoformat()
|
|
elif isinstance(value, date):
|
|
result[key] = value.isoformat()
|
|
return result
|
|
|
|
def get_table_columns(cursor, table_name):
|
|
"""Get column names for a table in a database-agnostic way"""
|
|
if 'postgres' in DATABASE_URL:
|
|
cursor.execute("SELECT column_name FROM information_schema.columns WHERE table_name = ?", (table_name,))
|
|
return [row['column_name'] for row in cursor.fetchall()]
|
|
else:
|
|
cursor.execute(f"PRAGMA table_info({table_name})")
|
|
return [col[1] for col in cursor.fetchall()]
|