Cognition ETL Provider

alembic/versions/9d5fb67e29f7_config_sets.py

@@ -0,0 +1,51 @@
+"""Config sets'
+
+
+Revision ID: 9d5fb67e29f7
+Revises: f428a22ecdb3
+Create Date: 2025-11-03 15:28:47.686657
+
+"""
+from alembic import op
+import sqlalchemy as sa
+from sqlalchemy.dialects import postgresql
+
+# revision identifiers, used by Alembic.
+revision = '9d5fb67e29f7'
+down_revision = 'f428a22ecdb3'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table('etl_config_preset',
+    sa.Column('id', postgresql.UUID(as_uuid=True), nullable=False),
+    sa.Column('organization_id', postgresql.UUID(as_uuid=True), nullable=True),
+    sa.Column('project_id', postgresql.UUID(as_uuid=True), nullable=True),
+    sa.Column('name', sa.String(), nullable=True),
+    sa.Column('description', sa.String(), nullable=True),
+    sa.Column('created_at', sa.DateTime(), nullable=True),
+    sa.Column('created_by', postgresql.UUID(as_uuid=True), nullable=True),
+    sa.Column('etl_config', sa.JSON(), nullable=True),
+    sa.Column('add_config', sa.JSON(), nullable=True),
+    sa.ForeignKeyConstraint(['created_by'], ['user.id'], ondelete='SET NULL'),
+    sa.ForeignKeyConstraint(['organization_id'], ['organization.id'], ondelete='CASCADE'),
+    sa.ForeignKeyConstraint(['project_id'], ['cognition.project.id'], ondelete='CASCADE'),
+    sa.PrimaryKeyConstraint('id'),
+    sa.UniqueConstraint('name'),
+    schema='cognition'
+    )
+    op.create_index(op.f('ix_cognition_etl_config_preset_created_by'), 'etl_config_preset', ['created_by'], unique=False, schema='cognition')
+    op.create_index(op.f('ix_cognition_etl_config_preset_organization_id'), 'etl_config_preset', ['organization_id'], unique=False, schema='cognition')
+    op.create_index(op.f('ix_cognition_etl_config_preset_project_id'), 'etl_config_preset', ['project_id'], unique=False, schema='cognition')
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_index(op.f('ix_cognition_etl_config_preset_project_id'), table_name='etl_config_preset', schema='cognition')
+    op.drop_index(op.f('ix_cognition_etl_config_preset_organization_id'), table_name='etl_config_preset', schema='cognition')
+    op.drop_index(op.f('ix_cognition_etl_config_preset_created_by'), table_name='etl_config_preset', schema='cognition')
+    op.drop_table('etl_config_preset', schema='cognition')
+    # ### end Alembic commands ###

alembic/versions/f428a22ecdb3_adds_etl_task_table.py

