From 40ade4889e23b18ed43eaaa39a0489483fa5bbdc Mon Sep 17 00:00:00 2001
From: Anish Sarkar <104695310+AnishSarkar22@users.noreply.github.com>
Date: Thu, 2 Apr 2026 21:01:31 +0530
Subject: [PATCH] feat: add LOCAL_FOLDER_FILE document type and update
 document_versions table management

---
 ...d_local_folder_connector_and_versioning.py | 33 +++++++------------
 1 file changed, 12 insertions(+), 21 deletions(-)

diff --git a/surfsense_backend/alembic/versions/117_add_local_folder_connector_and_versioning.py b/surfsense_backend/alembic/versions/117_add_local_folder_connector_and_versioning.py
index e97a4787c..a9da3beb4 100644
--- a/surfsense_backend/alembic/versions/117_add_local_folder_connector_and_versioning.py
+++ b/surfsense_backend/alembic/versions/117_add_local_folder_connector_and_versioning.py
@@ -1,4 +1,4 @@
-"""Add local folder connector enums and document_versions table
+"""Add LOCAL_FOLDER_FILE document type and document_versions table
 
 Revision ID: 117
 Revises: 116
@@ -21,23 +21,6 @@ PUBLICATION_NAME = "zero_publication"
 def upgrade() -> None:
     conn = op.get_bind()
 
-    # Add LOCAL_FOLDER_CONNECTOR to searchsourceconnectortype enum
-    op.execute(
-        """
-    DO $$
-    BEGIN
-        IF NOT EXISTS (
-            SELECT 1 FROM pg_type t
-            JOIN pg_enum e ON t.oid = e.enumtypid
-            WHERE t.typname = 'searchsourceconnectortype' AND e.enumlabel = 'LOCAL_FOLDER_CONNECTOR'
-        ) THEN
-            ALTER TYPE searchsourceconnectortype ADD VALUE 'LOCAL_FOLDER_CONNECTOR';
-        END IF;
-    END
-    $$;
-    """
-    )
-
     # Add LOCAL_FOLDER_FILE to documenttype enum
     op.execute(
         """
@@ -126,9 +109,17 @@ def downgrade() -> None:
         {"name": PUBLICATION_NAME},
     ).fetchone()
     if pub_exists:
-        op.execute(
-            f"ALTER PUBLICATION {PUBLICATION_NAME} DROP TABLE IF EXISTS document_versions"
-        )
+        already_in_pub = conn.execute(
+            sa.text(
+                "SELECT 1 FROM pg_publication_tables "
+                "WHERE pubname = :name AND tablename = 'document_versions'"
+            ),
+            {"name": PUBLICATION_NAME},
+        ).fetchone()
+        if already_in_pub:
+            op.execute(
+                f"ALTER PUBLICATION {PUBLICATION_NAME} DROP TABLE document_versions"
+            )
 
     op.execute("DROP INDEX IF EXISTS ix_document_versions_created_at")
     op.execute("DROP INDEX IF EXISTS ix_document_versions_document_id")