From a6fe7e583b9388652b59e711950e1119d4ed129b Mon Sep 17 00:00:00 2001 From: CREDO23 Date: Thu, 24 Jul 2025 11:38:49 +0200 Subject: [PATCH] generate new migration / fix migration files --- ..._add_llm_config_table_and_relationships.py | 114 +++++++++++------- .../alembic/versions/12_add_logs_table.py | 54 +++++---- .../versions/13_add_jira_connector_enums.py | 18 +-- .../versions/1_add_github_connector_enum.py | 41 ++++--- .../versions/2_add_linear_connector_enum.py | 37 +++--- ...3_add_linear_connector_to_documenttype_.py | 45 ++++--- .../alembic/versions/4_add_linkup_api_enum.py | 40 +++--- .../6_change_podcast_content_to_transcript.py | 44 ++++--- .../versions/7_remove_is_generated_column.py | 17 +-- .../8_add_content_hash_to_documents.py | 60 +++++---- ...discord_connector_enum_and_documenttype.py | 16 +-- ...1_add_github_connector_to_documenttype_.py | 37 +++--- 12 files changed, 303 insertions(+), 220 deletions(-) diff --git a/surfsense_backend/alembic/versions/11_add_llm_config_table_and_relationships.py b/surfsense_backend/alembic/versions/11_add_llm_config_table_and_relationships.py index d2f0420..007cd70 100644 --- a/surfsense_backend/alembic/versions/11_add_llm_config_table_and_relationships.py +++ b/surfsense_backend/alembic/versions/11_add_llm_config_table_and_relationships.py @@ -6,10 +6,8 @@ Revises: 10 from typing import Sequence, Union -from alembic import op import sqlalchemy as sa -from sqlalchemy.dialects.postgresql import UUID, JSON - +from alembic import op # revision identifiers, used by Alembic. revision: str = "11" @@ -22,7 +20,8 @@ def upgrade() -> None: """Upgrade schema - add LiteLLMProvider enum, LLMConfig table and user LLM preferences.""" # Create enum only if not exists - op.execute(""" + op.execute( + """ DO $$ BEGIN IF NOT EXISTS (SELECT 1 FROM pg_type WHERE typname = 'litellmprovider') THEN @@ -35,14 +34,16 @@ def upgrade() -> None: ); END IF; END$$; - """) + """ + ) # Create llm_configs table only if it doesn't already exist - op.execute(""" + op.execute( + """ DO $$ BEGIN IF NOT EXISTS ( - SELECT FROM information_schema.tables + SELECT FROM information_schema.tables WHERE table_name = 'llm_configs' ) THEN CREATE TABLE llm_configs ( @@ -59,78 +60,103 @@ def upgrade() -> None: ); END IF; END$$; - """) + """ + ) # Create indexes if they don't exist - op.execute(""" + op.execute( + """ DO $$ BEGIN IF NOT EXISTS ( - SELECT 1 FROM pg_indexes + SELECT 1 FROM pg_indexes WHERE tablename = 'llm_configs' AND indexname = 'ix_llm_configs_id' ) THEN CREATE INDEX ix_llm_configs_id ON llm_configs(id); END IF; IF NOT EXISTS ( - SELECT 1 FROM pg_indexes + SELECT 1 FROM pg_indexes WHERE tablename = 'llm_configs' AND indexname = 'ix_llm_configs_created_at' ) THEN CREATE INDEX ix_llm_configs_created_at ON llm_configs(created_at); END IF; IF NOT EXISTS ( - SELECT 1 FROM pg_indexes + SELECT 1 FROM pg_indexes WHERE tablename = 'llm_configs' AND indexname = 'ix_llm_configs_name' ) THEN CREATE INDEX ix_llm_configs_name ON llm_configs(name); END IF; END$$; - """) + """ + ) # Safely add columns to user table bind = op.get_bind() inspector = sa.inspect(bind) existing_columns = [col["name"] for col in inspector.get_columns("user")] - with op.batch_alter_table('user') as batch_op: - if 'long_context_llm_id' not in existing_columns: - batch_op.add_column(sa.Column('long_context_llm_id', sa.Integer(), nullable=True)) - batch_op.create_foreign_key(op.f('fk_user_long_context_llm_id_llm_configs'), - 'llm_configs', ['long_context_llm_id'], ['id'], - ondelete='SET NULL') + with op.batch_alter_table("user") as batch_op: + if "long_context_llm_id" not in existing_columns: + batch_op.add_column( + sa.Column("long_context_llm_id", sa.Integer(), nullable=True) + ) + batch_op.create_foreign_key( + op.f("fk_user_long_context_llm_id_llm_configs"), + "llm_configs", + ["long_context_llm_id"], + ["id"], + ondelete="SET NULL", + ) - if 'fast_llm_id' not in existing_columns: - batch_op.add_column(sa.Column('fast_llm_id', sa.Integer(), nullable=True)) - batch_op.create_foreign_key(op.f('fk_user_fast_llm_id_llm_configs'), - 'llm_configs', ['fast_llm_id'], ['id'], - ondelete='SET NULL') + if "fast_llm_id" not in existing_columns: + batch_op.add_column(sa.Column("fast_llm_id", sa.Integer(), nullable=True)) + batch_op.create_foreign_key( + op.f("fk_user_fast_llm_id_llm_configs"), + "llm_configs", + ["fast_llm_id"], + ["id"], + ondelete="SET NULL", + ) - if 'strategic_llm_id' not in existing_columns: - batch_op.add_column(sa.Column('strategic_llm_id', sa.Integer(), nullable=True)) - batch_op.create_foreign_key(op.f('fk_user_strategic_llm_id_llm_configs'), - 'llm_configs', ['strategic_llm_id'], ['id'], - ondelete='SET NULL') + if "strategic_llm_id" not in existing_columns: + batch_op.add_column( + sa.Column("strategic_llm_id", sa.Integer(), nullable=True) + ) + batch_op.create_foreign_key( + op.f("fk_user_strategic_llm_id_llm_configs"), + "llm_configs", + ["strategic_llm_id"], + ["id"], + ondelete="SET NULL", + ) def downgrade() -> None: """Downgrade schema - remove LLMConfig table and user LLM preferences.""" - + # Drop foreign key constraints - op.drop_constraint(op.f('fk_user_strategic_llm_id_llm_configs'), 'user', type_='foreignkey') - op.drop_constraint(op.f('fk_user_fast_llm_id_llm_configs'), 'user', type_='foreignkey') - op.drop_constraint(op.f('fk_user_long_context_llm_id_llm_configs'), 'user', type_='foreignkey') - + op.drop_constraint( + op.f("fk_user_strategic_llm_id_llm_configs"), "user", type_="foreignkey" + ) + op.drop_constraint( + op.f("fk_user_fast_llm_id_llm_configs"), "user", type_="foreignkey" + ) + op.drop_constraint( + op.f("fk_user_long_context_llm_id_llm_configs"), "user", type_="foreignkey" + ) + # Drop LLM preference columns from user table - op.drop_column('user', 'strategic_llm_id') - op.drop_column('user', 'fast_llm_id') - op.drop_column('user', 'long_context_llm_id') - + op.drop_column("user", "strategic_llm_id") + op.drop_column("user", "fast_llm_id") + op.drop_column("user", "long_context_llm_id") + # Drop indexes and table - op.drop_index(op.f('ix_llm_configs_name'), table_name='llm_configs') - op.drop_index(op.f('ix_llm_configs_created_at'), table_name='llm_configs') - op.drop_index(op.f('ix_llm_configs_id'), table_name='llm_configs') - op.drop_table('llm_configs') - + op.drop_index(op.f("ix_llm_configs_name"), table_name="llm_configs") + op.drop_index(op.f("ix_llm_configs_created_at"), table_name="llm_configs") + op.drop_index(op.f("ix_llm_configs_id"), table_name="llm_configs") + op.drop_table("llm_configs") + # Drop LiteLLMProvider enum - op.execute("DROP TYPE IF EXISTS litellmprovider") \ No newline at end of file + op.execute("DROP TYPE IF EXISTS litellmprovider") diff --git a/surfsense_backend/alembic/versions/12_add_logs_table.py b/surfsense_backend/alembic/versions/12_add_logs_table.py index c74d0e8..2fc8b2b 100644 --- a/surfsense_backend/alembic/versions/12_add_logs_table.py +++ b/surfsense_backend/alembic/versions/12_add_logs_table.py @@ -7,8 +7,6 @@ Revises: 11 from typing import Sequence, Union from alembic import op -import sqlalchemy as sa -from sqlalchemy.dialects.postgresql import JSONB from sqlalchemy import inspect # revision identifiers, used by Alembic. @@ -22,27 +20,32 @@ def upgrade() -> None: """Upgrade schema - add LogLevel and LogStatus enums and logs table.""" # Create LogLevel enum if it doesn't exist - op.execute(""" + op.execute( + """ DO $$ BEGIN IF NOT EXISTS (SELECT 1 FROM pg_type WHERE typname = 'loglevel') THEN CREATE TYPE loglevel AS ENUM ('DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL'); END IF; END$$; - """) + """ + ) # Create LogStatus enum if it doesn't exist - op.execute(""" + op.execute( + """ DO $$ BEGIN IF NOT EXISTS (SELECT 1 FROM pg_type WHERE typname = 'logstatus') THEN CREATE TYPE logstatus AS ENUM ('IN_PROGRESS', 'SUCCESS', 'FAILED'); END IF; END$$; - """) + """ + ) # Create logs table if it doesn't exist - op.execute(""" + op.execute( + """ CREATE TABLE IF NOT EXISTS logs ( id SERIAL PRIMARY KEY, created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), @@ -53,38 +56,39 @@ def upgrade() -> None: log_metadata JSONB DEFAULT '{}', search_space_id INTEGER NOT NULL REFERENCES searchspaces(id) ON DELETE CASCADE ); - """) + """ + ) # Get existing indexes conn = op.get_bind() inspector = inspect(conn) - existing_indexes = [idx['name'] for idx in inspector.get_indexes('logs')] + existing_indexes = [idx["name"] for idx in inspector.get_indexes("logs")] # Create indexes only if they don't already exist - if 'ix_logs_id' not in existing_indexes: - op.create_index('ix_logs_id', 'logs', ['id']) - if 'ix_logs_created_at' not in existing_indexes: - op.create_index('ix_logs_created_at', 'logs', ['created_at']) - if 'ix_logs_level' not in existing_indexes: - op.create_index('ix_logs_level', 'logs', ['level']) - if 'ix_logs_status' not in existing_indexes: - op.create_index('ix_logs_status', 'logs', ['status']) - if 'ix_logs_source' not in existing_indexes: - op.create_index('ix_logs_source', 'logs', ['source']) + if "ix_logs_id" not in existing_indexes: + op.create_index("ix_logs_id", "logs", ["id"]) + if "ix_logs_created_at" not in existing_indexes: + op.create_index("ix_logs_created_at", "logs", ["created_at"]) + if "ix_logs_level" not in existing_indexes: + op.create_index("ix_logs_level", "logs", ["level"]) + if "ix_logs_status" not in existing_indexes: + op.create_index("ix_logs_status", "logs", ["status"]) + if "ix_logs_source" not in existing_indexes: + op.create_index("ix_logs_source", "logs", ["source"]) def downgrade() -> None: """Downgrade schema - remove logs table and enums.""" # Drop indexes - op.drop_index('ix_logs_source', table_name='logs') - op.drop_index('ix_logs_status', table_name='logs') - op.drop_index('ix_logs_level', table_name='logs') - op.drop_index('ix_logs_created_at', table_name='logs') - op.drop_index('ix_logs_id', table_name='logs') + op.drop_index("ix_logs_source", table_name="logs") + op.drop_index("ix_logs_status", table_name="logs") + op.drop_index("ix_logs_level", table_name="logs") + op.drop_index("ix_logs_created_at", table_name="logs") + op.drop_index("ix_logs_id", table_name="logs") # Drop logs table - op.drop_table('logs') + op.drop_table("logs") # Drop enums op.execute("DROP TYPE IF EXISTS logstatus") diff --git a/surfsense_backend/alembic/versions/13_add_jira_connector_enums.py b/surfsense_backend/alembic/versions/13_add_jira_connector_enums.py index 2623cb9..78d95f1 100644 --- a/surfsense_backend/alembic/versions/13_add_jira_connector_enums.py +++ b/surfsense_backend/alembic/versions/13_add_jira_connector_enums.py @@ -5,12 +5,12 @@ Revises: 12 """ from typing import Sequence, Union + from alembic import op -import sqlalchemy as sa # revision identifiers, used by Alembic. -revision: str = '13' -down_revision: Union[str, None] = '12' +revision: str = "13" +down_revision: Union[str, None] = "12" branch_labels: Union[str, Sequence[str], None] = None depends_on: Union[str, Sequence[str], None] = None @@ -19,7 +19,8 @@ def upgrade() -> None: """Safely add 'JIRA_CONNECTOR' to enum types if missing.""" # Add to searchsourceconnectortype enum - op.execute(""" + op.execute( + """ DO $$ BEGIN IF NOT EXISTS ( @@ -31,10 +32,12 @@ def upgrade() -> None: END IF; END $$; - """) + """ + ) # Add to documenttype enum - op.execute(""" + op.execute( + """ DO $$ BEGIN IF NOT EXISTS ( @@ -46,7 +49,8 @@ def upgrade() -> None: END IF; END $$; - """) + """ + ) def downgrade() -> None: diff --git a/surfsense_backend/alembic/versions/1_add_github_connector_enum.py b/surfsense_backend/alembic/versions/1_add_github_connector_enum.py index 713621f..d094c99 100644 --- a/surfsense_backend/alembic/versions/1_add_github_connector_enum.py +++ b/surfsense_backend/alembic/versions/1_add_github_connector_enum.py @@ -1,19 +1,20 @@ """Add GITHUB_CONNECTOR to SearchSourceConnectorType enum Revision ID: 1 -Revises: +Revises: """ + from typing import Sequence, Union from alembic import op -import sqlalchemy as sa + # Import pgvector if needed for other types, though not for this ENUM change # import pgvector # revision identifiers, used by Alembic. -revision: str = '1' +revision: str = "1" down_revision: Union[str, None] = None branch_labels: Union[str, Sequence[str], None] = None depends_on: Union[str, Sequence[str], None] = None @@ -21,10 +22,11 @@ depends_on: Union[str, Sequence[str], None] = None def upgrade() -> None: # ### commands auto generated by Alembic - please adjust! ### - + # Manually add the command to add the enum value # Note: It's generally better to let autogenerate handle this, but we're bypassing it - op.execute(""" + op.execute( + """ DO $$ BEGIN IF NOT EXISTS ( @@ -38,9 +40,9 @@ BEGIN ALTER TYPE searchsourceconnectortype ADD VALUE 'GITHUB_CONNECTOR'; END IF; END$$; -""") +""" + ) - # Pass for the rest, as autogenerate didn't run to add other schema details pass # ### end Alembic commands ### @@ -48,20 +50,25 @@ END$$; def downgrade() -> None: # ### commands auto generated by Alembic - please adjust! ### - + # Downgrading removal of an enum value is complex and potentially dangerous # if the value is in use. Often omitted or requires manual SQL based on context. - # For now, we'll just pass. If you needed to reverse this, you'd likely + # For now, we'll just pass. If you needed to reverse this, you'd likely # have to manually check if 'GITHUB_CONNECTOR' is used in the table # and then potentially recreate the type without it. - op.execute("ALTER TYPE searchsourceconnectortype RENAME TO searchsourceconnectortype_old") - op.execute("CREATE TYPE searchsourceconnectortype AS ENUM('SERPER_API', 'TAVILY_API', 'SLACK_CONNECTOR', 'NOTION_CONNECTOR')") - op.execute(( - "ALTER TABLE search_source_connectors ALTER COLUMN connector_type TYPE searchsourceconnectortype USING " - "connector_type::text::searchsourceconnectortype" - )) + op.execute( + "ALTER TYPE searchsourceconnectortype RENAME TO searchsourceconnectortype_old" + ) + op.execute( + "CREATE TYPE searchsourceconnectortype AS ENUM('SERPER_API', 'TAVILY_API', 'SLACK_CONNECTOR', 'NOTION_CONNECTOR')" + ) + op.execute( + ( + "ALTER TABLE search_source_connectors ALTER COLUMN connector_type TYPE searchsourceconnectortype USING " + "connector_type::text::searchsourceconnectortype" + ) + ) op.execute("DROP TYPE searchsourceconnectortype_old") - pass - # ### end Alembic commands ### + # ### end Alembic commands ### diff --git a/surfsense_backend/alembic/versions/2_add_linear_connector_enum.py b/surfsense_backend/alembic/versions/2_add_linear_connector_enum.py index aa0ac3c..31fcee8 100644 --- a/surfsense_backend/alembic/versions/2_add_linear_connector_enum.py +++ b/surfsense_backend/alembic/versions/2_add_linear_connector_enum.py @@ -4,21 +4,21 @@ Revision ID: 2 Revises: e55302644c51 """ + from typing import Sequence, Union from alembic import op -import sqlalchemy as sa - # revision identifiers, used by Alembic. -revision: str = '2' -down_revision: Union[str, None] = 'e55302644c51' +revision: str = "2" +down_revision: Union[str, None] = "e55302644c51" branch_labels: Union[str, Sequence[str], None] = None depends_on: Union[str, Sequence[str], None] = None def upgrade() -> None: - op.execute(""" + op.execute( + """ DO $$ BEGIN IF NOT EXISTS ( @@ -31,21 +31,30 @@ def upgrade() -> None: ALTER TYPE searchsourceconnectortype ADD VALUE 'LINEAR_CONNECTOR'; END IF; END$$; - """) + """ + ) + + # def downgrade() -> None: # ### commands auto generated by Alembic - please adjust! ### - + # Downgrading removal of an enum value requires recreating the type - op.execute("ALTER TYPE searchsourceconnectortype RENAME TO searchsourceconnectortype_old") - op.execute("CREATE TYPE searchsourceconnectortype AS ENUM('SERPER_API', 'TAVILY_API', 'SLACK_CONNECTOR', 'NOTION_CONNECTOR', 'GITHUB_CONNECTOR')") - op.execute(( - "ALTER TABLE search_source_connectors ALTER COLUMN connector_type TYPE searchsourceconnectortype USING " - "connector_type::text::searchsourceconnectortype" - )) + op.execute( + "ALTER TYPE searchsourceconnectortype RENAME TO searchsourceconnectortype_old" + ) + op.execute( + "CREATE TYPE searchsourceconnectortype AS ENUM('SERPER_API', 'TAVILY_API', 'SLACK_CONNECTOR', 'NOTION_CONNECTOR', 'GITHUB_CONNECTOR')" + ) + op.execute( + ( + "ALTER TABLE search_source_connectors ALTER COLUMN connector_type TYPE searchsourceconnectortype USING " + "connector_type::text::searchsourceconnectortype" + ) + ) op.execute("DROP TYPE searchsourceconnectortype_old") pass - # ### end Alembic commands ### \ No newline at end of file + # ### end Alembic commands ### diff --git a/surfsense_backend/alembic/versions/3_add_linear_connector_to_documenttype_.py b/surfsense_backend/alembic/versions/3_add_linear_connector_to_documenttype_.py index 2b326f2..b108699 100644 --- a/surfsense_backend/alembic/versions/3_add_linear_connector_to_documenttype_.py +++ b/surfsense_backend/alembic/versions/3_add_linear_connector_to_documenttype_.py @@ -4,25 +4,26 @@ Revision ID: 3 Revises: 2 """ + from typing import Sequence, Union from alembic import op -import sqlalchemy as sa - # revision identifiers, used by Alembic. -revision: str = '3' -down_revision: Union[str, None] = '2' +revision: str = "3" +down_revision: Union[str, None] = "2" branch_labels: Union[str, Sequence[str], None] = None depends_on: Union[str, Sequence[str], None] = None # Define the ENUM type name and the new value -ENUM_NAME = 'documenttype' # Make sure this matches the name in your DB (usually lowercase class name) -NEW_VALUE = 'LINEAR_CONNECTOR' +ENUM_NAME = "documenttype" # Make sure this matches the name in your DB (usually lowercase class name) +NEW_VALUE = "LINEAR_CONNECTOR" + def upgrade() -> None: """Upgrade schema.""" - op.execute(f""" + op.execute( + f""" DO $$ BEGIN IF NOT EXISTS ( @@ -35,9 +36,9 @@ def upgrade() -> None: ALTER TYPE {ENUM_NAME} ADD VALUE '{NEW_VALUE}'; END IF; END$$; - """) + """ + ) - # Warning: This will delete all rows with the new value def downgrade() -> None: @@ -48,19 +49,19 @@ def downgrade() -> None: # Enum values *before* LINEAR_CONNECTOR was added old_values = ( - 'EXTENSION', - 'CRAWLED_URL', - 'FILE', - 'SLACK_CONNECTOR', - 'NOTION_CONNECTOR', - 'YOUTUBE_VIDEO', - 'GITHUB_CONNECTOR' + "EXTENSION", + "CRAWLED_URL", + "FILE", + "SLACK_CONNECTOR", + "NOTION_CONNECTOR", + "YOUTUBE_VIDEO", + "GITHUB_CONNECTOR", ) old_values_sql = ", ".join([f"'{v}'" for v in old_values]) # Table and column names (adjust if different) - table_name = 'documents' - column_name = 'document_type' + table_name = "documents" + column_name = "document_type" # 1. Rename the current enum type op.execute(f"ALTER TYPE {ENUM_NAME} RENAME TO {old_enum_name}") @@ -68,10 +69,8 @@ def downgrade() -> None: # 2. Create the new enum type with the old values op.execute(f"CREATE TYPE {ENUM_NAME} AS ENUM({old_values_sql})") - # 3. Update the table: - op.execute( - f"DELETE FROM {table_name} WHERE {column_name}::text = '{NEW_VALUE}'" - ) + # 3. Update the table: + op.execute(f"DELETE FROM {table_name} WHERE {column_name}::text = '{NEW_VALUE}'") # 4. Alter the column to use the new enum type (casting old values) op.execute( @@ -81,4 +80,4 @@ def downgrade() -> None: # 5. Drop the old enum type op.execute(f"DROP TYPE {old_enum_name}") - # ### end Alembic commands ### \ No newline at end of file + # ### end Alembic commands ### diff --git a/surfsense_backend/alembic/versions/4_add_linkup_api_enum.py b/surfsense_backend/alembic/versions/4_add_linkup_api_enum.py index 45d4ecf..6720ae7 100644 --- a/surfsense_backend/alembic/versions/4_add_linkup_api_enum.py +++ b/surfsense_backend/alembic/versions/4_add_linkup_api_enum.py @@ -4,24 +4,24 @@ Revision ID: 4 Revises: 3 """ + from typing import Sequence, Union from alembic import op -import sqlalchemy as sa - # revision identifiers, used by Alembic. -revision: str = '4' -down_revision: Union[str, None] = '3' +revision: str = "4" +down_revision: Union[str, None] = "3" branch_labels: Union[str, Sequence[str], None] = None depends_on: Union[str, Sequence[str], None] = None def upgrade() -> None: - ENUM_NAME = 'searchsourceconnectortype' - NEW_VALUE = 'LINKUP_API' + ENUM_NAME = "searchsourceconnectortype" + NEW_VALUE = "LINKUP_API" - op.execute(f""" + op.execute( + f""" DO $$ BEGIN IF NOT EXISTS ( @@ -34,21 +34,27 @@ def upgrade() -> None: ALTER TYPE {ENUM_NAME} ADD VALUE '{NEW_VALUE}'; END IF; END$$; - """) - + """ + ) def downgrade() -> None: # ### commands auto generated by Alembic - please adjust! ### - + # Downgrading removal of an enum value requires recreating the type - op.execute("ALTER TYPE searchsourceconnectortype RENAME TO searchsourceconnectortype_old") - op.execute("CREATE TYPE searchsourceconnectortype AS ENUM('SERPER_API', 'TAVILY_API', 'SLACK_CONNECTOR', 'NOTION_CONNECTOR', 'GITHUB_CONNECTOR', 'LINEAR_CONNECTOR')") - op.execute(( - "ALTER TABLE search_source_connectors ALTER COLUMN connector_type TYPE searchsourceconnectortype USING " - "connector_type::text::searchsourceconnectortype" - )) + op.execute( + "ALTER TYPE searchsourceconnectortype RENAME TO searchsourceconnectortype_old" + ) + op.execute( + "CREATE TYPE searchsourceconnectortype AS ENUM('SERPER_API', 'TAVILY_API', 'SLACK_CONNECTOR', 'NOTION_CONNECTOR', 'GITHUB_CONNECTOR', 'LINEAR_CONNECTOR')" + ) + op.execute( + ( + "ALTER TABLE search_source_connectors ALTER COLUMN connector_type TYPE searchsourceconnectortype USING " + "connector_type::text::searchsourceconnectortype" + ) + ) op.execute("DROP TYPE searchsourceconnectortype_old") pass - # ### end Alembic commands ### \ No newline at end of file + # ### end Alembic commands ### diff --git a/surfsense_backend/alembic/versions/6_change_podcast_content_to_transcript.py b/surfsense_backend/alembic/versions/6_change_podcast_content_to_transcript.py index edc65b7..3f0865f 100644 --- a/surfsense_backend/alembic/versions/6_change_podcast_content_to_transcript.py +++ b/surfsense_backend/alembic/versions/6_change_podcast_content_to_transcript.py @@ -4,48 +4,58 @@ Revision ID: 6 Revises: 5 """ + from typing import Sequence, Union -from alembic import op import sqlalchemy as sa -from sqlalchemy.dialects.postgresql import JSON +from alembic import op from sqlalchemy import inspect - +from sqlalchemy.dialects.postgresql import JSON # revision identifiers, used by Alembic. -revision: str = '6' -down_revision: Union[str, None] = '5' +revision: str = "6" +down_revision: Union[str, None] = "5" branch_labels: Union[str, Sequence[str], None] = None depends_on: Union[str, Sequence[str], None] = None + def upgrade() -> None: bind = op.get_bind() inspector = inspect(bind) columns = [col["name"] for col in inspector.get_columns("podcasts")] if "podcast_transcript" not in columns: - op.add_column('podcasts', sa.Column('podcast_transcript', JSON, nullable=False, server_default='{}')) + op.add_column( + "podcasts", + sa.Column("podcast_transcript", JSON, nullable=False, server_default="{}"), + ) # Copy data from old column to new column - op.execute(""" - UPDATE podcasts - SET podcast_transcript = jsonb_build_object('text', podcast_content) + op.execute( + """ + UPDATE podcasts + SET podcast_transcript = jsonb_build_object('text', podcast_content) WHERE podcast_content != '' - """) + """ + ) # Drop the old column only if it exists if "podcast_content" in columns: - op.drop_column('podcasts', 'podcast_content') - + op.drop_column("podcasts", "podcast_content") def downgrade() -> None: # Add back the original column - op.add_column('podcasts', sa.Column('podcast_content', sa.Text(), nullable=False, server_default='')) - + op.add_column( + "podcasts", + sa.Column("podcast_content", sa.Text(), nullable=False, server_default=""), + ) + # Copy data from JSON column back to text column # Extract the 'text' field if it exists, otherwise use empty string - op.execute("UPDATE podcasts SET podcast_content = COALESCE((podcast_transcript->>'text'), '')") - + op.execute( + "UPDATE podcasts SET podcast_content = COALESCE((podcast_transcript->>'text'), '')" + ) + # Drop the new column - op.drop_column('podcasts', 'podcast_transcript') \ No newline at end of file + op.drop_column("podcasts", "podcast_transcript") diff --git a/surfsense_backend/alembic/versions/7_remove_is_generated_column.py b/surfsense_backend/alembic/versions/7_remove_is_generated_column.py index 9fa4546..62b273b 100644 --- a/surfsense_backend/alembic/versions/7_remove_is_generated_column.py +++ b/surfsense_backend/alembic/versions/7_remove_is_generated_column.py @@ -4,19 +4,20 @@ Revision ID: 7 Revises: 6 """ + from typing import Sequence, Union -from alembic import op import sqlalchemy as sa +from alembic import op from sqlalchemy import inspect - # revision identifiers, used by Alembic. -revision: str = '7' -down_revision: Union[str, None] = '6' +revision: str = "7" +down_revision: Union[str, None] = "6" branch_labels: Union[str, Sequence[str], None] = None depends_on: Union[str, Sequence[str], None] = None + def upgrade() -> None: # Get the current database connection bind = op.get_bind() @@ -25,10 +26,12 @@ def upgrade() -> None: # Check if the column exists before attempting to drop it columns = [col["name"] for col in inspector.get_columns("podcasts")] if "is_generated" in columns: - op.drop_column('podcasts', 'is_generated') - + op.drop_column("podcasts", "is_generated") def downgrade() -> None: # Add back the is_generated column with its original constraints - op.add_column('podcasts', sa.Column('is_generated', sa.Boolean(), nullable=False, server_default='false')) \ No newline at end of file + op.add_column( + "podcasts", + sa.Column("is_generated", sa.Boolean(), nullable=False, server_default="false"), + ) diff --git a/surfsense_backend/alembic/versions/8_add_content_hash_to_documents.py b/surfsense_backend/alembic/versions/8_add_content_hash_to_documents.py index ed7ebc5..976c6d3 100644 --- a/surfsense_backend/alembic/versions/8_add_content_hash_to_documents.py +++ b/surfsense_backend/alembic/versions/8_add_content_hash_to_documents.py @@ -3,15 +3,16 @@ Revision ID: 8 Revises: 7 """ + from typing import Sequence, Union -from alembic import op + import sqlalchemy as sa +from alembic import op from sqlalchemy import inspect - # revision identifiers, used by Alembic. -revision: str = '8' -down_revision: Union[str, None] = '7' +revision: str = "8" +down_revision: Union[str, None] = "7" branch_labels: Union[str, Sequence[str], None] = None depends_on: Union[str, Sequence[str], None] = None @@ -19,38 +20,51 @@ depends_on: Union[str, Sequence[str], None] = None def upgrade() -> None: bind = op.get_bind() inspector = inspect(bind) - columns = [col['name'] for col in inspector.get_columns('documents')] + columns = [col["name"] for col in inspector.get_columns("documents")] # Only add the column if it doesn't already exist - if 'content_hash' not in columns: - op.add_column('documents', sa.Column('content_hash', sa.String(), nullable=True)) + if "content_hash" not in columns: + op.add_column( + "documents", sa.Column("content_hash", sa.String(), nullable=True) + ) # Populate the content_hash column - op.execute(""" - UPDATE documents + op.execute( + """ + UPDATE documents SET content_hash = encode(sha256(convert_to(content, 'UTF8')), 'hex') WHERE content_hash IS NULL - """) + """ + ) - op.execute(""" - DELETE FROM documents + op.execute( + """ + DELETE FROM documents WHERE id NOT IN ( - SELECT MIN(id) - FROM documents + SELECT MIN(id) + FROM documents GROUP BY content_hash ) - """) + """ + ) - op.alter_column('documents', 'content_hash', - existing_type=sa.String(), - nullable=False) - op.create_index(op.f('ix_documents_content_hash'), 'documents', ['content_hash'], unique=False) - op.create_unique_constraint(op.f('uq_documents_content_hash'), 'documents', ['content_hash']) + op.alter_column( + "documents", "content_hash", existing_type=sa.String(), nullable=False + ) + op.create_index( + op.f("ix_documents_content_hash"), + "documents", + ["content_hash"], + unique=False, + ) + op.create_unique_constraint( + op.f("uq_documents_content_hash"), "documents", ["content_hash"] + ) else: print("Column 'content_hash' already exists. Skipping column creation.") def downgrade() -> None: - op.drop_constraint(op.f('uq_documents_content_hash'), 'documents', type_='unique') - op.drop_index(op.f('ix_documents_content_hash'), table_name='documents') - op.drop_column('documents', 'content_hash') + op.drop_constraint(op.f("uq_documents_content_hash"), "documents", type_="unique") + op.drop_index(op.f("ix_documents_content_hash"), table_name="documents") + op.drop_column("documents", "content_hash") diff --git a/surfsense_backend/alembic/versions/9_add_discord_connector_enum_and_documenttype.py b/surfsense_backend/alembic/versions/9_add_discord_connector_enum_and_documenttype.py index ad77dad..4dec112 100644 --- a/surfsense_backend/alembic/versions/9_add_discord_connector_enum_and_documenttype.py +++ b/surfsense_backend/alembic/versions/9_add_discord_connector_enum_and_documenttype.py @@ -7,8 +7,6 @@ Revises: 8 from typing import Sequence, Union from alembic import op -import sqlalchemy as sa - # revision identifiers, used by Alembic. revision: str = "9" @@ -26,7 +24,8 @@ DOCUMENT_NEW_VALUE = "DISCORD_CONNECTOR" def upgrade() -> None: """Upgrade schema - add DISCORD_CONNECTOR to connector and document enum safely.""" # Add DISCORD_CONNECTOR to searchsourceconnectortype only if not exists - op.execute(f""" + op.execute( + f""" DO $$ BEGIN IF NOT EXISTS ( @@ -37,10 +36,12 @@ def upgrade() -> None: ALTER TYPE {CONNECTOR_ENUM} ADD VALUE '{CONNECTOR_NEW_VALUE}'; END IF; END$$; - """) + """ + ) # Add DISCORD_CONNECTOR to documenttype only if not exists - op.execute(f""" + op.execute( + f""" DO $$ BEGIN IF NOT EXISTS ( @@ -51,8 +52,8 @@ def upgrade() -> None: ALTER TYPE {DOCUMENT_ENUM} ADD VALUE '{DOCUMENT_NEW_VALUE}'; END IF; END$$; - """) - + """ + ) def downgrade() -> None: @@ -109,7 +110,6 @@ def downgrade() -> None: # 4. Drop the old connector enum type op.execute(f"DROP TYPE {old_connector_enum_name}") - # Document Enum Downgrade Steps # 1. Rename the current document enum type op.execute(f"ALTER TYPE {DOCUMENT_ENUM} RENAME TO {old_document_enum_name}") diff --git a/surfsense_backend/alembic/versions/e55302644c51_add_github_connector_to_documenttype_.py b/surfsense_backend/alembic/versions/e55302644c51_add_github_connector_to_documenttype_.py index ae0a60a..9c93eb7 100644 --- a/surfsense_backend/alembic/versions/e55302644c51_add_github_connector_to_documenttype_.py +++ b/surfsense_backend/alembic/versions/e55302644c51_add_github_connector_to_documenttype_.py @@ -1,21 +1,22 @@ from typing import Sequence, Union from alembic import op -import sqlalchemy as sa # revision identifiers, used by Alembic. -revision: str = 'e55302644c51' -down_revision: Union[str, None] = '1' +revision: str = "e55302644c51" +down_revision: Union[str, None] = "1" branch_labels: Union[str, Sequence[str], None] = None depends_on: Union[str, Sequence[str], None] = None # Define the ENUM type name and the new value -ENUM_NAME = 'documenttype' -NEW_VALUE = 'GITHUB_CONNECTOR' +ENUM_NAME = "documenttype" +NEW_VALUE = "GITHUB_CONNECTOR" + def upgrade() -> None: """Upgrade schema.""" - op.execute(f""" + op.execute( + f""" DO $$ BEGIN IF NOT EXISTS ( @@ -28,32 +29,32 @@ def upgrade() -> None: ALTER TYPE {ENUM_NAME} ADD VALUE '{NEW_VALUE}'; END IF; END$$; - """) + """ + ) + def downgrade() -> None: """Downgrade schema - remove GITHUB_CONNECTOR from enum.""" old_enum_name = f"{ENUM_NAME}_old" old_values = ( - 'EXTENSION', - 'CRAWLED_URL', - 'FILE', - 'SLACK_CONNECTOR', - 'NOTION_CONNECTOR', - 'YOUTUBE_VIDEO' + "EXTENSION", + "CRAWLED_URL", + "FILE", + "SLACK_CONNECTOR", + "NOTION_CONNECTOR", + "YOUTUBE_VIDEO", ) old_values_sql = ", ".join([f"'{v}'" for v in old_values]) - table_name = 'documents' - column_name = 'document_type' + table_name = "documents" + column_name = "document_type" # 1. Create the new enum type with the old values op.execute(f"CREATE TYPE {old_enum_name} AS ENUM({old_values_sql})") # 2. Delete rows using the new value - op.execute( - f"DELETE FROM {table_name} WHERE {column_name}::text = '{NEW_VALUE}'" - ) + op.execute(f"DELETE FROM {table_name} WHERE {column_name}::text = '{NEW_VALUE}'") # 3. Alter the column to use the old enum type op.execute(