@@ -0,0 +1,318 @@
+"""adds etl task table
+
+Revision ID: f428a22ecdb3
+Revises: 85bb3ebee137
+Create Date: 2025-10-30 10:45:20.843280
+
+"""
+
+from alembic import op
+import sqlalchemy as sa
+from sqlalchemy.dialects import postgresql
+
+# revision identifiers, used by Alembic.
+revision = "f428a22ecdb3"
+down_revision = "85bb3ebee137"
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table(
+        "etl_task",
+        sa.Column("id", postgresql.UUID(as_uuid=True), nullable=False),
+        sa.Column("organization_id", postgresql.UUID(as_uuid=True), nullable=True),
+        sa.Column("created_at", sa.DateTime(), nullable=True),
+        sa.Column("created_by", postgresql.UUID(as_uuid=True), nullable=True),
+        sa.Column("file_path", sa.String(), nullable=True),
+        sa.Column("file_size_bytes", sa.BigInteger(), nullable=True),
+        sa.Column("tokenizer", sa.String(), nullable=True),
+        sa.Column("full_config", sa.JSON(), nullable=True),
+        sa.Column("started_at", sa.DateTime(), nullable=True),
+        sa.Column("finished_at", sa.DateTime(), nullable=True),
+        sa.Column("state", sa.String(), nullable=True),
+        sa.Column("is_active", sa.Boolean(), nullable=True),
+        sa.Column("priority", sa.Integer(), nullable=True),
+        sa.Column("error_message", sa.String(), nullable=True),
+        sa.ForeignKeyConstraint(["created_by"], ["user.id"], ondelete="SET NULL"),
+        sa.ForeignKeyConstraint(
+            ["organization_id"], ["organization.id"], ondelete="CASCADE"
+        ),
+        sa.PrimaryKeyConstraint("id"),
+        schema="global",
+    )
+    op.create_index(
+        op.f("ix_global_etl_task_created_by"),
+        "etl_task",
+        ["created_by"],
+        unique=False,
+        schema="global",
+    )
+    op.create_index(
+        op.f("ix_global_etl_task_organization_id"),
+        "etl_task",
+        ["organization_id"],
+        unique=False,
+        schema="global",
+    )
+    op.add_column(
+        "markdown_file",
+        sa.Column("etl_task_id", postgresql.UUID(as_uuid=True), nullable=True),
+        schema="cognition",
+    )
+    op.create_index(
+        op.f("ix_cognition_markdown_file_etl_task_id"),
+        "markdown_file",
+        ["etl_task_id"],
+        unique=False,
+        schema="cognition",
+    )
+    op.create_unique_constraint(
+        "unique_markdown_file_etl_task_id",
+        "markdown_file",
+        ["id", "etl_task_id"],
+        schema="cognition",
+    )
+    op.create_foreign_key(
+        None,
+        "markdown_file",
+        "etl_task",
+        ["etl_task_id"],
+        ["id"],
+        source_schema="cognition",
+        referent_schema="global",
+        ondelete="CASCADE",
+    )
+    op.add_column(
+        "github_file",
+        sa.Column("etl_task_id", postgresql.UUID(as_uuid=True), nullable=True),
+        schema="integration",
+    )
+    op.drop_constraint(
+        "unique_github_file_source", "github_file", schema="integration", type_="unique"
+    )
+    op.create_unique_constraint(
+        "unique_github_file_source",
+        "github_file",
+        ["integration_id", "running_id", "source", "etl_task_id"],
+        schema="integration",
+    )
+    op.create_index(
+        op.f("ix_integration_github_file_etl_task_id"),
+        "github_file",
+        ["etl_task_id"],
+        unique=False,
+        schema="integration",
+    )
+    op.create_foreign_key(
+        None,
+        "github_file",
+        "etl_task",
+        ["etl_task_id"],
+        ["id"],
+        source_schema="integration",
+        referent_schema="global",
+        ondelete="CASCADE",
+    )
+    op.add_column(
+        "github_issue",
+        sa.Column("etl_task_id", postgresql.UUID(as_uuid=True), nullable=True),
+        schema="integration",
+    )
+    op.drop_constraint(
+        "unique_github_issue_source",
+        "github_issue",
+        schema="integration",
+        type_="unique",
+    )
+    op.create_unique_constraint(
+        "unique_github_issue_source",
+        "github_issue",
+        ["integration_id", "running_id", "source", "etl_task_id"],
+        schema="integration",
+    )
+    op.create_index(
+        op.f("ix_integration_github_issue_etl_task_id"),
+        "github_issue",
+        ["etl_task_id"],
+        unique=False,
+        schema="integration",
+    )
+    op.create_foreign_key(
+        None,
+        "github_issue",
+        "etl_task",
+        ["etl_task_id"],
+        ["id"],
+        source_schema="integration",
+        referent_schema="global",
+        ondelete="CASCADE",
+    )
+    op.add_column(
+        "pdf",
+        sa.Column("etl_task_id", postgresql.UUID(as_uuid=True), nullable=True),
+        schema="integration",
+    )
+    op.drop_constraint("unique_pdf_source", "pdf", schema="integration", type_="unique")
+    op.create_unique_constraint(
+        "unique_pdf_source",
+        "pdf",
+        ["integration_id", "running_id", "source", "etl_task_id"],
+        schema="integration",
+    )
+    op.create_index(
+        op.f("ix_integration_pdf_etl_task_id"),
+        "pdf",
+        ["etl_task_id"],
+        unique=False,
+        schema="integration",
+    )
+    op.create_foreign_key(
+        None,
+        "pdf",
+        "etl_task",
+        ["etl_task_id"],
+        ["id"],
+        source_schema="integration",
+        referent_schema="global",
+        ondelete="CASCADE",
+    )
+    op.add_column(
+        "sharepoint",
+        sa.Column("etl_task_id", postgresql.UUID(as_uuid=True), nullable=True),
+        schema="integration",
+    )
+    op.drop_constraint(
+        "unique_sharepoint_source", "sharepoint", schema="integration", type_="unique"
+    )
+    op.create_unique_constraint(
+        "unique_sharepoint_source",
+        "sharepoint",
+        ["integration_id", "running_id", "source", "etl_task_id"],
+        schema="integration",
+    )
+    op.create_index(
+        op.f("ix_integration_sharepoint_etl_task_id"),
+        "sharepoint",
+        ["etl_task_id"],
+        unique=False,
+        schema="integration",
+    )
+    op.create_foreign_key(
+        None,
+        "sharepoint",
+        "etl_task",
+        ["etl_task_id"],
+        ["id"],
+        source_schema="integration",
+        referent_schema="global",
+        ondelete="CASCADE",
+    )
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_constraint(
+        "sharepoint_etl_task_id_fkey",
+        "sharepoint",
+        schema="integration",
+        type_="foreignkey",
+    )
+    op.drop_index(
+        op.f("ix_integration_sharepoint_etl_task_id"),
+        table_name="sharepoint",
+        schema="integration",
+    )
+    op.drop_constraint(
+        "unique_sharepoint_source", "sharepoint", schema="integration", type_="unique"
+    )
+    op.create_unique_constraint(
+        "unique_sharepoint_source",
+        "sharepoint",
+        ["integration_id", "running_id", "source"],
+        schema="integration",
+    )
+    op.drop_column("sharepoint", "etl_task_id", schema="integration")
+    op.drop_constraint(
+        "pdf_etl_task_id_fkey", "pdf", schema="integration", type_="foreignkey"
+    )
+    op.drop_index(
+        op.f("ix_integration_pdf_etl_task_id"), table_name="pdf", schema="integration"
+    )
+    op.drop_constraint("unique_pdf_source", "pdf", schema="integration", type_="unique")
+    op.create_unique_constraint(
+        "unique_pdf_source",
+        "pdf",
+        ["integration_id", "running_id", "source"],
+        schema="integration",
+    )
+    op.drop_column("pdf", "etl_task_id", schema="integration")
+    op.drop_constraint(
+        "github_issue_etl_task_id_fkey",
+        "github_issue",
+        schema="integration",
+        type_="foreignkey",
+    )
+    op.drop_index(
+        op.f("ix_integration_github_issue_etl_task_id"),
+        table_name="github_issue",
+        schema="integration",
+    )
+    op.drop_constraint(
+        "unique_github_issue_source",
+        "github_issue",
+        schema="integration",
+        type_="unique",
+    )
+    op.create_unique_constraint(
+        "unique_github_issue_source",
+        "github_issue",
+        ["integration_id", "running_id", "source"],
+        schema="integration",
+    )
+    op.drop_column("github_issue", "etl_task_id", schema="integration")
+    op.drop_constraint(
+        "github_file_etl_task_id_fkey",
+        "github_file",
+        schema="integration",
+        type_="foreignkey",
+    )
+    op.drop_index(
+        op.f("ix_integration_github_file_etl_task_id"),
+        table_name="github_file",
+        schema="integration",
+    )
+    op.drop_constraint(
+        "unique_github_file_source", "github_file", schema="integration", type_="unique"
+    )
+    op.create_unique_constraint(
+        "unique_github_file_source",
+        "github_file",
+        ["integration_id", "running_id", "source"],
+        schema="integration",
+    )
+    op.drop_column("github_file", "etl_task_id", schema="integration")
+    op.drop_constraint(
+        "markdown_file_etl_task_id_fkey",
+        "markdown_file",
+        schema="cognition",
+        type_="foreignkey",
+    )
+    op.drop_index(
+        op.f("ix_cognition_markdown_file_etl_task_id"),
+        table_name="markdown_file",
+        schema="cognition",
+    )
+    op.drop_column("markdown_file", "etl_task_id", schema="cognition")
+    op.drop_index(
+        op.f("ix_global_etl_task_organization_id"),
+        table_name="etl_task",
+        schema="global",
+    )
+    op.drop_index(
+        op.f("ix_global_etl_task_created_by"), table_name="etl_task", schema="global"
+    )
+    op.drop_table("etl_task", schema="global")
+    # ### end Alembic commands ###

controller/monitor/manager.py

@@ -126,3 +126,14 @@ def cancel_integration_task(
     task_monitor.set_integration_task_to_failed(
         integration_id, error_message="Cancelled by task manager"
     )
+
+
+def cancel_etl_task(
+    task_info: Dict[str, Any],
+) -> None:
+
+    etl_task_id = task_info.get("etlTaskId")
+
+    task_monitor.set_etl_task_to_failed(
+        etl_task_id, error_message="Cancelled by task manager"
+    )