feat(tables): enforce model-a RLS on dynamic workspace schemas

daryllimyt · daryllimyt · commit 37f760c9596b · 2026-03-09T12:39:25.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
@@ -0,0 +1,215 @@
+"""apply model a rls to dynamic workspace schemas
+
+Revision ID: b5fc4168fe22
+Revises: c76f9b01fad7
+Create Date: 2026-02-27 11:43:07.059399
+
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from uuid import UUID
+
+import sqlalchemy as sa
+
+from alembic import op
+from tracecat.identifiers.workflow import WorkspaceUUID
+
+# revision identifiers, used by Alembic.
+revision: str = "b5fc4168fe22"
+down_revision: str | None = "c76f9b01fad7"
+branch_labels: str | Sequence[str] | None = None
+depends_on: str | Sequence[str] | None = None
+
+INTERNAL_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"
+DYNAMIC_WORKSPACE_SCHEMA_PREFIXES = ("tables_", "custom_fields_")
+
+
+def _workspace_id_from_schema(schema_name: str) -> UUID | None:
+    """Resolve workspace UUID from a dynamic workspace schema name."""
+    for prefix in DYNAMIC_WORKSPACE_SCHEMA_PREFIXES:
+        if not schema_name.startswith(prefix):
+            continue
+        short_workspace_id = schema_name.removeprefix(prefix)
+        try:
+            return WorkspaceUUID.new(short_workspace_id)
+        except ValueError:
+            return None
+    return None
+
+
+def _qualified_table(
+    preparer: sa.sql.compiler.IdentifierPreparer, schema_name: str, table_name: str
+) -> str:
+    """Return a properly quoted schema-qualified table reference."""
+    return f"{preparer.quote_schema(schema_name)}.{preparer.quote(table_name)}"
+
+
+def _workspace_dynamic_tables() -> list[tuple[str, str, UUID]]:
+    """Collect all physical tables in workspace-scoped dynamic schemas."""
+    bind = op.get_bind()
+    inspector = sa.inspect(bind)
+    tables: list[tuple[str, str, UUID]] = []
+    for schema_name in sorted(inspector.get_schema_names()):
+        workspace_id = _workspace_id_from_schema(schema_name)
+        if workspace_id is None:
+            continue
+        for table_name in sorted(inspector.get_table_names(schema=schema_name)):
+            tables.append((schema_name, table_name, workspace_id))
+    return tables
+
+
+def _policy_expression() -> str:
+    """Build the RLS tenant policy expression for dynamic workspace tables."""
+    return (
+        f"current_setting('{RLS_BYPASS_VAR}', true) = '{RLS_BYPASS_ON}' "
+        f'OR "{INTERNAL_TENANT_COLUMN}" = '
+        f"NULLIF(current_setting('{RLS_WORKSPACE_VAR}', true), '')::uuid"
+    )
+
+
+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."""
+    bind = op.get_bind()
+    inspector = sa.inspect(bind)
+    qualified_table = _qualified_table(preparer, schema_name, table_name)
+
+    has_tenant_column = any(
+        column["name"] == INTERNAL_TENANT_COLUMN
+        for column in inspector.get_columns(table_name, schema=schema_name)
+    )
+    if not has_tenant_column:
+        op.execute(
+            sa.text(
+                f"""
+                ALTER TABLE {qualified_table}
+                ADD COLUMN "{INTERNAL_TENANT_COLUMN}" UUID
+                """
+            )
+        )
+
+    bind.execute(
+        sa.text(
+            f"""
+            UPDATE {qualified_table}
+            SET "{INTERNAL_TENANT_COLUMN}" = :workspace_id
+            WHERE "{INTERNAL_TENANT_COLUMN}" IS NULL
+            """
+        ),
+        {"workspace_id": str(workspace_id)},
+    )
+
+    op.execute(
+        sa.text(
+            f"""
+            ALTER TABLE {qualified_table}
+            ALTER COLUMN "{INTERNAL_TENANT_COLUMN}"
+            SET DEFAULT '{workspace_id}'::uuid
+            """
+        )
+    )
+    op.execute(
+        sa.text(
+            f"""
+            ALTER TABLE {qualified_table}
+            ALTER COLUMN "{INTERNAL_TENANT_COLUMN}"
+            SET NOT NULL
+            """
+        )
+    )
+
+
+def _enable_table_rls(
+    *, schema_name: str, table_name: str, preparer: sa.sql.compiler.IdentifierPreparer
+) -> None:
+    """Enable RLS policy for a physical dynamic table."""
+    qualified_table = _qualified_table(preparer, schema_name, table_name)
+    policy_expr = _policy_expression()
+
+    op.execute(sa.text(f"ALTER TABLE {qualified_table} ENABLE ROW LEVEL SECURITY"))
+    op.execute(
+        sa.text(
+            f"""
+            DROP POLICY IF EXISTS {DYNAMIC_WORKSPACE_RLS_POLICY}
+            ON {qualified_table}
+            """
+        )
+    )
+    op.execute(
+        sa.text(
+            f"""
+            CREATE POLICY {DYNAMIC_WORKSPACE_RLS_POLICY} ON {qualified_table}
+                FOR ALL
+                USING ({policy_expr})
+                WITH CHECK ({policy_expr})
+            """
+        )
+    )
+
+
+def _disable_table_rls(
+    *, schema_name: str, table_name: str, preparer: sa.sql.compiler.IdentifierPreparer
+) -> None:
+    """Disable RLS policy for a physical dynamic table."""
+    qualified_table = _qualified_table(preparer, schema_name, table_name)
+    op.execute(
+        sa.text(
+            f"""
+            DROP POLICY IF EXISTS {DYNAMIC_WORKSPACE_RLS_POLICY}
+            ON {qualified_table}
+            """
+        )
+    )
+    op.execute(sa.text(f"ALTER TABLE {qualified_table} DISABLE ROW LEVEL SECURITY"))
+
+
+def upgrade() -> None:
+    """Apply workspace-tenant column + RLS to dynamic workspace schemas."""
+    bind = op.get_bind()
+    preparer = sa.sql.compiler.IdentifierPreparer(bind.dialect)
+
+    for schema_name, table_name, workspace_id in _workspace_dynamic_tables():
+        _ensure_tenant_column(
+            schema_name=schema_name,
+            table_name=table_name,
+            workspace_id=workspace_id,
+            preparer=preparer,
+        )
+        _enable_table_rls(
+            schema_name=schema_name,
+            table_name=table_name,
+            preparer=preparer,
+        )
+
+
+def downgrade() -> None:
+    """Rollback workspace-tenant column + RLS on dynamic workspace schemas."""
+    bind = op.get_bind()
+    preparer = sa.sql.compiler.IdentifierPreparer(bind.dialect)
+
+    for schema_name, table_name, _ in _workspace_dynamic_tables():
+        _disable_table_rls(
+            schema_name=schema_name,
+            table_name=table_name,
+            preparer=preparer,
+        )
+        qualified_table = _qualified_table(preparer, schema_name, table_name)
+        op.execute(
+            sa.text(
+                f"""
+                ALTER TABLE {qualified_table}
+                DROP COLUMN IF EXISTS "{INTERNAL_TENANT_COLUMN}"
+                """
+            )
+        )
diff --git a/tracecat/cases/service.py b/tracecat/cases/service.py
@@ -114,7 +114,7 @@
 from tracecat.service import BaseWorkspaceService, requires_entitlement
 from tracecat.tables.common import normalize_column_options
 from tracecat.tables.enums import SqlType
