Agent Zero
112 lines
4.4 KiB
Python
112 lines
4.4 KiB
Python
"""Add embedding column to contacts + timestamp columns to audit_log.
|
|
|
|
Fixes two issues found by API integration tests:
|
|
1. contacts.embedding (vector(768)) — ORM model was updated in Phase 5.3 but
|
|
the plugin migration 0002_embeddings.sql was never run as an Alembic migration.
|
|
2. audit_log.created_at, updated_at, deleted_at — AuditLog model inherits TenantMixin
|
|
which expects these columns, but they were never added to the DB table.
|
|
|
|
Revision ID: 0104
|
|
"""
|
|
|
|
from alembic import op
|
|
import sqlalchemy as sa
|
|
|
|
revision = "0104"
|
|
down_revision = "0103"
|
|
|
|
|
|
def upgrade() -> None:
|
|
conn = op.get_bind()
|
|
|
|
# 0. Ensure pgvector extension is installed
|
|
op.execute("CREATE EXTENSION IF NOT EXISTS vector")
|
|
|
|
# 1. Add embedding column to contacts (if not exists)
|
|
result = conn.execute(sa.text(
|
|
"SELECT column_name FROM information_schema.columns "
|
|
"WHERE table_name = 'contacts' AND column_name = 'embedding'"
|
|
))
|
|
if result.fetchone() is None:
|
|
op.execute("ALTER TABLE contacts ADD COLUMN embedding vector(768)")
|
|
op.execute(
|
|
"CREATE INDEX IF NOT EXISTS ix_contacts_embedding "
|
|
"ON contacts USING hnsw(embedding vector_cosine_ops)"
|
|
)
|
|
|
|
# 2. Add embedding columns to other tables (from plugin migration 0002)
|
|
for table in ["mails", "files", "calendar_entries"]:
|
|
result = conn.execute(sa.text(
|
|
f"SELECT column_name FROM information_schema.columns "
|
|
f"WHERE table_name = '{table}' AND column_name = 'embedding'"
|
|
))
|
|
if result.fetchone() is None:
|
|
# Check if table exists
|
|
table_exists = conn.execute(sa.text(
|
|
f"SELECT EXISTS (SELECT 1 FROM information_schema.tables WHERE table_name = '{table}')"
|
|
)).scalar()
|
|
if table_exists:
|
|
op.execute(f"ALTER TABLE {table} ADD COLUMN embedding vector(768)")
|
|
op.execute(
|
|
f"CREATE INDEX IF NOT EXISTS ix_{table}_embedding "
|
|
f"ON {table} USING hnsw(embedding vector_cosine_ops)"
|
|
)
|
|
|
|
# Tags use 384-dim embeddings
|
|
result = conn.execute(sa.text(
|
|
"SELECT column_name FROM information_schema.columns "
|
|
"WHERE table_name = 'tags' AND column_name = 'embedding'"
|
|
))
|
|
if result.fetchone() is None:
|
|
table_exists = conn.execute(sa.text(
|
|
"SELECT EXISTS (SELECT 1 FROM information_schema.tables WHERE table_name = 'tags')"
|
|
)).scalar()
|
|
if table_exists:
|
|
op.execute("ALTER TABLE tags ADD COLUMN embedding vector(384)")
|
|
op.execute(
|
|
"CREATE INDEX IF NOT EXISTS ix_tags_embedding "
|
|
"ON tags USING hnsw(embedding vector_cosine_ops)"
|
|
)
|
|
|
|
# 3. Add timestamp columns to audit_log (if not exists)
|
|
for col in ["created_at", "updated_at", "deleted_at"]:
|
|
result = conn.execute(sa.text(
|
|
f"SELECT column_name FROM information_schema.columns "
|
|
f"WHERE table_name = 'audit_log' AND column_name = '{col}'"
|
|
))
|
|
if result.fetchone() is None:
|
|
op.execute(
|
|
f"ALTER TABLE audit_log ADD COLUMN {col} "
|
|
f"TIMESTAMPTZ DEFAULT NOW()"
|
|
)
|
|
|
|
# 4. Add timestamp columns to deletion_log (if not exists)
|
|
for col in ["created_at", "updated_at", "deleted_at"]:
|
|
result = conn.execute(sa.text(
|
|
f"SELECT column_name FROM information_schema.columns "
|
|
f"WHERE table_name = 'deletion_log' AND column_name = '{col}'"
|
|
))
|
|
if result.fetchone() is None:
|
|
op.execute(
|
|
f"ALTER TABLE deletion_log ADD COLUMN {col} "
|
|
f"TIMESTAMPTZ DEFAULT NOW()"
|
|
)
|
|
|
|
|
|
def downgrade() -> None:
|
|
# Drop embedding columns
|
|
for table in ["contacts", "mails", "companies", "files", "calendar_entries"]:
|
|
op.execute(f"DROP INDEX IF EXISTS ix_{table}_embedding")
|
|
op.execute(f"ALTER TABLE {table} DROP COLUMN IF EXISTS embedding")
|
|
|
|
op.execute("DROP INDEX IF EXISTS ix_tags_embedding")
|
|
op.execute("ALTER TABLE tags DROP COLUMN IF EXISTS embedding")
|
|
|
|
# Drop timestamp columns from audit_log
|
|
for col in ["created_at", "updated_at", "deleted_at"]:
|
|
op.execute(f"ALTER TABLE audit_log DROP COLUMN IF EXISTS {col}")
|
|
|
|
# Drop timestamp columns from deletion_log
|
|
for col in ["created_at", "updated_at", "deleted_at"]:
|
|
op.execute(f"ALTER TABLE deletion_log DROP COLUMN IF EXISTS {col}")
|