Odysseus v1.0
This commit is contained in:
pewdiepie-archdaemon
281
scripts/update_database.py
Normal file
281
scripts/update_database.py
Normal file
@@ -0,0 +1,281 @@
|
||||
"""
|
||||
update_database.py
|
||||
|
||||
This script updates the database schema by adding new columns to the sessions table
|
||||
and populating them with appropriate values. It handles SQLite's limitations
|
||||
with ALTER TABLE operations by checking if columns exist before attempting to add them.
|
||||
|
||||
The following columns are added:
|
||||
- last_accessed (DateTime): Set to created_at for existing records
|
||||
- is_important (Boolean): Set to False for existing records
|
||||
- message_count (Integer): Calculated from the number of messages in chat_messages table
|
||||
|
||||
Usage:
|
||||
python update_database.py
|
||||
"""
|
||||
|
||||
import sqlite3
|
||||
import os
|
||||
from datetime import datetime
|
||||
from sqlalchemy import create_engine, inspect, text
|
||||
from database import DATABASE_URL, SessionLocal, Base
|
||||
|
||||
def check_column_exists(engine, table_name, column_name):
|
||||
"""Check if a column exists in a table."""
|
||||
inspector = inspect(engine)
|
||||
columns = inspector.get_columns(table_name)
|
||||
return any(col['name'] == column_name for col in columns)
|
||||
|
||||
def add_column_sqlite(db_path, table_name, column_name, column_type, default_value=None):
|
||||
"""
|
||||
Add a column to a SQLite table by creating a new table, copying data, and renaming.
|
||||
This is necessary because SQLite has limited ALTER TABLE support.
|
||||
"""
|
||||
conn = sqlite3.connect(db_path)
|
||||
cursor = conn.cursor()
|
||||
|
||||
# Get current table info
|
||||
cursor.execute(f"PRAGMA table_info({table_name})")
|
||||
columns = cursor.fetchall()
|
||||
column_names = [col[1] for col in columns]
|
||||
|
||||
# Create new table with additional column
|
||||
new_table_name = f"{table_name}_new"
|
||||
|
||||
# Build new column list
|
||||
new_columns = []
|
||||
for col in columns:
|
||||
new_columns.append(f"{col[1]} {col[2]}")
|
||||
|
||||
# Add the new column
|
||||
new_column_def = f"{column_name} {column_type}"
|
||||
if default_value is not None:
|
||||
new_column_def += f" DEFAULT {default_value}"
|
||||
new_columns.append(new_column_def)
|
||||
|
||||
# Create new table
|
||||
columns_sql = ", ".join(new_columns)
|
||||
create_sql = f"CREATE TABLE {new_table_name} ({columns_sql})"
|
||||
cursor.execute(create_sql)
|
||||
|
||||
# Copy data from old table to new table
|
||||
column_names_str = ", ".join(column_names)
|
||||
insert_sql = f"INSERT INTO {new_table_name} ({column_names_str}) SELECT {column_names_str} FROM {table_name}"
|
||||
cursor.execute(insert_sql)
|
||||
|
||||
# Drop old table and rename new table
|
||||
cursor.execute(f"DROP TABLE {table_name}")
|
||||
cursor.execute(f"ALTER TABLE {new_table_name} RENAME TO {table_name}")
|
||||
|
||||
conn.commit()
|
||||
conn.close()
|
||||
|
||||
def update_database():
|
||||
"""Update the database schema and populate new columns."""
|
||||
# Create engine from DATABASE_URL
|
||||
engine = create_engine(DATABASE_URL)
|
||||
|
||||
# Extract database path from DATABASE_URL for SQLite
|
||||
db_path = None
|
||||
if "sqlite" in DATABASE_URL:
|
||||
db_path = DATABASE_URL.replace("sqlite:///", "")
|
||||
# Handle relative paths
|
||||
if not os.path.isabs(db_path):
|
||||
db_path = os.path.join(os.path.dirname(__file__), db_path)
|
||||
|
||||
print(f"Updating database at: {DATABASE_URL}")
|
||||
|
||||
# Start a transaction
|
||||
db = SessionLocal()
|
||||
try:
|
||||
# Add last_accessed column if it doesn't exist
|
||||
if not check_column_exists(engine, 'sessions', 'last_accessed'):
|
||||
print("Adding last_accessed column...")
|
||||
if db_path: # SQLite
|
||||
add_column_sqlite(db_path, 'sessions', 'last_accessed', 'DATETIME')
|
||||
else: # Other databases
|
||||
with engine.connect() as conn:
|
||||
conn.execute(text("ALTER TABLE sessions ADD COLUMN last_accessed DATETIME"))
|
||||
conn.commit()
|
||||
|
||||
# Add is_important column if it doesn't exist
|
||||
if not check_column_exists(engine, 'sessions', 'is_important'):
|
||||
print("Adding is_important column...")
|
||||
if db_path: # SQLite
|
||||
add_column_sqlite(db_path, 'sessions', 'is_important', 'BOOLEAN', '0')
|
||||
else: # Other databases
|
||||
with engine.connect() as conn:
|
||||
conn.execute(text("ALTER TABLE sessions ADD COLUMN is_important BOOLEAN DEFAULT FALSE"))
|
||||
conn.commit()
|
||||
|
||||
# Add message_count column if it doesn't exist
|
||||
if not check_column_exists(engine, 'sessions', 'message_count'):
|
||||
print("Adding message_count column...")
|
||||
if db_path: # SQLite
|
||||
add_column_sqlite(db_path, 'sessions', 'message_count', 'INTEGER', '0')
|
||||
else: # Other databases
|
||||
with engine.connect() as conn:
|
||||
conn.execute(text("ALTER TABLE sessions ADD COLUMN message_count INTEGER DEFAULT 0"))
|
||||
conn.commit()
|
||||
|
||||
# Populate last_accessed with created_at for existing records where last_accessed is NULL
|
||||
print("Populating last_accessed column...")
|
||||
with engine.connect() as conn:
|
||||
conn.execute(text("""
|
||||
UPDATE sessions
|
||||
SET last_accessed = created_at
|
||||
WHERE last_accessed IS NULL
|
||||
"""))
|
||||
conn.commit()
|
||||
|
||||
# Populate is_important with FALSE for existing records where is_important is NULL
|
||||
print("Populating is_important column...")
|
||||
with engine.connect() as conn:
|
||||
conn.execute(text("""
|
||||
UPDATE sessions
|
||||
SET is_important = 0
|
||||
WHERE is_important IS NULL
|
||||
"""))
|
||||
conn.commit()
|
||||
|
||||
# Calculate and populate message_count from chat_messages table
|
||||
print("Calculating and populating message_count column...")
|
||||
with engine.connect() as conn:
|
||||
# First, set all message_count to 0
|
||||
conn.execute(text("UPDATE sessions SET message_count = 0"))
|
||||
|
||||
# Then, count messages for each session and update
|
||||
conn.execute(text("""
|
||||
UPDATE sessions
|
||||
SET message_count = (
|
||||
SELECT COUNT(*)
|
||||
FROM chat_messages
|
||||
WHERE chat_messages.session_id = sessions.id
|
||||
)
|
||||
"""))
|
||||
conn.commit()
|
||||
|
||||
print("Database update completed successfully!")
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error updating database: {e}")
|
||||
db.rollback()
|
||||
raise
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
if __name__ == "__main__":
|
||||
update_database()
|
||||
"""
|
||||
update_database.py
|
||||
|
||||
This script updates the database schema by adding new columns to the sessions table
|
||||
if they don't already exist. It uses raw SQL ALTER TABLE statements to modify
|
||||
the existing SQLite database.
|
||||
|
||||
The following columns are added:
|
||||
- last_accessed (DateTime): Set to created_at for existing records
|
||||
- is_important (Boolean): Set to False for existing records
|
||||
- message_count (Integer): Calculated from the number of messages in chat_messages table
|
||||
|
||||
Usage:
|
||||
python update_database.py
|
||||
"""
|
||||
|
||||
import os
|
||||
from datetime import datetime
|
||||
from sqlalchemy import create_engine, text
|
||||
from database import DATABASE_URL, SessionLocal
|
||||
|
||||
def update_database():
|
||||
"""Update the database schema and populate new columns."""
|
||||
# Create engine from DATABASE_URL
|
||||
engine = create_engine(DATABASE_URL)
|
||||
|
||||
# Start a transaction
|
||||
db = SessionLocal()
|
||||
try:
|
||||
# Add last_accessed column if it doesn't exist
|
||||
try:
|
||||
with engine.connect() as conn:
|
||||
conn.execute(text("ALTER TABLE sessions ADD COLUMN last_accessed DATETIME"))
|
||||
conn.commit()
|
||||
print("Added last_accessed column to sessions table")
|
||||
except Exception as e:
|
||||
if "duplicate column name" in str(e).lower():
|
||||
print("last_accessed column already exists")
|
||||
else:
|
||||
print(f"Error adding last_accessed column: {e}")
|
||||
|
||||
# Add is_important column if it doesn't exist
|
||||
try:
|
||||
with engine.connect() as conn:
|
||||
conn.execute(text("ALTER TABLE sessions ADD COLUMN is_important BOOLEAN DEFAULT FALSE"))
|
||||
conn.commit()
|
||||
print("Added is_important column to sessions table")
|
||||
except Exception as e:
|
||||
if "duplicate column name" in str(e).lower():
|
||||
print("is_important column already exists")
|
||||
else:
|
||||
print(f"Error adding is_important column: {e}")
|
||||
|
||||
# Add message_count column if it doesn't exist
|
||||
try:
|
||||
with engine.connect() as conn:
|
||||
conn.execute(text("ALTER TABLE sessions ADD COLUMN message_count INTEGER DEFAULT 0"))
|
||||
conn.commit()
|
||||
print("Added message_count column to sessions table")
|
||||
except Exception as e:
|
||||
if "duplicate column name" in str(e).lower():
|
||||
print("message_count column already exists")
|
||||
else:
|
||||
print(f"Error adding message_count column: {e}")
|
||||
|
||||
# Populate last_accessed with created_at for existing records where last_accessed is NULL
|
||||
print("Populating last_accessed column...")
|
||||
with engine.connect() as conn:
|
||||
conn.execute(text("""
|
||||
UPDATE sessions
|
||||
SET last_accessed = created_at
|
||||
WHERE last_accessed IS NULL
|
||||
"""))
|
||||
conn.commit()
|
||||
|
||||
# Populate is_important with FALSE for existing records where is_important is NULL
|
||||
print("Populating is_important column...")
|
||||
with engine.connect() as conn:
|
||||
conn.execute(text("""
|
||||
UPDATE sessions
|
||||
SET is_important = 0
|
||||
WHERE is_important IS NULL
|
||||
"""))
|
||||
conn.commit()
|
||||
|
||||
# Calculate and populate message_count from chat_messages table
|
||||
print("Calculating and populating message_count column...")
|
||||
with engine.connect() as conn:
|
||||
# First, set all message_count to 0
|
||||
conn.execute(text("UPDATE sessions SET message_count = 0"))
|
||||
|
||||
# Then, count messages for each session and update
|
||||
conn.execute(text("""
|
||||
UPDATE sessions
|
||||
SET message_count = (
|
||||
SELECT COUNT(*)
|
||||
FROM chat_messages
|
||||
WHERE chat_messages.session_id = sessions.id
|
||||
)
|
||||
"""))
|
||||
conn.commit()
|
||||
|
||||
print("Database update completed successfully!")
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error updating database: {e}")
|
||||
db.rollback()
|
||||
raise
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
if __name__ == "__main__":
|
||||
update_database()
|
||||
Reference in New Issue
Block a user