76 lines
2.2 KiB
Python
76 lines
2.2 KiB
Python
"""add_reasoning_tree_columns
|
|
|
|
Add source_ids column to documents table for reasoning tree traversal.
|
|
This enables linking observations (deductive, inductive, contradiction)
|
|
to their source observations.
|
|
|
|
Revision ID: f1a2b3c4d5e6
|
|
Revises: 110bdf470272
|
|
Create Date: 2025-12-11
|
|
|
|
"""
|
|
|
|
from collections.abc import Sequence
|
|
|
|
import sqlalchemy as sa
|
|
from alembic import op
|
|
from sqlalchemy.dialects.postgresql import JSONB
|
|
|
|
from migrations.utils import column_exists, get_schema, index_exists
|
|
|
|
# revision identifiers, used by Alembic.
|
|
revision: str = "f1a2b3c4d5e6"
|
|
down_revision: str | None = "110bdf470272"
|
|
branch_labels: str | Sequence[str] | None = None
|
|
depends_on: str | Sequence[str] | None = None
|
|
|
|
schema = get_schema()
|
|
|
|
|
|
def upgrade() -> None:
|
|
"""Add source_ids column with GIN index for tree traversal."""
|
|
connection = op.get_bind()
|
|
inspector = sa.inspect(connection)
|
|
|
|
# Add source_ids column (nullable, for linking to parent observations)
|
|
if not column_exists("documents", "source_ids", inspector):
|
|
op.add_column(
|
|
"documents",
|
|
sa.Column(
|
|
"source_ids",
|
|
JSONB,
|
|
nullable=True,
|
|
server_default=sa.text("NULL"),
|
|
),
|
|
schema=schema,
|
|
)
|
|
|
|
# Add GIN index on source_ids for efficient child lookups
|
|
# (finding all observations that have a given observation as a source)
|
|
if not index_exists("documents", "ix_documents_source_ids_gin", inspector):
|
|
op.create_index(
|
|
"ix_documents_source_ids_gin",
|
|
"documents",
|
|
["source_ids"],
|
|
postgresql_using="gin",
|
|
schema=schema,
|
|
)
|
|
|
|
|
|
def downgrade() -> None:
|
|
"""Remove source_ids column and its index."""
|
|
connection = op.get_bind()
|
|
inspector = sa.inspect(connection)
|
|
|
|
# Drop GIN index first
|
|
if index_exists("documents", "ix_documents_source_ids_gin", inspector):
|
|
op.drop_index(
|
|
"ix_documents_source_ids_gin",
|
|
table_name="documents",
|
|
schema=schema,
|
|
)
|
|
|
|
# Drop column
|
|
if column_exists("documents", "source_ids", inspector):
|
|
op.drop_column("documents", "source_ids", schema=schema)
|