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
pool = await asyncpg.create_pool(DATABASE_URL, min_size=2, max_size=10)
async with pool.acquire() as conn:
await conn.execute("""
CREATE TABLE IF NOT EXISTS memories (
id SERIAL PRIMARY KEY,
user_id VARCHAR(100) NOT NULL DEFAULT 'default',
content TEXT NOT NULL,
category VARCHAR(50) DEFAULT 'facts',
tags TEXT DEFAULT '',
expanded_keywords TEXT DEFAULT '',
importance REAL DEFAULT 0.5,
is_sensitive BOOLEAN DEFAULT FALSE,
vault_path TEXT DEFAULT NULL,
encrypted_content BYTEA DEFAULT NULL,
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
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
)
""")
# 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 memories (
id SERIAL PRIMARY KEY,
user_id VARCHAR(100) NOT NULL DEFAULT 'default',
content TEXT NOT NULL,
category VARCHAR(50) DEFAULT 'facts',
tags TEXT DEFAULT '',
expanded_keywords TEXT DEFAULT '',
importance REAL DEFAULT 0.5,
is_sensitive BOOLEAN DEFAULT FALSE,
vault_path TEXT DEFAULT NULL,
encrypted_content BYTEA DEFAULT NULL,
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
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(
"CREATE INDEX IF NOT EXISTS idx_memories_search ON memories USING GIN(search_vector)"
)