fix: add migration logic for existing databases without new columns

This commit is contained in:
Viktor Barzin 2026-03-14 10:22:14 +00:00
parent c67a8336cf
commit 63205dbd0c
No known key found for this signature in database
GPG key ID: 0EB088298288D958

View file

@ -11,8 +11,13 @@ async def init_pool() -> asyncpg.Pool:
global pool
pool = await asyncpg.create_pool(DATABASE_URL, min_size=2, max_size=10)
async with pool.acquire() as conn:
# Check if table exists
exists = await conn.fetchval(
"SELECT EXISTS(SELECT 1 FROM information_schema.tables WHERE table_name = 'memories')"
)
if not exists:
await conn.execute("""
CREATE TABLE IF NOT EXISTS memories (
CREATE TABLE memories (
id SERIAL PRIMARY KEY,
user_id VARCHAR(100) NOT NULL DEFAULT 'default',
content TEXT NOT NULL,
@ -33,6 +38,27 @@ async def init_pool() -> asyncpg.Pool:
) STORED
)
""")
else:
# Migrate existing table: add new columns if missing
columns = [row["column_name"] for row in await conn.fetch(
"SELECT column_name FROM information_schema.columns WHERE table_name = 'memories'"
)]
if "user_id" not in columns:
await conn.execute(
"ALTER TABLE memories ADD COLUMN user_id VARCHAR(100) NOT NULL DEFAULT 'default'"
)
if "is_sensitive" not in columns:
await conn.execute(
"ALTER TABLE memories ADD COLUMN is_sensitive BOOLEAN DEFAULT FALSE"
)
if "vault_path" not in columns:
await conn.execute(
"ALTER TABLE memories ADD COLUMN vault_path TEXT DEFAULT NULL"
)
if "encrypted_content" not in columns:
await conn.execute(
"ALTER TABLE memories ADD COLUMN encrypted_content BYTEA DEFAULT NULL"
)
await conn.execute(
"CREATE INDEX IF NOT EXISTS idx_memories_search ON memories USING GIN(search_vector)"
)