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,28 +11,54 @@ async def init_pool() -> asyncpg.Pool:
global pool global pool
pool = await asyncpg.create_pool(DATABASE_URL, min_size=2, max_size=10) pool = await asyncpg.create_pool(DATABASE_URL, min_size=2, max_size=10)
async with pool.acquire() as conn: async with pool.acquire() as conn:
await conn.execute(""" # Check if table exists
CREATE TABLE IF NOT EXISTS memories ( exists = await conn.fetchval(
id SERIAL PRIMARY KEY, "SELECT EXISTS(SELECT 1 FROM information_schema.tables WHERE table_name = 'memories')"
user_id VARCHAR(100) NOT NULL DEFAULT 'default', )
content TEXT NOT NULL, if not exists:
category VARCHAR(50) DEFAULT 'facts', await conn.execute("""
tags TEXT DEFAULT '', CREATE TABLE memories (
expanded_keywords TEXT DEFAULT '', id SERIAL PRIMARY KEY,
importance REAL DEFAULT 0.5, user_id VARCHAR(100) NOT NULL DEFAULT 'default',
is_sensitive BOOLEAN DEFAULT FALSE, content TEXT NOT NULL,
vault_path TEXT DEFAULT NULL, category VARCHAR(50) DEFAULT 'facts',
encrypted_content BYTEA DEFAULT NULL, tags TEXT DEFAULT '',
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), expanded_keywords TEXT DEFAULT '',
updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), importance REAL DEFAULT 0.5,
search_vector tsvector GENERATED ALWAYS AS ( is_sensitive BOOLEAN DEFAULT FALSE,
setweight(to_tsvector('english', coalesce(content, '')), 'A') || vault_path TEXT DEFAULT NULL,
setweight(to_tsvector('english', coalesce(expanded_keywords, '')), 'B') || encrypted_content BYTEA DEFAULT NULL,
setweight(to_tsvector('english', coalesce(tags, '')), 'C') || created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
setweight(to_tsvector('english', coalesce(category, '')), 'D') updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
) STORED search_vector tsvector GENERATED ALWAYS AS (
) setweight(to_tsvector('english', coalesce(content, '')), 'A') ||
""") setweight(to_tsvector('english', coalesce(expanded_keywords, '')), 'B') ||
setweight(to_tsvector('english', coalesce(tags, '')), 'C') ||
setweight(to_tsvector('english', coalesce(category, '')), 'D')
) 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( await conn.execute(
"CREATE INDEX IF NOT EXISTS idx_memories_search ON memories USING GIN(search_vector)" "CREATE INDEX IF NOT EXISTS idx_memories_search ON memories USING GIN(search_vector)"
) )