Skip to content

Commit 1b1ad96

Browse files
ppluleepplulee
committed
fix: enhance db migration
1 parent 2efba2b commit 1b1ad96

File tree

1 file changed

+112
-77
lines changed

1 file changed

+112
-77
lines changed

‎db_migrate/20251112_spam_topic.py‎

Lines changed: 112 additions & 77 deletions
Original file line numberDiff line numberDiff line change
@@ -1,89 +1,124 @@
11
import sqlite3
2+
import logging
3+
4+
logger = logging.getLogger()
25

36

47
def upgrade(db_path):
58
with sqlite3.connect(db_path) as conn:
69
db_cursor = conn.cursor()
7-
8-
# Add spam_topic setting to store the spam topic ID
9-
db_cursor.execute("""
10-
INSERT OR IGNORE INTO settings (key, value)
11-
VALUES ('spam_topic', NULL)
12-
""")
13-
10+
11+
# Add spam_topic setting
12+
try:
13+
db_cursor.execute("""
14+
INSERT OR IGNORE INTO settings (key, value)
15+
VALUES ('spam_topic', NULL)
16+
""")
17+
conn.commit()
18+
logger.info("Added spam_topic setting")
19+
except Exception as e:
20+
logger.error(f"Failed to add spam_topic setting: {e}")
21+
conn.rollback()
22+
1423
# Create blocked_users table
15-
db_cursor.execute("""
16-
CREATE TABLE IF NOT EXISTS blocked_users (
17-
user_id INTEGER PRIMARY KEY,
18-
username TEXT,
19-
first_name TEXT,
20-
last_name TEXT,
21-
blocked_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
22-
)
23-
""")
24-
25-
# Migrate existing banned users from topics table
26-
db_cursor.execute("""
27-
INSERT OR IGNORE INTO blocked_users (user_id)
28-
SELECT user_id FROM topics WHERE ban = 1 AND user_id IS NOT NULL
29-
""")
30-
31-
# Add settings for blocked user auto-reply
32-
db_cursor.execute("""
33-
INSERT OR IGNORE INTO settings (key, value)
34-
VALUES ('blocked_user_reply_enabled', 'disable')
35-
""")
36-
37-
db_cursor.execute("""
38-
INSERT OR IGNORE INTO settings (key, value)
39-
VALUES ('blocked_user_reply_message', NULL)
40-
""")
41-
42-
# Get current table schema
43-
db_cursor.execute("SELECT sql FROM sqlite_master WHERE type='table' AND name='topics'")
44-
result = db_cursor.fetchone()
45-
46-
if result is None:
47-
# Table doesn't exist, nothing to do
24+
try:
25+
db_cursor.execute("""
26+
CREATE TABLE IF NOT EXISTS blocked_users (
27+
user_id INTEGER PRIMARY KEY,
28+
username TEXT,
29+
first_name TEXT,
30+
last_name TEXT,
31+
blocked_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
32+
)
33+
""")
4834
conn.commit()
49-
return
35+
logger.info("Created blocked_users table")
36+
except Exception as e:
37+
logger.error(f"Failed to create blocked_users table: {e}")
38+
conn.rollback()
5039

51-
# Check if ban column exists
52-
db_cursor.execute("PRAGMA table_info(topics)")
53-
columns = db_cursor.fetchall()
54-
has_ban_column = any(col[1] == 'ban' for col in columns)
40+
# Migrate existing banned users
41+
try:
42+
db_cursor.execute("""
43+
INSERT OR IGNORE INTO blocked_users (user_id)
44+
SELECT user_id FROM topics WHERE ban = 1 AND user_id IS NOT NULL
45+
""")
46+
migrated = db_cursor.rowcount
47+
conn.commit()
48+
logger.info(f"Migrated {migrated} banned users to blocked_users table")
49+
except Exception as e:
50+
logger.error(f"Failed to migrate banned users: {e}")
51+
conn.rollback()
5552

