Handle exact tenant column collisions

daryllimyt · daryllimyt · commit f63407a04999 · 2026-03-09T23:15:57.000-04:00
diff --git a/alembic/versions/b5fc4168fe22_apply_model_a_rls_to_dynamic_workspace_.py b/alembic/versions/b5fc4168fe22_apply_model_a_rls_to_dynamic_workspace_.py
@@ -12,6 +12,7 @@
 from uuid import UUID
 
 import sqlalchemy as sa
+from sqlalchemy.dialects.postgresql import JSONB
 
 from alembic import op
 from tracecat.identifiers.workflow import WorkspaceUUID
@@ -23,6 +24,7 @@
 depends_on: str | Sequence[str] | None = None
 
 INTERNAL_TENANT_COLUMN = "__tc_workspace_id"
+LEGACY_TENANT_COLUMN_PREFIX = "migrated_tc_workspace_id"
 DYNAMIC_WORKSPACE_RLS_POLICY = "rls_policy_dynamic_workspace"
 RLS_WORKSPACE_VAR = "app.current_workspace_id"
 RLS_BYPASS_VAR = "app.rls_bypass"
@@ -73,18 +75,147 @@ def _policy_expression() -> str:
     )
 
 
-def _ensure_tenant_column(
+def _next_legacy_tenant_column_name(existing_column_names: set[str]) -> str:
+    """Generate a non-conflicting rename target for legacy tenant collisions."""
+    candidate = LEGACY_TENANT_COLUMN_PREFIX
+    suffix = 1
+    while candidate in existing_column_names:
+        candidate = f"{LEGACY_TENANT_COLUMN_PREFIX}_{suffix}"
+        suffix += 1
+    return candidate
+
+
+def _rename_table_column_metadata(
+    *,
+    workspace_id: UUID,
+    table_name: str,
+    old_name: str,
+    new_name: str,
+) -> None:
+    """Rename matching table-column metadata when a physical column moves."""
+    bind = op.get_bind()
+    inspector = sa.inspect(bind)
+    if not (
+        inspector.has_table("tables", schema="public")
+        and inspector.has_table("table_column", schema="public")
+    ):
+        return
+
+    bind.execute(
+        sa.text(
+            """
+            UPDATE table_column AS tc
+            SET name = :new_name
+            FROM tables AS t
+            WHERE tc.table_id = t.id
+              AND t.workspace_id = :workspace_id
+              AND t.name = :table_name
+              AND tc.name = :old_name
+            """
+        ),
+        {
+            "workspace_id": str(workspace_id),
+            "table_name": table_name,
+            "old_name": old_name,
+            "new_name": new_name,
+        },
+    )
+
+
+def _rename_case_field_schema_key(
+    *,
+    workspace_id: UUID,
+    old_name: str,
+    new_name: str,
+) -> None:
+    """Rename a case-field schema key when the physical column moves."""
+    bind = op.get_bind()
+    inspector = sa.inspect(bind)
+    if not inspector.has_table("case_field", schema="public"):
+        return
+
+    case_field_table = sa.table(
+        "case_field",
+        sa.column("workspace_id", sa.UUID()),
+        sa.column("schema", JSONB),
+    )
+    current_schema = bind.execute(
+        sa.select(case_field_table.c.schema).where(
+            case_field_table.c.workspace_id == workspace_id
+        )
+    ).scalar_one_or_none()
+    if not isinstance(current_schema, dict) or old_name not in current_schema:
+        return
+
+    updated_schema = dict(current_schema)
+    updated_schema[new_name] = updated_schema.pop(old_name)
+    bind.execute(
+        sa.update(case_field_table)
+        .where(case_field_table.c.workspace_id == workspace_id)
+        .values(schema=updated_schema)
+    )
+
+
+def _rename_legacy_tenant_column(
     *,
     schema_name: str,
     table_name: str,
     workspace_id: UUID,
     preparer: sa.sql.compiler.IdentifierPreparer,
 ) -> None:
-    """Add/backfill the internal tenant column on a physical dynamic table."""
+    """Rename any pre-existing user column that collides with the tenant column."""
     bind = op.get_bind()
     inspector = sa.inspect(bind)
+    columns = inspector.get_columns(table_name, schema=schema_name)
+    column_names = {column["name"] for column in columns}
+    if INTERNAL_TENANT_COLUMN not in column_names:
+        return
+
+    qualified_table = _qualified_table(preparer, schema_name, table_name)
+    legacy_column_name = _next_legacy_tenant_column_name(column_names)
+    op.execute(
+        sa.text(
+            f"""
+            ALTER TABLE {qualified_table}
+            RENAME COLUMN "{INTERNAL_TENANT_COLUMN}" TO "{legacy_column_name}"
+            """
+        )
+    )
+
+    if schema_name.startswith("tables_"):
+        _rename_table_column_metadata(
+            workspace_id=workspace_id,
+            table_name=table_name,
+            old_name=INTERNAL_TENANT_COLUMN,
+            new_name=legacy_column_name,
+        )
+    elif schema_name.startswith("custom_fields_") and table_name == "case_fields":
+        _rename_case_field_schema_key(
+            workspace_id=workspace_id,
+            old_name=INTERNAL_TENANT_COLUMN,
+            new_name=legacy_column_name,
+        )
+
+
+def _ensure_tenant_column(
+    *,
+    schema_name: str,
+    table_name: str,
+    workspace_id: UUID,
+    preparer: sa.sql.compiler.IdentifierPreparer,
+) -> None:
+    """Add/backfill the internal tenant column on a physical dynamic table."""
     qualified_table = _qualified_table(preparer, schema_name, table_name)
 
+    _rename_legacy_tenant_column(
+        schema_name=schema_name,
+        table_name=table_name,
+        workspace_id=workspace_id,
+        preparer=preparer,
+    )
+
+    bind = op.get_bind()
+    inspector = sa.inspect(bind)
     has_tenant_column = any(
         column["name"] == INTERNAL_TENANT_COLUMN
         for column in inspector.get_columns(table_name, schema=schema_name)
diff --git a/tests/migration/test_dynamic_workspace_rls_migration.py b/tests/migration/test_dynamic_workspace_rls_migration.py
@@ -20,12 +20,13 @@
 from tracecat.identifiers.workflow import WorkspaceUUID
 
 MIGRATION_REVISION = "b5fc4168fe22"
-PREVIOUS_REVISION = "c76f9b01fad7"
+PREVIOUS_REVISION = "9a6d0e0ec5b1"
 TABLES_PREFIX = "tables_"
 CUSTOM_FIELDS_PREFIX = "custom_fields_"
 TABLES_TABLE = "alerts"
 CUSTOM_FIELDS_TABLE = "case_fields"
 INTERNAL_TENANT_COLUMN = "__tc_workspace_id"
+LEGACY_TENANT_COLUMN = "migrated_tc_workspace_id"
 DYNAMIC_WORKSPACE_RLS_POLICY = "rls_policy_dynamic_workspace"
 
 
@@ -263,6 +264,80 @@ def test_db():
 
 
 class TestDynamicWorkspaceRlsMigration:
+    def test_upgrade_renames_legacy_exact_tenant_column_collision(
+        self, test_db
+    ) -> None:
+        engine = create_engine(test_db["db_url"])
+        try:
+            with engine.begin() as conn:
+                workspace_id = test_db["workspace_ids"][0]
+                for prefix, table_name in (
+                    (TABLES_PREFIX, TABLES_TABLE),
+                    (CUSTOM_FIELDS_PREFIX, CUSTOM_FIELDS_TABLE),
+                ):
+                    schema_name = _workspace_schema(prefix, workspace_id)
+                    conn.execute(
+                        text(
+                            f'''
+                            ALTER TABLE "{schema_name}"."{table_name}"
+                            ADD COLUMN "{INTERNAL_TENANT_COLUMN}" TEXT
+                            '''
+                        )
+                    )
+                    conn.execute(
+                        text(
+                            f'''
+                            UPDATE "{schema_name}"."{table_name}"
+                            SET "{INTERNAL_TENANT_COLUMN}" = :legacy_value
+                            '''
+                        ),
+                        {"legacy_value": f"legacy-{table_name}"},
+                    )
+        finally:
+            engine.dispose()
+
+        _run_alembic_upgrade(test_db["db_url"])
+
+        engine = create_engine(test_db["db_url"])
+        try:
+            with engine.begin() as conn:
+                workspace_id = test_db["workspace_ids"][0]
+                for prefix, table_name in (
+                    (TABLES_PREFIX, TABLES_TABLE),
+                    (CUSTOM_FIELDS_PREFIX, CUSTOM_FIELDS_TABLE),
+                ):
+                    schema_name = _workspace_schema(prefix, workspace_id)
+                    columns = conn.execute(
+                        text(
+                            """
+                            SELECT column_name
+                            FROM information_schema.columns
+                            WHERE table_schema = :schema_name
+                              AND table_name = :table_name
+                            """
+                        ),
+                        {
+                            "schema_name": schema_name,
+                            "table_name": table_name,
+                        },
+                    ).fetchall()
+                    column_names = {row[0] for row in columns}
+                    assert INTERNAL_TENANT_COLUMN in column_names
+                    assert LEGACY_TENANT_COLUMN in column_names
+
+                    tenant_values = conn.execute(
+                        text(
+                            f'''
+                            SELECT "{INTERNAL_TENANT_COLUMN}", "{LEGACY_TENANT_COLUMN}"
+                            FROM "{schema_name}"."{table_name}"
+                            '''
+                        )
+                    ).one()
+                    assert tenant_values[0] == workspace_id
+                    assert tenant_values[1] == f"legacy-{table_name}"
+        finally:
+            engine.dispose()
+
     def test_upgrade_backfills_tenant_column_and_sets_default(self, test_db) -> None:
         _run_alembic_upgrade(test_db["db_url"])
 
diff --git a/tests/unit/test_case_fields_service.py b/tests/unit/test_case_fields_service.py
@@ -180,6 +180,17 @@ async def test_create_field_rejects_internal_name(
                 CaseFieldCreate(name="__TC_workspace_id", type=SqlType.TEXT)
             )
 
+    async def test_create_field_allows_legacy_tc_prefix_name(
+        self, case_fields_service: CaseFieldsService
+    ) -> None:
+        """Only the exact internal tenant column name should be reserved."""
+        with patch.object(
+            case_fields_service.editor, "create_column"
+        ) as mock_create_column:
+            field_params = CaseFieldCreate(name="__tc_shadow", type=SqlType.TEXT)
+            await case_fields_service.create_field(field_params)
+            mock_create_column.assert_called_once_with(field_params)
+
     async def test_update_field(self, case_fields_service: CaseFieldsService) -> None:
         """Test updating a case field."""
         # Create field update parameters
diff --git a/tests/unit/test_tables_service.py b/tests/unit/test_tables_service.py
@@ -30,6 +30,7 @@
     TableEditorService,
     TablesService,
     is_internal_column_name,
+    visible_column_names,
 )
 
 pytestmark = pytest.mark.usefixtures("db")
@@ -151,9 +152,10 @@ async def _list_rows(
 @pytest.mark.anyio
 class TestTablesService:
     async def test_internal_column_name_check_is_case_insensitive(self) -> None:
-        """Internal column detection should normalize case before prefix checks."""
+        """Internal column detection should normalize case for exact matches."""
         assert is_internal_column_name("__tc_workspace_id") is True
         assert is_internal_column_name("__TC_workspace_id") is True
+        assert is_internal_column_name("__tc_shadow") is False
         assert is_internal_column_name("user_field") is False
 
     async def test_create_and_get_table(self, tables_service: TablesService) -> None:
@@ -942,6 +944,15 @@ async def test_table_editor_visible_columns_refresh_after_create_column(
         retrieved = await editor.get_row(inserted["id"])
         assert retrieved["city"] == "NYC"
 
+    async def test_visible_column_names_preserves_legacy_underscore_names(self) -> None:
+        """Visibility filtering should not reject legacy underscore-prefixed names."""
+        assert visible_column_names(["_legacy_field", "__tc_workspace_id"]) == [
+            "id",
+            "created_at",
+            "updated_at",
+            "_legacy_field",
+        ]
+
     async def test_table_editor_insert_row_rejects_internal_column(
         self, tables_service: TablesService, table: Table
     ) -> None:
diff --git a/tracecat/tables/service.py b/tracecat/tables/service.py
@@ -72,27 +72,28 @@
     InFailedSQLTransactionError,
 )
 
-INTERNAL_COLUMN_PREFIX = "__tc_"
 DYNAMIC_WORKSPACE_TENANT_COLUMN = "__tc_workspace_id"
 DYNAMIC_WORKSPACE_RLS_POLICY = "rls_policy_dynamic_workspace"
 RLS_WORKSPACE_VAR = "app.current_workspace_id"
 RLS_BYPASS_VAR = "app.rls_bypass"
 RLS_BYPASS_ON = "on"
+INTERNAL_COLUMN_NAMES: frozenset[str] = frozenset({DYNAMIC_WORKSPACE_TENANT_COLUMN})
 SYSTEM_VISIBLE_COLUMN_NAMES: tuple[str, ...] = ("id", "created_at", "updated_at")
 
 
 def visible_column_names(column_names: Sequence[str]) -> list[str]:
     """Build the API-facing column order for dynamic table row payloads."""
     visible_names: list[str] = list(SYSTEM_VISIBLE_COLUMN_NAMES)
-    seen: set[str] = set(SYSTEM_VISIBLE_COLUMN_NAMES)
+    seen_lower: set[str] = {
+        column_name.lower() for column_name in SYSTEM_VISIBLE_COLUMN_NAMES
+    }
     for column_name in column_names:
         if is_internal_column_name(column_name):
             continue
-        sanitized_name = sanitize_identifier(column_name)
-        if sanitized_name in seen:
+        if (normalized_name := column_name.lower()) in seen_lower:
             continue
-        visible_names.append(sanitized_name)
-        seen.add(sanitized_name)
+        visible_names.append(column_name)
+        seen_lower.add(normalized_name)
     return visible_names
 
 
@@ -2239,4 +2240,4 @@ def sanitize_identifier(identifier: str) -> str:
 
 def is_internal_column_name(column_name: str) -> bool:
     """Check whether a column is internal/system-managed for dynamic schemas."""
-    return column_name.lower().startswith(INTERNAL_COLUMN_PREFIX)
+    return column_name.lower() in INTERNAL_COLUMN_NAMES