-from tracecat.tables.service import TablesService
+from tracecat.tables.service import DYNAMIC_WORKSPACE_TENANT_COLUMN, TablesService
 from tracecat.tiers.enums import Entitlement
 
 
@@ -1002,7 +1002,14 @@ class CaseFieldsService(CustomFieldsService):
     # Hardcoded to preserve existing workspace-scoped table names
     # (metadata table was renamed from case_fields to case_field)
     _table = "case_fields"
-    _reserved_columns = {"id", "case_id", "created_at", "updated_at", "workspace_id"}
+    _reserved_columns = {
+        "id",
+        "case_id",
+        "created_at",
+        "updated_at",
+        "workspace_id",
+        DYNAMIC_WORKSPACE_TENANT_COLUMN,
+    }
 
     def _table_definition(self) -> sa.Table:
         """Return the SQLAlchemy Table definition for the case_fields workspace table."""
@@ -1023,6 +1030,12 @@ def _table_definition(self) -> sa.Table:
                 nullable=False,
                 server_default=sa.func.now(),
             ),
+            sa.Column(
+                DYNAMIC_WORKSPACE_TENANT_COLUMN,
+                UUID(as_uuid=True),
+                nullable=False,
+                server_default=self._workspace_tenant_default_sql(),
+            ),
             # Use the actual Case table column to avoid metadata resolution issues
             sa.ForeignKeyConstraint(
                 ["case_id"],
diff --git a/tracecat/custom_fields/service.py b/tracecat/custom_fields/service.py
@@ -13,7 +13,16 @@
 from tracecat.db.locks import derive_lock_key_from_parts, pg_advisory_lock
 from tracecat.identifiers.workflow import WorkspaceUUID
 from tracecat.service import BaseWorkspaceService
-from tracecat.tables.service import TableEditorService, sanitize_identifier
+from tracecat.tables.service import (
+    DYNAMIC_WORKSPACE_RLS_POLICY,
+    DYNAMIC_WORKSPACE_TENANT_COLUMN,
+    RLS_BYPASS_ON,
+    RLS_BYPASS_VAR,
+    RLS_WORKSPACE_VAR,
+    TableEditorService,
+    is_internal_column_name,
+    sanitize_identifier,
+)
 
 
 class CustomFieldsService(BaseWorkspaceService, ABC):
@@ -50,6 +59,49 @@ def _get_schema_name(self) -> str:
 
         return f"{self._schema_prefix}{self._workspace_uuid.short()}"
 
+    def _full_table_name(self) -> str:
+        """Get the fully-qualified workspace table name."""
+        return f'"{self.schema_name}".{self.sanitized_table_name}'
+
+    def _workspace_tenant_default_sql(self) -> sa.TextClause:
+        """Build the server default expression for the tenant column."""
+        return sa.text(f"'{self._workspace_uuid}'::uuid")
+
+    async def _enable_workspace_rls_for_base_table(self) -> None:
+        """Enable workspace RLS policy on the workspace base table."""
+        conn = await self.session.connection()
+        policy_expr = (
+            f"current_setting('{RLS_BYPASS_VAR}', true) = '{RLS_BYPASS_ON}' "
+            f'OR "{DYNAMIC_WORKSPACE_TENANT_COLUMN}" = '
+            f"NULLIF(current_setting('{RLS_WORKSPACE_VAR}', true), '')::uuid"
+        )
+        full_table_name = self._full_table_name()
+        await conn.execute(
+            sa.DDL("ALTER TABLE %s ENABLE ROW LEVEL SECURITY", full_table_name)
+        )
+        await conn.execute(
+            sa.DDL(
+                f"DROP POLICY IF EXISTS {DYNAMIC_WORKSPACE_RLS_POLICY} ON %s",
+                full_table_name,
+            )
+        )
+        await conn.execute(
+            sa.DDL(
+                f"""
+                CREATE POLICY {DYNAMIC_WORKSPACE_RLS_POLICY} ON %s
+                    FOR ALL
+                    USING ({policy_expr})
+                    WITH CHECK ({policy_expr})
+                """,
+                full_table_name,
+            )
+        )
+
+    def _assert_user_field_name_allowed(self, field_name: str) -> None:
+        """Reject operations on internal/system-managed field names."""
+        if is_internal_column_name(field_name):
+            raise ValueError(f"Field {field_name} is reserved for internal use")
+
     @abstractmethod
     def _table_definition(self) -> sa.Table:
         """Return the SQLAlchemy Table definition for the workspace table.
@@ -74,6 +126,7 @@ async def initialize_workspace_schema(self) -> None:
             await self.session.execute(
                 CreateTable(self._table_definition(), if_not_exists=True)
             )
+            await self._enable_workspace_rls_for_base_table()
         self._schema_initialized = True
 
     async def _ensure_schema_ready(self) -> None:
@@ -108,12 +161,14 @@ async def list_fields(self) -> Sequence[sa.engine.interfaces.ReflectedColumn]:
         """List all custom fields for the workspace."""
 
         await self._ensure_schema_ready()
-        return await self.editor.get_columns()
+        columns = await self.editor.get_columns()
+        return [col for col in columns if not is_internal_column_name(col["name"])]
 
     async def create_field(self, params: CustomFieldCreate) -> None:
         """Create a new custom field column."""
 
         await self._ensure_schema_ready()
+        self._assert_user_field_name_allowed(params.name)
         params.nullable = True  # Custom fields remain nullable by default
         await self.editor.create_column(params)
         await self.session.commit()
@@ -122,6 +177,9 @@ async def update_field(self, field_id: str, params: CustomFieldUpdate) -> None:
         """Update a custom field column."""
 
         await self._ensure_schema_ready()
+        self._assert_user_field_name_allowed(field_id)
+        if params.name is not None:
+            self._assert_user_field_name_allowed(params.name)
         await self.editor.update_column(field_id, params)
         await self.session.commit()
 
@@ -136,6 +194,7 @@ async def delete_field(self, field_id: str) -> None:
         """
 
         await self._ensure_schema_ready()
+        self._assert_user_field_name_allowed(field_id)
         if field_id in self._reserved_columns:
             raise ValueError(f"Field {field_id} is a reserved field")
         await self.editor.delete_column(field_id)
diff --git a/tracecat/tables/service.py b/tracecat/tables/service.py