56-
if not has_ban_column:
57-
# Column already removed or never existed
53+
# Add blocked user auto-reply settings
54+
try:
55+
db_cursor.execute("""
56+
INSERT OR IGNORE INTO settings (key, value)
57+
VALUES ('blocked_user_reply_enabled', 'disable')
58+
""")
59+
db_cursor.execute("""
60+
INSERT OR IGNORE INTO settings (key, value)
61+
VALUES ('blocked_user_reply_message', NULL)
62+
""")
5863
conn.commit()
59-
return
60-
61-
# Create new table without ban column
62-
db_cursor.execute("""
63-
CREATE TABLE topics_new
64-
(
65-
id INTEGER PRIMARY KEY,
66-
user_id INTEGER,
67-
thread_id INTEGER
68-
)
69-
""")
70-
71-
# Copy data from old table (excluding ban column)
72-
db_cursor.execute("""
73-
INSERT INTO topics_new (id, user_id, thread_id)
74-
SELECT id, user_id, thread_id
75-
FROM topics
76-
""")
77-
78-
# Drop old table
79-
db_cursor.execute("DROP TABLE topics")
80-
81-
# Rename new table to original name
82-
db_cursor.execute("ALTER TABLE topics_new RENAME TO topics")
83-
84-
# Recreate indexes
85-
db_cursor.execute("CREATE INDEX IF NOT EXISTS idx_user_id ON topics(user_id)")
86-
db_cursor.execute("CREATE INDEX IF NOT EXISTS idx_thread_id ON topics(thread_id)")
87-
88-
conn.commit()
64+
logger.info("Added blocked user reply settings")
65+
except Exception as e:
66+
logger.error(f"Failed to add blocked user reply settings: {e}")
67+
conn.rollback()
68+
69+
# Remove ban column from topics table
70+
try:
71+
# Get current table schema
72+
db_cursor.execute("SELECT sql FROM sqlite_master WHERE type='table' AND name='topics'")
73+
result = db_cursor.fetchone()
74+
75+
if result is None:
76+
logger.info("Topics table doesn't exist, skipping ban column removal")
77+
return
78+
79+
# Check if ban column exists
80+
db_cursor.execute("PRAGMA table_info(topics)")
81+
columns = db_cursor.fetchall()
82+
has_ban_column = any(col[1] == 'ban' for col in columns)
8983

84+
if not has_ban_column:
85+
logger.info("Ban column doesn't exist in topics table, skipping removal")
86+
return
87+
88+
# Create new table without ban column
89+
db_cursor.execute("""
90+
CREATE TABLE topics_new (
91+
id INTEGER PRIMARY KEY,
92+
user_id INTEGER,
93+
thread_id INTEGER
94+
)
95+
""")
96+
97+
# Copy data from old table (excluding ban column)
98+
db_cursor.execute("""
99+
INSERT INTO topics_new (id, user_id, thread_id)
100+
SELECT id, user_id, thread_id
101+
FROM topics
102+
""")
103+
104+
# Drop old table
105+
db_cursor.execute("DROP TABLE topics")
106+
107+
# Rename new table to original name
108+
db_cursor.execute("ALTER TABLE topics_new RENAME TO topics")
109+
110+
# Recreate indexes
111+
db_cursor.execute("CREATE INDEX IF NOT EXISTS idx_user_id ON topics(user_id)")
112+
db_cursor.execute("CREATE INDEX IF NOT EXISTS idx_thread_id ON topics(thread_id)")
113+
114+
conn.commit()
115+
logger.info("Successfully removed ban column from topics table")
116+
except Exception as e:
117+
logger.error(f"Failed to remove ban column from topics table: {e}")
118+
# Try to cleanup if table creation failed
119+
try:
120+
db_cursor.execute("DROP TABLE IF EXISTS topics_new")
121+
conn.commit()
122+
except Exception:
123+
pass
124+
conn.rollback()

0 commit comments

Comments
 (0)