vrr/SurfSense
2
0
Fork 0
mirror of https://github.com/MODSetter/SurfSense.git synced 2025-09-04 19:49:09 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 3

generate new migration / fix migration files

Browse source
This commit is contained in:
CREDO23 2025-07-24 11:38:49 +02:00
parent b2a19af1f7
commit a6fe7e583b
12 changed files with 303 additions and 220 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

96
surfsense_backend/alembic/versions/11_add_llm_config_table_and_relationships.py
View file

@ -6,10 +6,8 @@ Revises: 10
from typing import Sequence, Union from typing import Sequence, Union
from alembic import op
import sqlalchemy as sa import sqlalchemy as sa
from sqlalchemy.dialects.postgresql import UUID, JSON from alembic import op
# revision identifiers, used by Alembic. # revision identifiers, used by Alembic.
revision: str = "11" revision: str = "11"
@ -22,7 +20,8 @@ def upgrade() -> None:
"""Upgrade schema - add LiteLLMProvider enum, LLMConfig table and user LLM preferences.""" """Upgrade schema - add LiteLLMProvider enum, LLMConfig table and user LLM preferences."""
# Create enum only if not exists # Create enum only if not exists
op.execute(""" op.execute(
"""
DO $$ DO $$
BEGIN BEGIN
IF NOT EXISTS (SELECT 1 FROM pg_type WHERE typname = 'litellmprovider') THEN IF NOT EXISTS (SELECT 1 FROM pg_type WHERE typname = 'litellmprovider') THEN
@ -35,10 +34,12 @@ def upgrade() -> None:
); );
END IF; END IF;
END$$; END$$;
""") """
)
# Create llm_configs table only if it doesn't already exist # Create llm_configs table only if it doesn't already exist
op.execute(""" op.execute(
"""
DO $$ DO $$
BEGIN BEGIN
IF NOT EXISTS ( IF NOT EXISTS (
@ -59,10 +60,12 @@ def upgrade() -> None:
); );
END IF; END IF;
END$$; END$$;
""") """
)
# Create indexes if they don't exist # Create indexes if they don't exist
op.execute(""" op.execute(
"""
DO $$ DO $$
BEGIN BEGIN
IF NOT EXISTS ( IF NOT EXISTS (
@ -86,51 +89,74 @@ def upgrade() -> None:
CREATE INDEX ix_llm_configs_name ON llm_configs(name); CREATE INDEX ix_llm_configs_name ON llm_configs(name);
END IF; END IF;
END$$; END$$;
""") """
)
# Safely add columns to user table # Safely add columns to user table
bind = op.get_bind() bind = op.get_bind()
inspector = sa.inspect(bind) inspector = sa.inspect(bind)
existing_columns = [col["name"] for col in inspector.get_columns("user")] existing_columns = [col["name"] for col in inspector.get_columns("user")]
with op.batch_alter_table('user') as batch_op: with op.batch_alter_table("user") as batch_op:
if 'long_context_llm_id' not in existing_columns: 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.add_column(
batch_op.create_foreign_key(op.f('fk_user_long_context_llm_id_llm_configs'), sa.Column("long_context_llm_id", sa.Integer(), nullable=True)
'llm_configs', ['long_context_llm_id'], ['id'], )
ondelete='SET NULL') 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: if "fast_llm_id" not in existing_columns:
batch_op.add_column(sa.Column('fast_llm_id', sa.Integer(), nullable=True)) 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'), batch_op.create_foreign_key(
'llm_configs', ['fast_llm_id'], ['id'], op.f("fk_user_fast_llm_id_llm_configs"),
ondelete='SET NULL') "llm_configs",
["fast_llm_id"],
["id"],
ondelete="SET NULL",
)
if 'strategic_llm_id' not in existing_columns: if "strategic_llm_id" not in existing_columns:
batch_op.add_column(sa.Column('strategic_llm_id', sa.Integer(), nullable=True)) batch_op.add_column(
batch_op.create_foreign_key(op.f('fk_user_strategic_llm_id_llm_configs'), sa.Column("strategic_llm_id", sa.Integer(), nullable=True)
'llm_configs', ['strategic_llm_id'], ['id'], )
ondelete='SET NULL') 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: def downgrade() -> None:
"""Downgrade schema - remove LLMConfig table and user LLM preferences.""" """Downgrade schema - remove LLMConfig table and user LLM preferences."""
# Drop foreign key constraints # Drop foreign key constraints
op.drop_constraint(op.f('fk_user_strategic_llm_id_llm_configs'), 'user', type_='foreignkey') op.drop_constraint(
op.drop_constraint(op.f('fk_user_fast_llm_id_llm_configs'), 'user', type_='foreignkey') op.f("fk_user_strategic_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_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 # Drop LLM preference columns from user table
op.drop_column('user', 'strategic_llm_id') op.drop_column("user", "strategic_llm_id")
op.drop_column('user', 'fast_llm_id') op.drop_column("user", "fast_llm_id")
op.drop_column('user', 'long_context_llm_id') op.drop_column("user", "long_context_llm_id")
# Drop indexes and table # 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_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_created_at"), table_name="llm_configs")
op.drop_index(op.f('ix_llm_configs_id'), table_name='llm_configs') op.drop_index(op.f("ix_llm_configs_id"), table_name="llm_configs")
op.drop_table('llm_configs') op.drop_table("llm_configs")
# Drop LiteLLMProvider enum # Drop LiteLLMProvider enum
op.execute("DROP TYPE IF EXISTS litellmprovider") op.execute("DROP TYPE IF EXISTS litellmprovider")

54
surfsense_backend/alembic/versions/12_add_logs_table.py
View file

@ -7,8 +7,6 @@ Revises: 11
from typing import Sequence, Union from typing import Sequence, Union
from alembic import op from alembic import op
import sqlalchemy as sa
from sqlalchemy.dialects.postgresql import JSONB
from sqlalchemy import inspect from sqlalchemy import inspect
# revision identifiers, used by Alembic. # revision identifiers, used by Alembic.
@ -22,27 +20,32 @@ def upgrade() -> None:
"""Upgrade schema - add LogLevel and LogStatus enums and logs table.""" """Upgrade schema - add LogLevel and LogStatus enums and logs table."""
# Create LogLevel enum if it doesn't exist # Create LogLevel enum if it doesn't exist
op.execute(""" op.execute(
"""
DO $$ DO $$
BEGIN BEGIN
IF NOT EXISTS (SELECT 1 FROM pg_type WHERE typname = 'loglevel') THEN IF NOT EXISTS (SELECT 1 FROM pg_type WHERE typname = 'loglevel') THEN
CREATE TYPE loglevel AS ENUM ('DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL'); CREATE TYPE loglevel AS ENUM ('DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL');
END IF; END IF;
END$$; END$$;
""") """
)
# Create LogStatus enum if it doesn't exist # Create LogStatus enum if it doesn't exist
op.execute(""" op.execute(
"""
DO $$ DO $$
BEGIN BEGIN
IF NOT EXISTS (SELECT 1 FROM pg_type WHERE typname = 'logstatus') THEN IF NOT EXISTS (SELECT 1 FROM pg_type WHERE typname = 'logstatus') THEN
CREATE TYPE logstatus AS ENUM ('IN_PROGRESS', 'SUCCESS', 'FAILED'); CREATE TYPE logstatus AS ENUM ('IN_PROGRESS', 'SUCCESS', 'FAILED');
END IF; END IF;
END$$; END$$;
""") """
)
# Create logs table if it doesn't exist # Create logs table if it doesn't exist
op.execute(""" op.execute(
"""
CREATE TABLE IF NOT EXISTS logs ( CREATE TABLE IF NOT EXISTS logs (
id SERIAL PRIMARY KEY, id SERIAL PRIMARY KEY,
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
@ -53,38 +56,39 @@ def upgrade() -> None:
log_metadata JSONB DEFAULT '{}', log_metadata JSONB DEFAULT '{}',
search_space_id INTEGER NOT NULL REFERENCES searchspaces(id) ON DELETE CASCADE search_space_id INTEGER NOT NULL REFERENCES searchspaces(id) ON DELETE CASCADE
); );
""") """
)
# Get existing indexes # Get existing indexes
conn = op.get_bind() conn = op.get_bind()
inspector = inspect(conn) 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 # Create indexes only if they don't already exist
if 'ix_logs_id' not in existing_indexes: if "ix_logs_id" not in existing_indexes:
op.create_index('ix_logs_id', 'logs', ['id']) op.create_index("ix_logs_id", "logs", ["id"])
if 'ix_logs_created_at' not in existing_indexes: if "ix_logs_created_at" not in existing_indexes:
op.create_index('ix_logs_created_at', 'logs', ['created_at']) op.create_index("ix_logs_created_at", "logs", ["created_at"])
if 'ix_logs_level' not in existing_indexes: if "ix_logs_level" not in existing_indexes:
op.create_index('ix_logs_level', 'logs', ['level']) op.create_index("ix_logs_level", "logs", ["level"])
if 'ix_logs_status' not in existing_indexes: if "ix_logs_status" not in existing_indexes:
op.create_index('ix_logs_status', 'logs', ['status']) op.create_index("ix_logs_status", "logs", ["status"])
if 'ix_logs_source' not in existing_indexes: if "ix_logs_source" not in existing_indexes:
op.create_index('ix_logs_source', 'logs', ['source']) op.create_index("ix_logs_source", "logs", ["source"])
def downgrade() -> None: def downgrade() -> None:
"""Downgrade schema - remove logs table and enums.""" """Downgrade schema - remove logs table and enums."""
# Drop indexes # Drop indexes
op.drop_index('ix_logs_source', 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_status", table_name="logs")
op.drop_index('ix_logs_level', 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_created_at", table_name="logs")
op.drop_index('ix_logs_id', table_name='logs') op.drop_index("ix_logs_id", table_name="logs")
# Drop logs table # Drop logs table
op.drop_table('logs') op.drop_table("logs")
# Drop enums # Drop enums
op.execute("DROP TYPE IF EXISTS logstatus") op.execute("DROP TYPE IF EXISTS logstatus")

18
surfsense_backend/alembic/versions/13_add_jira_connector_enums.py
View file

@ -5,12 +5,12 @@ Revises: 12
""" """
from typing import Sequence, Union from typing import Sequence, Union
from alembic import op from alembic import op
import sqlalchemy as sa
# revision identifiers, used by Alembic. # revision identifiers, used by Alembic.
revision: str = '13' revision: str = "13"
down_revision: Union[str, None] = '12' down_revision: Union[str, None] = "12"
branch_labels: Union[str, Sequence[str], None] = None branch_labels: Union[str, Sequence[str], None] = None
depends_on: 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.""" """Safely add 'JIRA_CONNECTOR' to enum types if missing."""
# Add to searchsourceconnectortype enum # Add to searchsourceconnectortype enum
op.execute(""" op.execute(
"""
DO $$ DO $$
BEGIN BEGIN
IF NOT EXISTS ( IF NOT EXISTS (
@ -31,10 +32,12 @@ def upgrade() -> None:
END IF; END IF;
END END
$$; $$;
""") """
)
# Add to documenttype enum # Add to documenttype enum
op.execute(""" op.execute(
"""
DO $$ DO $$
BEGIN BEGIN
IF NOT EXISTS ( IF NOT EXISTS (
@ -46,7 +49,8 @@ def upgrade() -> None:
END IF; END IF;
END END
$$; $$;
""") """
)
def downgrade() -> None: def downgrade() -> None:

31
surfsense_backend/alembic/versions/1_add_github_connector_enum.py
View file

@ -4,16 +4,17 @@ Revision ID: 1
Revises: Revises:
""" """
from typing import Sequence, Union from typing import Sequence, Union
from alembic import op from alembic import op
import sqlalchemy as sa
# Import pgvector if needed for other types, though not for this ENUM change # Import pgvector if needed for other types, though not for this ENUM change
# import pgvector # import pgvector
# revision identifiers, used by Alembic. # revision identifiers, used by Alembic.
revision: str = '1' revision: str = "1"
down_revision: Union[str, None] = None down_revision: Union[str, None] = None
branch_labels: Union[str, Sequence[str], None] = None branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None depends_on: Union[str, Sequence[str], None] = None
@ -24,7 +25,8 @@ def upgrade() -> None:
# Manually add the command to add the enum value # Manually add the command to add the enum value
# Note: It's generally better to let autogenerate handle this, but we're bypassing it # Note: It's generally better to let autogenerate handle this, but we're bypassing it
op.execute(""" op.execute(
"""
DO $$ DO $$
BEGIN BEGIN
IF NOT EXISTS ( IF NOT EXISTS (
@ -38,8 +40,8 @@ BEGIN
ALTER TYPE searchsourceconnectortype ADD VALUE 'GITHUB_CONNECTOR'; ALTER TYPE searchsourceconnectortype ADD VALUE 'GITHUB_CONNECTOR';
END IF; END IF;
END$$; END$$;
""") """
)
# Pass for the rest, as autogenerate didn't run to add other schema details # Pass for the rest, as autogenerate didn't run to add other schema details
pass pass
@ -54,14 +56,19 @@ def downgrade() -> None:
# 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 # have to manually check if 'GITHUB_CONNECTOR' is used in the table
# and then potentially recreate the type without it. # and then potentially recreate the type without it.
op.execute("ALTER TYPE searchsourceconnectortype RENAME TO searchsourceconnectortype_old") op.execute(
op.execute("CREATE TYPE searchsourceconnectortype AS ENUM('SERPER_API', 'TAVILY_API', 'SLACK_CONNECTOR', 'NOTION_CONNECTOR')") "ALTER TYPE searchsourceconnectortype RENAME TO searchsourceconnectortype_old"
op.execute(( )
"ALTER TABLE search_source_connectors ALTER COLUMN connector_type TYPE searchsourceconnectortype USING " op.execute(
"connector_type::text::searchsourceconnectortype" "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") op.execute("DROP TYPE searchsourceconnectortype_old")
pass pass
# ### end Alembic commands ### # ### end Alembic commands ###

33
surfsense_backend/alembic/versions/2_add_linear_connector_enum.py
View file

@ -4,21 +4,21 @@ Revision ID: 2
Revises: e55302644c51 Revises: e55302644c51
""" """
from typing import Sequence, Union from typing import Sequence, Union
from alembic import op from alembic import op
import sqlalchemy as sa
# revision identifiers, used by Alembic. # revision identifiers, used by Alembic.
revision: str = '2' revision: str = "2"
down_revision: Union[str, None] = 'e55302644c51' down_revision: Union[str, None] = "e55302644c51"
branch_labels: Union[str, Sequence[str], None] = None branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None depends_on: Union[str, Sequence[str], None] = None
def upgrade() -> None: def upgrade() -> None:
op.execute(""" op.execute(
"""
DO $$ DO $$
BEGIN BEGIN
IF NOT EXISTS ( IF NOT EXISTS (
@ -31,7 +31,10 @@ def upgrade() -> None:
ALTER TYPE searchsourceconnectortype ADD VALUE 'LINEAR_CONNECTOR'; ALTER TYPE searchsourceconnectortype ADD VALUE 'LINEAR_CONNECTOR';
END IF; END IF;
END$$; END$$;
""") """
)
# #
@ -39,12 +42,18 @@ def downgrade() -> None:
# ### commands auto generated by Alembic - please adjust! ### # ### commands auto generated by Alembic - please adjust! ###
# Downgrading removal of an enum value requires recreating the type # Downgrading removal of an enum value requires recreating the type
op.execute("ALTER TYPE searchsourceconnectortype RENAME TO searchsourceconnectortype_old") op.execute(
op.execute("CREATE TYPE searchsourceconnectortype AS ENUM('SERPER_API', 'TAVILY_API', 'SLACK_CONNECTOR', 'NOTION_CONNECTOR', 'GITHUB_CONNECTOR')") "ALTER TYPE searchsourceconnectortype RENAME TO searchsourceconnectortype_old"
op.execute(( )
"ALTER TABLE search_source_connectors ALTER COLUMN connector_type TYPE searchsourceconnectortype USING " op.execute(
"connector_type::text::searchsourceconnectortype" "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") op.execute("DROP TYPE searchsourceconnectortype_old")
pass pass

41
surfsense_backend/alembic/versions/3_add_linear_connector_to_documenttype_.py
View file

@ -4,25 +4,26 @@ Revision ID: 3
Revises: 2 Revises: 2
""" """
from typing import Sequence, Union from typing import Sequence, Union
from alembic import op from alembic import op
import sqlalchemy as sa
# revision identifiers, used by Alembic. # revision identifiers, used by Alembic.
revision: str = '3' revision: str = "3"
down_revision: Union[str, None] = '2' down_revision: Union[str, None] = "2"
branch_labels: Union[str, Sequence[str], None] = None branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None depends_on: Union[str, Sequence[str], None] = None
# Define the ENUM type name and the new value # 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) ENUM_NAME = "documenttype" # Make sure this matches the name in your DB (usually lowercase class name)
NEW_VALUE = 'LINEAR_CONNECTOR' NEW_VALUE = "LINEAR_CONNECTOR"
def upgrade() -> None: def upgrade() -> None:
"""Upgrade schema.""" """Upgrade schema."""
op.execute(f""" op.execute(
f"""
DO $$ DO $$
BEGIN BEGIN
IF NOT EXISTS ( IF NOT EXISTS (
@ -35,8 +36,8 @@ def upgrade() -> None:
ALTER TYPE {ENUM_NAME} ADD VALUE '{NEW_VALUE}'; ALTER TYPE {ENUM_NAME} ADD VALUE '{NEW_VALUE}';
END IF; END IF;
END$$; END$$;
""") """
)
# Warning: This will delete all rows with the new value # Warning: This will delete all rows with the new value
@ -48,19 +49,19 @@ def downgrade() -> None:
# Enum values *before* LINEAR_CONNECTOR was added # Enum values *before* LINEAR_CONNECTOR was added
old_values = ( old_values = (
'EXTENSION', "EXTENSION",
'CRAWLED_URL', "CRAWLED_URL",
'FILE', "FILE",
'SLACK_CONNECTOR', "SLACK_CONNECTOR",
'NOTION_CONNECTOR', "NOTION_CONNECTOR",
'YOUTUBE_VIDEO', "YOUTUBE_VIDEO",
'GITHUB_CONNECTOR' "GITHUB_CONNECTOR",
) )
old_values_sql = ", ".join([f"'{v}'" for v in old_values]) old_values_sql = ", ".join([f"'{v}'" for v in old_values])
# Table and column names (adjust if different) # Table and column names (adjust if different)
table_name = 'documents' table_name = "documents"
column_name = 'document_type' column_name = "document_type"
# 1. Rename the current enum type # 1. Rename the current enum type
op.execute(f"ALTER TYPE {ENUM_NAME} RENAME TO {old_enum_name}") op.execute(f"ALTER TYPE {ENUM_NAME} RENAME TO {old_enum_name}")
@ -69,9 +70,7 @@ def downgrade() -> None:
op.execute(f"CREATE TYPE {ENUM_NAME} AS ENUM({old_values_sql})") op.execute(f"CREATE TYPE {ENUM_NAME} AS ENUM({old_values_sql})")
# 3. Update the table: # 3. Update the table:
op.execute( op.execute(f"DELETE FROM {table_name} WHERE {column_name}::text = '{NEW_VALUE}'")
f"DELETE FROM {table_name} WHERE {column_name}::text = '{NEW_VALUE}'"
)
# 4. Alter the column to use the new enum type (casting old values) # 4. Alter the column to use the new enum type (casting old values)
op.execute( op.execute(

36
surfsense_backend/alembic/versions/4_add_linkup_api_enum.py
View file

@ -4,24 +4,24 @@ Revision ID: 4
Revises: 3 Revises: 3
""" """
from typing import Sequence, Union from typing import Sequence, Union
from alembic import op from alembic import op
import sqlalchemy as sa
# revision identifiers, used by Alembic. # revision identifiers, used by Alembic.
revision: str = '4' revision: str = "4"
down_revision: Union[str, None] = '3' down_revision: Union[str, None] = "3"
branch_labels: Union[str, Sequence[str], None] = None branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None depends_on: Union[str, Sequence[str], None] = None
def upgrade() -> None: def upgrade() -> None:
ENUM_NAME = 'searchsourceconnectortype' ENUM_NAME = "searchsourceconnectortype"
NEW_VALUE = 'LINKUP_API' NEW_VALUE = "LINKUP_API"
op.execute(f""" op.execute(
f"""
DO $$ DO $$
BEGIN BEGIN
IF NOT EXISTS ( IF NOT EXISTS (
@ -34,20 +34,26 @@ def upgrade() -> None:
ALTER TYPE {ENUM_NAME} ADD VALUE '{NEW_VALUE}'; ALTER TYPE {ENUM_NAME} ADD VALUE '{NEW_VALUE}';
END IF; END IF;
END$$; END$$;
""") """
)
def downgrade() -> None: def downgrade() -> None:
# ### commands auto generated by Alembic - please adjust! ### # ### commands auto generated by Alembic - please adjust! ###
# Downgrading removal of an enum value requires recreating the type # Downgrading removal of an enum value requires recreating the type
op.execute("ALTER TYPE searchsourceconnectortype RENAME TO searchsourceconnectortype_old") op.execute(
op.execute("CREATE TYPE searchsourceconnectortype AS ENUM('SERPER_API', 'TAVILY_API', 'SLACK_CONNECTOR', 'NOTION_CONNECTOR', 'GITHUB_CONNECTOR', 'LINEAR_CONNECTOR')") "ALTER TYPE searchsourceconnectortype RENAME TO searchsourceconnectortype_old"
op.execute(( )
"ALTER TABLE search_source_connectors ALTER COLUMN connector_type TYPE searchsourceconnectortype USING " op.execute(
"connector_type::text::searchsourceconnectortype" "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") op.execute("DROP TYPE searchsourceconnectortype_old")
pass pass

36
surfsense_backend/alembic/versions/6_change_podcast_content_to_transcript.py
View file

@ -4,48 +4,58 @@ Revision ID: 6
Revises: 5 Revises: 5
""" """
from typing import Sequence, Union from typing import Sequence, Union
from alembic import op
import sqlalchemy as sa import sqlalchemy as sa
from sqlalchemy.dialects.postgresql import JSON from alembic import op
from sqlalchemy import inspect from sqlalchemy import inspect
from sqlalchemy.dialects.postgresql import JSON
# revision identifiers, used by Alembic. # revision identifiers, used by Alembic.
revision: str = '6' revision: str = "6"
down_revision: Union[str, None] = '5' down_revision: Union[str, None] = "5"
branch_labels: Union[str, Sequence[str], None] = None branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None depends_on: Union[str, Sequence[str], None] = None
def upgrade() -> None: def upgrade() -> None:
bind = op.get_bind() bind = op.get_bind()
inspector = inspect(bind) inspector = inspect(bind)
columns = [col["name"] for col in inspector.get_columns("podcasts")] columns = [col["name"] for col in inspector.get_columns("podcasts")]
if "podcast_transcript" not in columns: 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 # Copy data from old column to new column
op.execute(""" op.execute(
"""
UPDATE podcasts UPDATE podcasts
SET podcast_transcript = jsonb_build_object('text', podcast_content) SET podcast_transcript = jsonb_build_object('text', podcast_content)
WHERE podcast_content != '' WHERE podcast_content != ''
""") """
)
# Drop the old column only if it exists # Drop the old column only if it exists
if "podcast_content" in columns: if "podcast_content" in columns:
op.drop_column('podcasts', 'podcast_content') op.drop_column("podcasts", "podcast_content")
def downgrade() -> None: def downgrade() -> None:
# Add back the original column # 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 # Copy data from JSON column back to text column
# Extract the 'text' field if it exists, otherwise use empty string # 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 # Drop the new column
op.drop_column('podcasts', 'podcast_transcript') op.drop_column("podcasts", "podcast_transcript")

17
surfsense_backend/alembic/versions/7_remove_is_generated_column.py
View file

@ -4,19 +4,20 @@ Revision ID: 7
Revises: 6 Revises: 6
""" """
from typing import Sequence, Union from typing import Sequence, Union
from alembic import op
import sqlalchemy as sa import sqlalchemy as sa
from alembic import op
from sqlalchemy import inspect from sqlalchemy import inspect
# revision identifiers, used by Alembic. # revision identifiers, used by Alembic.
revision: str = '7' revision: str = "7"
down_revision: Union[str, None] = '6' down_revision: Union[str, None] = "6"
branch_labels: Union[str, Sequence[str], None] = None branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None depends_on: Union[str, Sequence[str], None] = None
def upgrade() -> None: def upgrade() -> None:
# Get the current database connection # Get the current database connection
bind = op.get_bind() bind = op.get_bind()
@ -25,10 +26,12 @@ def upgrade() -> None:
# Check if the column exists before attempting to drop it # Check if the column exists before attempting to drop it
columns = [col["name"] for col in inspector.get_columns("podcasts")] columns = [col["name"] for col in inspector.get_columns("podcasts")]
if "is_generated" in columns: if "is_generated" in columns:
op.drop_column('podcasts', 'is_generated') op.drop_column("podcasts", "is_generated")
def downgrade() -> None: def downgrade() -> None:
# Add back the is_generated column with its original constraints # 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')) op.add_column(
"podcasts",
sa.Column("is_generated", sa.Boolean(), nullable=False, server_default="false"),
)

52
surfsense_backend/alembic/versions/8_add_content_hash_to_documents.py
View file

@ -3,15 +3,16 @@
Revision ID: 8 Revision ID: 8
Revises: 7 Revises: 7
""" """
from typing import Sequence, Union from typing import Sequence, Union
from alembic import op
import sqlalchemy as sa import sqlalchemy as sa
from alembic import op
from sqlalchemy import inspect from sqlalchemy import inspect
# revision identifiers, used by Alembic. # revision identifiers, used by Alembic.
revision: str = '8' revision: str = "8"
down_revision: Union[str, None] = '7' down_revision: Union[str, None] = "7"
branch_labels: Union[str, Sequence[str], None] = None branch_labels: Union[str, Sequence[str], None] = None
depends_on: 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: def upgrade() -> None:
bind = op.get_bind() bind = op.get_bind()
inspector = inspect(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 # Only add the column if it doesn't already exist
if 'content_hash' not in columns: if "content_hash" not in columns:
op.add_column('documents', sa.Column('content_hash', sa.String(), nullable=True)) op.add_column(
"documents", sa.Column("content_hash", sa.String(), nullable=True)
)
# Populate the content_hash column # Populate the content_hash column
op.execute(""" op.execute(
"""
UPDATE documents UPDATE documents
SET content_hash = encode(sha256(convert_to(content, 'UTF8')), 'hex') SET content_hash = encode(sha256(convert_to(content, 'UTF8')), 'hex')
WHERE content_hash IS NULL WHERE content_hash IS NULL
""") """
)
op.execute(""" op.execute(
"""
DELETE FROM documents DELETE FROM documents
WHERE id NOT IN ( WHERE id NOT IN (
SELECT MIN(id) SELECT MIN(id)
FROM documents FROM documents
GROUP BY content_hash GROUP BY content_hash
) )
""") """
)
op.alter_column('documents', 'content_hash', op.alter_column(
existing_type=sa.String(), "documents", "content_hash", existing_type=sa.String(), nullable=False
nullable=False) )
op.create_index(op.f('ix_documents_content_hash'), 'documents', ['content_hash'], unique=False) op.create_index(
op.create_unique_constraint(op.f('uq_documents_content_hash'), 'documents', ['content_hash']) 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: else:
print("Column 'content_hash' already exists. Skipping column creation.") print("Column 'content_hash' already exists. Skipping column creation.")
def downgrade() -> None: def downgrade() -> None:
op.drop_constraint(op.f('uq_documents_content_hash'), 'documents', type_='unique') 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_index(op.f("ix_documents_content_hash"), table_name="documents")
op.drop_column('documents', 'content_hash') op.drop_column("documents", "content_hash")

16
surfsense_backend/alembic/versions/9_add_discord_connector_enum_and_documenttype.py
View file

@ -7,8 +7,6 @@ Revises: 8
from typing import Sequence, Union from typing import Sequence, Union
from alembic import op from alembic import op
import sqlalchemy as sa
# revision identifiers, used by Alembic. # revision identifiers, used by Alembic.
revision: str = "9" revision: str = "9"
@ -26,7 +24,8 @@ DOCUMENT_NEW_VALUE = "DISCORD_CONNECTOR"
def upgrade() -> None: def upgrade() -> None:
"""Upgrade schema - add DISCORD_CONNECTOR to connector and document enum safely.""" """Upgrade schema - add DISCORD_CONNECTOR to connector and document enum safely."""
# Add DISCORD_CONNECTOR to searchsourceconnectortype only if not exists # Add DISCORD_CONNECTOR to searchsourceconnectortype only if not exists
op.execute(f""" op.execute(
f"""
DO $$ DO $$
BEGIN BEGIN
IF NOT EXISTS ( IF NOT EXISTS (
@ -37,10 +36,12 @@ def upgrade() -> None:
ALTER TYPE {CONNECTOR_ENUM} ADD VALUE '{CONNECTOR_NEW_VALUE}'; ALTER TYPE {CONNECTOR_ENUM} ADD VALUE '{CONNECTOR_NEW_VALUE}';
END IF; END IF;
END$$; END$$;
""") """
)
# Add DISCORD_CONNECTOR to documenttype only if not exists # Add DISCORD_CONNECTOR to documenttype only if not exists
op.execute(f""" op.execute(
f"""
DO $$ DO $$
BEGIN BEGIN
IF NOT EXISTS ( IF NOT EXISTS (
@ -51,8 +52,8 @@ def upgrade() -> None:
ALTER TYPE {DOCUMENT_ENUM} ADD VALUE '{DOCUMENT_NEW_VALUE}'; ALTER TYPE {DOCUMENT_ENUM} ADD VALUE '{DOCUMENT_NEW_VALUE}';
END IF; END IF;
END$$; END$$;
""") """
)
def downgrade() -> None: def downgrade() -> None:
@ -109,7 +110,6 @@ def downgrade() -> None:
# 4. Drop the old connector enum type # 4. Drop the old connector enum type
op.execute(f"DROP TYPE {old_connector_enum_name}") op.execute(f"DROP TYPE {old_connector_enum_name}")
# Document Enum Downgrade Steps # Document Enum Downgrade Steps
# 1. Rename the current document enum type # 1. Rename the current document enum type
op.execute(f"ALTER TYPE {DOCUMENT_ENUM} RENAME TO {old_document_enum_name}") op.execute(f"ALTER TYPE {DOCUMENT_ENUM} RENAME TO {old_document_enum_name}")

37
surfsense_backend/alembic/versions/e55302644c51_add_github_connector_to_documenttype_.py
View file

@ -1,21 +1,22 @@
from typing import Sequence, Union from typing import Sequence, Union
from alembic import op from alembic import op
import sqlalchemy as sa
# revision identifiers, used by Alembic. # revision identifiers, used by Alembic.
revision: str = 'e55302644c51' revision: str = "e55302644c51"
down_revision: Union[str, None] = '1' down_revision: Union[str, None] = "1"
branch_labels: Union[str, Sequence[str], None] = None branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None depends_on: Union[str, Sequence[str], None] = None
# Define the ENUM type name and the new value # Define the ENUM type name and the new value
ENUM_NAME = 'documenttype' ENUM_NAME = "documenttype"
NEW_VALUE = 'GITHUB_CONNECTOR' NEW_VALUE = "GITHUB_CONNECTOR"
def upgrade() -> None: def upgrade() -> None:
"""Upgrade schema.""" """Upgrade schema."""
op.execute(f""" op.execute(
f"""
DO $$ DO $$
BEGIN BEGIN
IF NOT EXISTS ( IF NOT EXISTS (
@ -28,32 +29,32 @@ def upgrade() -> None:
ALTER TYPE {ENUM_NAME} ADD VALUE '{NEW_VALUE}'; ALTER TYPE {ENUM_NAME} ADD VALUE '{NEW_VALUE}';
END IF; END IF;
END$$; END$$;
""") """
)
def downgrade() -> None: def downgrade() -> None:
"""Downgrade schema - remove GITHUB_CONNECTOR from enum.""" """Downgrade schema - remove GITHUB_CONNECTOR from enum."""
old_enum_name = f"{ENUM_NAME}_old" old_enum_name = f"{ENUM_NAME}_old"
old_values = ( old_values = (
'EXTENSION', "EXTENSION",
'CRAWLED_URL', "CRAWLED_URL",
'FILE', "FILE",
'SLACK_CONNECTOR', "SLACK_CONNECTOR",
'NOTION_CONNECTOR', "NOTION_CONNECTOR",
'YOUTUBE_VIDEO' "YOUTUBE_VIDEO",
) )
old_values_sql = ", ".join([f"'{v}'" for v in old_values]) old_values_sql = ", ".join([f"'{v}'" for v in old_values])
table_name = 'documents' table_name = "documents"
column_name = 'document_type' column_name = "document_type"
# 1. Create the new enum type with the old values # 1. Create the new enum type with the old values
op.execute(f"CREATE TYPE {old_enum_name} AS ENUM({old_values_sql})") op.execute(f"CREATE TYPE {old_enum_name} AS ENUM({old_values_sql})")
# 2. Delete rows using the new value # 2. Delete rows using the new value
op.execute( op.execute(f"DELETE FROM {table_name} WHERE {column_name}::text = '{NEW_VALUE}'")
f"DELETE FROM {table_name} WHERE {column_name}::text = '{NEW_VALUE}'"
)
# 3. Alter the column to use the old enum type # 3. Alter the column to use the old enum type
op.execute( op.execute(