"""add project knowledge graph tables Revision ID: add_project_knowledge_graph_tables Revises: add_qa_sessions_and_turns Create Date: 2026-03-02 """ from __future__ import annotations from alembic import op import sqlalchemy as sa try: from sqlalchemy.dialects.postgresql import JSONB except Exception: # pragma: no cover JSONB = sa.JSON # revision identifiers, used by Alembic. revision = "add_project_knowledge_graph_tables" down_revision = "add_qa_sessions_and_turns" branch_labels = None depends_on = None def _json_type(bind): return JSONB() if bind.dialect.name == "postgresql" else sa.JSON() def _table_exists(bind, table_name: str) -> bool: inspector = sa.inspect(bind) return inspector.has_table(table_name) def _index_exists(bind, table_name: str, index_name: str) -> bool: inspector = sa.inspect(bind) indexes = inspector.get_indexes(table_name) return any(idx.get("name") == index_name for idx in indexes) def _create_index_if_not_exists(bind, index_name: str, table_name: str, columns: list[str]) -> None: if not _index_exists(bind, table_name, index_name): op.create_index(index_name, table_name, columns) def _drop_index_if_exists(bind, index_name: str, table_name: str) -> None: if _index_exists(bind, table_name, index_name): op.drop_index(index_name, table_name=table_name) def upgrade() -> None: bind = op.get_bind() if not _table_exists(bind, "kg_entities"): op.create_table( "kg_entities", sa.Column("id", sa.Integer(), primary_key=True, autoincrement=True), sa.Column("tenant_id", sa.BigInteger(), nullable=False, server_default="1"), sa.Column("created_at", sa.TIMESTAMP(), nullable=False, server_default=sa.func.now()), sa.Column("updated_at", sa.TIMESTAMP(), nullable=False, server_default=sa.func.now(), onupdate=sa.func.now()), sa.Column("entity_type", sa.String(length=40), nullable=False), sa.Column("entity_key", sa.String(length=255), nullable=False), sa.Column("display_name", sa.String(length=255), nullable=False), sa.Column("payload", _json_type(bind), nullable=True), sa.Column("source_type", sa.String(length=32), nullable=False, server_default="system"), sa.Column("source_ref", sa.String(length=255), nullable=True), sa.Column("last_seen_at", sa.TIMESTAMP(), nullable=False, server_default=sa.func.now()), sa.UniqueConstraint("tenant_id", "entity_type", "entity_key", name="uq_kg_entities_tenant_type_key"), ) _create_index_if_not_exists(bind, "ix_kg_entities_tenant_id", "kg_entities", ["tenant_id"]) _create_index_if_not_exists(bind, "ix_kg_entities_entity_type", "kg_entities", ["entity_type"]) _create_index_if_not_exists(bind, "ix_kg_entities_entity_key", "kg_entities", ["entity_key"]) _create_index_if_not_exists(bind, "ix_kg_entities_source_type", "kg_entities", ["source_type"]) _create_index_if_not_exists(bind, "ix_kg_entities_last_seen_at", "kg_entities", ["last_seen_at"]) if not _table_exists(bind, "kg_relations"): op.create_table( "kg_relations", sa.Column("id", sa.Integer(), primary_key=True, autoincrement=True), sa.Column("tenant_id", sa.BigInteger(), nullable=False, server_default="1"), sa.Column("created_at", sa.TIMESTAMP(), nullable=False, server_default=sa.func.now()), sa.Column("updated_at", sa.TIMESTAMP(), nullable=False, server_default=sa.func.now(), onupdate=sa.func.now()), sa.Column("relation_type", sa.String(length=40), nullable=False), sa.Column("subject_entity_id", sa.Integer(), sa.ForeignKey("kg_entities.id", ondelete="CASCADE"), nullable=False), sa.Column("object_entity_id", sa.Integer(), sa.ForeignKey("kg_entities.id", ondelete="CASCADE"), nullable=False), sa.Column("weight", sa.Float(), nullable=False, server_default="1"), sa.Column("payload", _json_type(bind), nullable=True), sa.Column("source_type", sa.String(length=32), nullable=False, server_default="system"), sa.Column("source_ref", sa.String(length=255), nullable=True), sa.UniqueConstraint( "tenant_id", "relation_type", "subject_entity_id", "object_entity_id", name="uq_kg_relations_tenant_type_subject_object", ), ) _create_index_if_not_exists(bind, "ix_kg_relations_tenant_id", "kg_relations", ["tenant_id"]) _create_index_if_not_exists(bind, "ix_kg_relations_relation_type", "kg_relations", ["relation_type"]) _create_index_if_not_exists(bind, "ix_kg_relations_subject_entity_id", "kg_relations", ["subject_entity_id"]) _create_index_if_not_exists(bind, "ix_kg_relations_object_entity_id", "kg_relations", ["object_entity_id"]) _create_index_if_not_exists(bind, "ix_kg_relations_source_type", "kg_relations", ["source_type"]) if not _table_exists(bind, "kg_evidences"): op.create_table( "kg_evidences", sa.Column("id", sa.Integer(), primary_key=True, autoincrement=True), sa.Column("tenant_id", sa.BigInteger(), nullable=False, server_default="1"), sa.Column("created_at", sa.TIMESTAMP(), nullable=False, server_default=sa.func.now()), sa.Column("updated_at", sa.TIMESTAMP(), nullable=False, server_default=sa.func.now(), onupdate=sa.func.now()), sa.Column("entity_id", sa.Integer(), sa.ForeignKey("kg_entities.id", ondelete="CASCADE"), nullable=True), sa.Column("relation_id", sa.Integer(), sa.ForeignKey("kg_relations.id", ondelete="CASCADE"), nullable=True), sa.Column("evidence_type", sa.String(length=40), nullable=False), sa.Column("content_hash", sa.String(length=64), nullable=False), sa.Column("content_excerpt", sa.Text(), nullable=False), sa.Column("source_type", sa.String(length=32), nullable=False, server_default="system"), sa.Column("source_ref", sa.String(length=255), nullable=True), sa.Column("metadata", _json_type(bind), nullable=True), sa.Column("occurred_at", sa.TIMESTAMP(), nullable=False, server_default=sa.func.now()), ) _create_index_if_not_exists(bind, "ix_kg_evidences_tenant_id", "kg_evidences", ["tenant_id"]) _create_index_if_not_exists(bind, "ix_kg_evidences_entity_id", "kg_evidences", ["entity_id"]) _create_index_if_not_exists(bind, "ix_kg_evidences_relation_id", "kg_evidences", ["relation_id"]) _create_index_if_not_exists(bind, "ix_kg_evidences_evidence_type", "kg_evidences", ["evidence_type"]) _create_index_if_not_exists(bind, "ix_kg_evidences_content_hash", "kg_evidences", ["content_hash"]) _create_index_if_not_exists(bind, "ix_kg_evidences_source_type", "kg_evidences", ["source_type"]) _create_index_if_not_exists(bind, "ix_kg_evidences_occurred_at", "kg_evidences", ["occurred_at"]) if not _table_exists(bind, "ingest_events"): op.create_table( "ingest_events", sa.Column("id", sa.Integer(), primary_key=True, autoincrement=True), sa.Column("tenant_id", sa.BigInteger(), nullable=False, server_default="1"), sa.Column("created_at", sa.TIMESTAMP(), nullable=False, server_default=sa.func.now()), sa.Column("updated_at", sa.TIMESTAMP(), nullable=False, server_default=sa.func.now(), onupdate=sa.func.now()), sa.Column("channel", sa.String(length=32), nullable=False), sa.Column("event_key", sa.String(length=128), nullable=False), sa.Column("status", sa.String(length=20), nullable=False, server_default="pending"), sa.Column("payload", _json_type(bind), nullable=True), sa.Column("retry_count", sa.Integer(), nullable=False, server_default="0"), sa.Column("error_message", sa.Text(), nullable=True), sa.Column("processed_at", sa.TIMESTAMP(), nullable=True), sa.UniqueConstraint("tenant_id", "channel", "event_key", name="uq_ingest_events_tenant_channel_key"), ) _create_index_if_not_exists(bind, "ix_ingest_events_tenant_id", "ingest_events", ["tenant_id"]) _create_index_if_not_exists(bind, "ix_ingest_events_channel", "ingest_events", ["channel"]) _create_index_if_not_exists(bind, "ix_ingest_events_event_key", "ingest_events", ["event_key"]) _create_index_if_not_exists(bind, "ix_ingest_events_status", "ingest_events", ["status"]) _create_index_if_not_exists(bind, "ix_ingest_events_processed_at", "ingest_events", ["processed_at"]) if not _table_exists(bind, "ingest_offsets"): op.create_table( "ingest_offsets", sa.Column("id", sa.Integer(), primary_key=True, autoincrement=True), sa.Column("tenant_id", sa.BigInteger(), nullable=False, server_default="1"), sa.Column("created_at", sa.TIMESTAMP(), nullable=False, server_default=sa.func.now()), sa.Column("updated_at", sa.TIMESTAMP(), nullable=False, server_default=sa.func.now(), onupdate=sa.func.now()), sa.Column("channel", sa.String(length=32), nullable=False), sa.Column("offset_key", sa.String(length=128), nullable=False), sa.Column("offset_value", sa.String(length=255), nullable=False), sa.Column("payload", _json_type(bind), nullable=True), sa.UniqueConstraint("tenant_id", "channel", "offset_key", name="uq_ingest_offsets_tenant_channel_key"), ) _create_index_if_not_exists(bind, "ix_ingest_offsets_tenant_id", "ingest_offsets", ["tenant_id"]) _create_index_if_not_exists(bind, "ix_ingest_offsets_channel", "ingest_offsets", ["channel"]) _create_index_if_not_exists(bind, "ix_ingest_offsets_offset_key", "ingest_offsets", ["offset_key"]) def downgrade() -> None: bind = op.get_bind() if _table_exists(bind, "ingest_offsets"): _drop_index_if_exists(bind, "ix_ingest_offsets_offset_key", "ingest_offsets") _drop_index_if_exists(bind, "ix_ingest_offsets_channel", "ingest_offsets") _drop_index_if_exists(bind, "ix_ingest_offsets_tenant_id", "ingest_offsets") op.drop_table("ingest_offsets") if _table_exists(bind, "ingest_events"): _drop_index_if_exists(bind, "ix_ingest_events_processed_at", "ingest_events") _drop_index_if_exists(bind, "ix_ingest_events_status", "ingest_events") _drop_index_if_exists(bind, "ix_ingest_events_event_key", "ingest_events") _drop_index_if_exists(bind, "ix_ingest_events_channel", "ingest_events") _drop_index_if_exists(bind, "ix_ingest_events_tenant_id", "ingest_events") op.drop_table("ingest_events") if _table_exists(bind, "kg_evidences"): _drop_index_if_exists(bind, "ix_kg_evidences_occurred_at", "kg_evidences") _drop_index_if_exists(bind, "ix_kg_evidences_source_type", "kg_evidences") _drop_index_if_exists(bind, "ix_kg_evidences_content_hash", "kg_evidences") _drop_index_if_exists(bind, "ix_kg_evidences_evidence_type", "kg_evidences") _drop_index_if_exists(bind, "ix_kg_evidences_relation_id", "kg_evidences") _drop_index_if_exists(bind, "ix_kg_evidences_entity_id", "kg_evidences") _drop_index_if_exists(bind, "ix_kg_evidences_tenant_id", "kg_evidences") op.drop_table("kg_evidences") if _table_exists(bind, "kg_relations"): _drop_index_if_exists(bind, "ix_kg_relations_source_type", "kg_relations") _drop_index_if_exists(bind, "ix_kg_relations_object_entity_id", "kg_relations") _drop_index_if_exists(bind, "ix_kg_relations_subject_entity_id", "kg_relations") _drop_index_if_exists(bind, "ix_kg_relations_relation_type", "kg_relations") _drop_index_if_exists(bind, "ix_kg_relations_tenant_id", "kg_relations") op.drop_table("kg_relations") if _table_exists(bind, "kg_entities"): _drop_index_if_exists(bind, "ix_kg_entities_last_seen_at", "kg_entities") _drop_index_if_exists(bind, "ix_kg_entities_source_type", "kg_entities") _drop_index_if_exists(bind, "ix_kg_entities_entity_key", "kg_entities") _drop_index_if_exists(bind, "ix_kg_entities_entity_type", "kg_entities") _drop_index_if_exists(bind, "ix_kg_entities_tenant_id", "kg_entities") op.drop_table("kg_entities")