diff --git a/logistics_chat/db.py b/logistics_chat/db.py index 1937715..fbe7d08 100644 --- a/logistics_chat/db.py +++ b/logistics_chat/db.py @@ -37,6 +37,8 @@ def init_db(): Path(get_setting('DATABASE')).parent.mkdir(parents=True, exist_ok=True) conn = connect() cur = conn.cursor() + # Serialize schema inspection and migration when multiple workers start together. + cur.execute('BEGIN IMMEDIATE') cur.execute(''' CREATE TABLE IF NOT EXISTS chat_history ( id INTEGER PRIMARY KEY AUTOINCREMENT, diff --git a/tests/test_storage.py b/tests/test_storage.py index 230616a..7483502 100644 --- a/tests/test_storage.py +++ b/tests/test_storage.py @@ -6,6 +6,8 @@ from pathlib import Path import tempfile import time import unittest +from concurrent.futures import ThreadPoolExecutor +from threading import Barrier from flask import Flask from werkzeug.datastructures import FileStorage @@ -29,6 +31,27 @@ class StorageTests(unittest.TestCase): self.addCleanup(self.context.pop) db.init_db() + def test_concurrent_startup_migrates_legacy_history_once_without_claiming_it(self): + connection = db.connect() + connection.executescript('DROP TABLE chat_history; CREATE TABLE chat_history ' + '(id INTEGER PRIMARY KEY, title TEXT NOT NULL, ctime TEXT NOT NULL);' + "INSERT INTO chat_history VALUES(1,'旧记录','now');") + connection.close() + gate = Barrier(4) + + def initialize(): + with self.app.app_context(): + gate.wait(timeout=5) + db.init_db() + + with ThreadPoolExecutor(max_workers=4) as workers: + for result in workers.map(lambda _: initialize(), range(4)): + self.assertIsNone(result) + connection = db.connect() + self.assertEqual(connection.execute('SELECT title, visitor_id FROM chat_history').fetchall(), [('旧记录', None)]) + self.assertEqual(connection.execute('SELECT COUNT(*) FROM messages WHERE conv_id=1').fetchone()[0], 1) + connection.close() + def test_old_schema_is_migrated_and_legacy_title_backfilled(self): connection = db.connect() connection.executescript('DROP TABLE chat_history; CREATE TABLE chat_history '