feat(studio): add Warehouse table management and disable (#50195)

## What kind of change does this PR introduce? Feature and UI polish. ## What is the current behavior? Warehouse setup uses a schema accordion for table selection. Once Warehouse is enabled, users cannot remove replicated tables or disable Warehouse from Studio. ## What is the new behavior? - Replaces the schema accordion with one grouped, searchable table selector. - Still allows for **Select all** and **Clear** actions for each schema. - Starts first-time setup with no tables selected and preselects current replicated tables when editing. - Adds support for removing previously replicated tables. - Adds a confirmed **Disable Warehouse** action. - Tracks successful Warehouse enable and disable actions. Disabling Warehouse removes its replication pipeline, publication, catalogue access, and foreign tables. Copied data remains in DuckLake storage until the user deletes it. Re-enabling a table rebuilds its data rather than reusing the retained copy. | Before | After | | --- | --- | | <img width="1024" height="759" alt="Integrations Test US East 1 testdw Supabase" src="https://github.com/user-attachments/assets/bded025b-1d45-41dc-8a35-9159baf8f9b7" /> | <img width="1024" height="759" alt="Integrations test Teamer Supabase" src="https://github.com/user-attachments/assets/69026d94-98a0-4878-ab58-2e9697296d93" /> | | <img width="1280" height="1323" alt="Integrations Test testdw Supabase" src="https://github.com/user-attachments/assets/3f71e754-1a87-4d58-a7b9-dd39d3e0ac5a" /> | <img width="1280" height="1323" alt="Integrations Regular AWS Teamer Supabase" src="https://github.com/user-attachments/assets/758ed48e-9ed6-45d3-ae94-e171147a21d5" /> | | _Feature did not exist_ | <img width="1024" height="759" alt="Integrations Regular AWS Teamer Supabase" src="https://github.com/user-attachments/assets/c977ac57-8b0c-4482-882b-69ad7602b5df" /> | ## Additional context Platform support for updating and disabling Warehouse was added in [supabase/platform#38190](https://github.com/supabase/platform/pull/38190). ### To test 1. Open `/project/{ref}/integrations/warehouse/overview` before setup. 2. Confirm **Tables to replicate** starts at zero and **Enable Warehouse** is disabled until a table is selected. 3. Confirm each schema's **Select all** and **Clear** actions update every table in that schema. 4. Enable Warehouse with a partial selection and wait for setup to complete. 5. Edit the selection, add and remove replicated tables, then confirm the saved selection is reflected in the publication. 6. Disable Warehouse, confirm the retention warning, and verify the integration returns to its initial state. 7. Re-enable Warehouse and confirm selected tables are rebuilt. 8. Trigger a replication pipeline limit error and confirm the inline guidance links to Database Replication.  ## Summary by CodeRabbit - **New Features** - Added the ability to disable Warehouse from the setup panel. - Warehouse setup now starts with no table selections. - Editing a setup preselects replicated tables and supports updating selections, including removing tables. - Added searchable schema and table selection with screen-reader count announcements. - Added telemetry tracking for initial Warehouse enablement. - **Bug Fixes** - Warehouse disable failures now show an error while keeping the confirmation dialog open for retry. - Configuration updates now refresh related data automatically.
2026-10-05 09:25:06 +03:00 · 2026-09-17 14:17:00 +10:00
1 parent 7880c2f079
commit 3e2d54eccb
@@ -1,10 +1,11 @@
-import { screen } from '@testing-library/react'
+import { fireEvent, screen, waitFor, within } from '@testing-library/react'
 import userEvent from '@testing-library/user-event'
 import { platformComponents as components } from 'api-types'
 import dayjs from 'dayjs'
 import duration from 'dayjs/plugin/duration'
 import { mockAnimationsApi } from 'jsdom-testing-mocks'
 import { HttpResponse } from 'msw'
+import { toast } from 'sonner'
 import { beforeEach, describe, expect, test, vi } from 'vitest'

 import { WarehouseOverviewTab } from './OverviewTab'
@@ -12,25 +13,47 @@ import { customRender } from '@/tests/lib/custom-render'
 import { addAPIMock, type APIErrorBody } from '@/tests/lib/msw'

 type WarehouseSetupStatusResponse = components['schemas']['WarehouseSetupStatusResponse_Output']
+type WarehouseSetupBody = components['schemas']['WarehouseSetupBody']
+type WarehouseSetupResponse = components['schemas']['WarehouseSetupResponse_Output']

 // Both integration shells are live, and the flag reads a context plus ConfigCat that
 // `customRender` doesn't provide.
 const mockIsMarketplaceEnabled = vi.fn(() => false)
+const mockTrack = vi.fn()
 vi.mock('@/components/interfaces/App/FeaturePreview/FeaturePreviewContext', () => ({
  useIsMarketplaceEnabled: () => mockIsMarketplaceEnabled(),
 }))
+vi.mock('@/lib/telemetry/track', () => ({ useTrack: () => mockTrack }))

 vi.mock('../Integration/IntegrationOverviewTab', () => ({
  IntegrationOverviewTab: ({ children }: { children: React.ReactNode }) => <div>{children}</div>,
 }))

+vi.mock('sonner', () => ({
+  toast: { success: vi.fn(), error: vi.fn() },
+}))
+
 // Exercised by its own unit tests, and it fires four upstream queries of its own.
 vi.mock('./WarehouseSchemaTablePicker', () => ({
-  WarehouseSchemaTablePicker: ({ error }: { error?: { message: string } | null }) => (
+  WarehouseSchemaTablePicker: ({
+    error,
+    isEditing,
+    onSubmit,
+  }: {
+    error?: { message: string } | null
+    isEditing?: boolean
+    onSubmit: (targets: [{ type: 'table'; schema: string; name: string }]) => void
+  }) => (
    <section>
      <h2>Tables</h2>
      <span>Replicated tables picker</span>
      {!!error && <span>Picker error: {error.message}</span>}
+      <button
+        tabIndex={0}
+        onClick={() => onSubmit([{ type: 'table', schema: 'public', name: 'orders' }])}
+      >
+        {isEditing ? 'Submit edited tables' : 'Submit initial tables'}
+      </button>
    </section>
  ),
 }))
@@ -81,6 +104,7 @@ const mockProject = () =>
 describe('WarehouseOverviewTab', () => {
  beforeEach(() => {
    mockIsMarketplaceEnabled.mockReturnValue(true)
+    mockTrack.mockClear()
  })

  test.each([false, true])(
@@ -167,7 +191,7 @@ describe('WarehouseOverviewTab', () => {
    expect(screen.queryByText('Initial setup failed')).not.toBeInTheDocument()
  })

-  test('shows Status, Tables, then Connect once setup is complete', async () => {
+  test('shows Status, Tables, Connect, then Disable once setup is complete', async () => {
    mockSetupStatus({
      setup_status: 'complete',
      tables: [
@@ -192,7 +216,7 @@ describe('WarehouseOverviewTab', () => {
    customRender(<WarehouseOverviewTab />)

    // findByRole throws on duplicates, so this also guards the section titles staying distinct.
-    for (const name of ['Status', 'Tables', 'Connect']) {
+    for (const name of ['Status', 'Tables', 'Connect', 'Disable']) {
      expect(await screen.findByRole('heading', { name })).toBeInTheDocument()
    }
    expect(screen.getByText('Replicated tables picker')).toBeInTheDocument()
@@ -203,10 +227,110 @@ describe('WarehouseOverviewTab', () => {
    const headings = screen
      .getAllByRole('heading')
      .map((heading) => heading.textContent)
-      .filter((heading) => ['Status', 'Tables', 'Connect'].includes(heading ?? ''))
+      .filter((heading) => ['Status', 'Tables', 'Connect', 'Disable'].includes(heading ?? ''))

-    expect(headings).toEqual(['Status', 'Tables', 'Connect'])
-    expect(screen.queryByRole('button', { name: 'Disable Warehouse' })).not.toBeInTheDocument()
+    expect(headings).toEqual(['Status', 'Tables', 'Connect', 'Disable'])
+  })
+
+  test('tracks initial setup but not table selection edits as enablement', async () => {
+    mockSetupStatus({ setup_status: 'not_started' })
+    addAPIMock({
+      method: 'post',
+      path: '/platform/warehouse/:ref/setup',
+      response: () => HttpResponse.json<WarehouseSetupResponse>({ pipeline_id: 1, tables: [] }),
+    })
+
+    customRender(<WarehouseOverviewTab />)
+
+    await userEvent.click(await screen.findByRole('button', { name: 'Submit initial tables' }))
+    await waitFor(() =>
+      expect(mockTrack).toHaveBeenCalledWith('warehouse_enabled', {
+        source: 'integrations_overview',
+        schemaTargetCount: 0,
+        tableTargetCount: 1,
+      })
+    )
+  })
+
+  test('does not track an edited table selection as enablement', async () => {
+    mockSetupStatus({ setup_status: 'complete' })
+    addAPIMock({
+      method: 'post',
+      path: '/platform/warehouse/:ref/setup',
+      response: () => HttpResponse.json<WarehouseSetupResponse>({ pipeline_id: 1, tables: [] }),
+    })
+
+    customRender(<WarehouseOverviewTab />)
+
+    await userEvent.click(await screen.findByRole('button', { name: 'Submit edited tables' }))
+    await waitFor(() => expect(screen.getByText('Replicated tables picker')).toBeInTheDocument())
+    expect(mockTrack).not.toHaveBeenCalledWith('warehouse_enabled', expect.anything())
+  })
+
+  test('disables Warehouse with an empty target list after confirmation', async () => {
+    mockSetupStatus({ setup_status: 'complete' })
+    const setupRequests: WarehouseSetupBody[] = []
+    addAPIMock({
+      method: 'post',
+      path: '/platform/warehouse/:ref/setup',
+      response: async ({ request }) => {
+        setupRequests.push((await request.json()) as WarehouseSetupBody)
+        return HttpResponse.json<WarehouseSetupResponse>({ pipeline_id: 1, tables: [] })
+      },
+    })
+
+    customRender(<WarehouseOverviewTab />)
+
+    await userEvent.click(await screen.findByRole('button', { name: 'Disable Warehouse' }))
+    expect(setupRequests).toEqual([])
+
+    const dialog = await screen.findByRole('alertdialog')
+    expect(dialog).toHaveTextContent('Copied data remains in DuckLake storage until deleted')
+    expect(screen.queryByPlaceholderText('Type the project ref to confirm')).not.toBeInTheDocument()
+
+    fireEvent.click(within(dialog).getByRole('button', { name: 'Disable Warehouse' }))
+
+    await waitFor(() => expect(setupRequests).toEqual([{ targets: [] }]))
+  })
+
+  test('shows a disable error and allows retrying from the open confirmation', async () => {
+    mockSetupStatus({ setup_status: 'complete' })
+    let attempts = 0
+    addAPIMock({
+      method: 'post',
+      path: '/platform/warehouse/:ref/setup',
+      response: () => {
+        attempts += 1
+        if (attempts === 1) {
+          return HttpResponse.json<APIErrorBody>(
+            { message: 'Disable request failed' },
+            { status: 500 }
+          )
+        }
+        return HttpResponse.json<WarehouseSetupResponse>({ pipeline_id: 1, tables: [] })
+      },
+    })
+
+    customRender(<WarehouseOverviewTab />)
+
+    await userEvent.click(await screen.findByRole('button', { name: 'Disable Warehouse' }))
+    const dialog = await screen.findByRole('alertdialog')
+    const confirm = within(dialog).getByRole('button', { name: 'Disable Warehouse' })
+
+    fireEvent.click(confirm)
+
+    await waitFor(() =>
+      expect(toast.error).toHaveBeenCalledWith(
+        'Failed to disable Warehouse: Disable request failed'
+      )
+    )
+    expect(dialog).toBeVisible()
+    expect(confirm).toBeEnabled()
+
+    fireEvent.click(confirm)
+
+    await waitFor(() => expect(attempts).toBe(2))
+    await waitFor(() => expect(screen.queryByRole('alertdialog')).not.toBeInTheDocument())
  })

  test('shows a status query failure without blocking an unrelated route', async () => {
@@ -413,7 +413,7 @@ export const WarehouseConnectionCard = ({ variant = 'default' }: WarehouseConnec

 export const WarehouseConnectSection = () => {
  return (
-    <PageSection className="pt-0!">
+    <PageSection className="pt-5!">
      <PageSectionMeta>
        <PageSectionSummary>
          <PageSectionTitle>Connect</PageSectionTitle>
@@ -0,0 +1,108 @@
+import { useQueryClient } from '@tanstack/react-query'
+import { useParams } from 'common'
+import { useState } from 'react'
+import { toast } from 'sonner'
+import {
+  AlertDialog,
+  AlertDialogAction,
+  AlertDialogBody,
+  AlertDialogCancel,
+  AlertDialogContent,
+  AlertDialogDescription,
+  AlertDialogFooter,
+  AlertDialogHeader,
+  AlertDialogTitle,
+  Button,
+  Card,
+  CardContent,
+} from 'ui'
+import { FormLayout } from 'ui-patterns/form/Layout/FormLayout'
+import {
+  PageSection,
+  PageSectionContent,
+  PageSectionMeta,
+  PageSectionSummary,
+  PageSectionTitle,
+} from 'ui-patterns/PageSection'
+
+import { warehouseKeys } from '@/data/warehouse/keys'
+import { useWarehouseSetupMutation } from '@/data/warehouse/warehouse-setup-mutation'
+import { useTrack } from '@/lib/telemetry/track'
+
+export const WarehouseDisableCard = () => {
+  const { ref: projectRef } = useParams()
+  const queryClient = useQueryClient()
+  const track = useTrack()
+  const [isConfirming, setIsConfirming] = useState(false)
+
+  const setupMutation = useWarehouseSetupMutation({
+    onSuccess: async () => {
+      track('warehouse_disabled', {})
+      // Disabling also turns off catalog access server-side, which the setup mutation doesn't know
+      // to invalidate on its own.
+      await queryClient.invalidateQueries({ queryKey: warehouseKeys.catalog(projectRef) })
+      setIsConfirming(false)
+      toast.success('Warehouse disabled')
+    },
+    onError: (error) => toast.error(`Failed to disable Warehouse: ${error.message}`),
+  })
+
+  return (
+    <PageSection className="pt-5!">
+      <PageSectionMeta>
+        <PageSectionSummary>
+          <PageSectionTitle>Disable</PageSectionTitle>
+        </PageSectionSummary>
+      </PageSectionMeta>
+      <PageSectionContent>
+        <Card>
+          <CardContent>
+            <FormLayout
+              layout="flex-row-reverse"
+              label="Disable Warehouse for this project"
+              description="Stops replication and removes Warehouse-managed resources."
+            >
+              <Button variant="danger" onClick={() => setIsConfirming(true)}>
+                Disable Warehouse
+              </Button>
+            </FormLayout>
+          </CardContent>
+        </Card>
+      </PageSectionContent>
+
+      <AlertDialog open={isConfirming} onOpenChange={setIsConfirming}>
+        <AlertDialogContent size="small">
+          <AlertDialogHeader>
+            <AlertDialogTitle>Disable Warehouse</AlertDialogTitle>
+          </AlertDialogHeader>
+          <AlertDialogBody>
+            <AlertDialogDescription asChild>
+              <div className="space-y-2">
+                <p>
+                  Disabling Warehouse stops replication and connections to its endpoint. Its
+                  pipeline, publication, catalog access, and foreign tables are removed.
+                </p>
+                <p>Copied data remains in DuckLake storage until deleted.</p>
+              </div>
+            </AlertDialogDescription>
+          </AlertDialogBody>
+          <AlertDialogFooter>
+            <AlertDialogCancel>Cancel</AlertDialogCancel>
+            <AlertDialogAction
+              variant="danger"
+              loading={setupMutation.isPending}
+              disabled={!projectRef}
+              onClick={() =>
+                projectRef
+                  ? setupMutation.mutateAsync({ projectRef, body: { targets: [] } })
+                  : undefined
+              }
+            >
+              Disable Warehouse
+            </AlertDialogAction>
+          </AlertDialogFooter>
+        </AlertDialogContent>
+      </AlertDialog>
+    </PageSection>
+  )
+}
@@ -0,0 +1,291 @@
+import { fireEvent, screen, waitFor } from '@testing-library/react'
+import userEvent from '@testing-library/user-event'
+import type { components } from 'api-types'
+import { mockAnimationsApi } from 'jsdom-testing-mocks'
+import { HttpResponse } from 'msw'
+import { describe, expect, test } from 'vitest'
+
+import { WarehouseSchemaTablePicker } from './WarehouseSchemaTablePicker'
+import type { SchemasData } from '@/data/database/schemas-query'
+import type { ProjectDetail } from '@/data/projects/project-detail-query'
+import type { TablesData } from '@/data/tables/tables-query'
+import { useWarehouseSetupMutation } from '@/data/warehouse/warehouse-setup-mutation'
+import { customRender } from '@/tests/lib/custom-render'
+import { addAPIMock } from '@/tests/lib/msw'
+
+type PublicationDetailsResponse = components['schemas']['PublicationDetailsResponse_Output']
+type ReplicationSourcesResponse = components['schemas']['SourcesResponse_Output']
+type RunQueryBody = components['schemas']['RunQueryBody']
+type WarehouseSetupBody = components['schemas']['WarehouseSetupBody']
+type WarehouseSetupResponse = components['schemas']['WarehouseSetupResponse_Output']
+
+mockAnimationsApi()
+
+const PROJECT: ProjectDetail = {
+  cloud_provider: 'AWS',
+  connectionString: 'postgresql://postgres@localhost:5432/postgres',
+  db_host: 'db.default.supabase.co',
+  high_availability: false,
+  id: 1,
+  inserted_at: '2026-01-01T00:00:00.000Z',
+  integration_source: null,
+  is_branch_enabled: false,
+  is_hibernating: false,
+  is_physical_backups_enabled: false,
+  name: 'Test project',
+  organization_id: 1,
+  ref: 'default',
+  region: 'us-east-1',
+  restUrl: 'https://default.supabase.co/rest/v1',
+  status: 'ACTIVE_HEALTHY',
+  subscription_id: 'subscription-1',
+  updated_at: '2026-01-01T00:00:00.000Z',
+}
+
+const SCHEMAS: SchemasData = [
+  { id: 1, name: 'public', owner: 'postgres', comment: null },
+  { id: 2, name: 'analytics', owner: 'postgres', comment: null },
+]
+
+const TABLES: TablesData = [
+  {
+    id: 1,
+    schema: 'public',
+    name: 'orders',
+    rls_enabled: false,
+    rls_forced: false,
+    replica_identity: 'DEFAULT',
+    bytes: 1024,
+    size: '1024 bytes',
+    live_rows_estimate: 10,
+    dead_rows_estimate: 0,
+    comment: null,
+    primary_keys: [],
+    relationships: [],
+  },
+  {
+    id: 2,
+    schema: 'public',
+    name: 'customers',
+    rls_enabled: false,
+    rls_forced: false,
+    replica_identity: 'DEFAULT',
+    bytes: 1024,
+    size: '1024 bytes',
+    live_rows_estimate: 10,
+    dead_rows_estimate: 0,
+    comment: null,
+    primary_keys: [],
+    relationships: [],
+  },
+  {
+    id: 3,
+    schema: 'analytics',
+    name: 'events',
+    rls_enabled: false,
+    rls_forced: false,
+    replica_identity: 'DEFAULT',
+    bytes: 1024,
+    size: '1024 bytes',
+    live_rows_estimate: 10,
+    dead_rows_estimate: 0,
+    comment: null,
+    primary_keys: [],
+    relationships: [],
+  },
+]
+
+const SOURCES: ReplicationSourcesResponse = {
+  sources: [
+    {
+      id: 1,
+      name: 'default',
+      tenant_id: 'tenant',
+      config: { host: 'db.internal', name: 'main-db', port: 5432, username: 'etl_user' },
+    },
+  ],
+}
+
+const PUBLICATION: PublicationDetailsResponse = {
+  name: 'supabase_warehouse',
+  config: {
+    type: 'tables',
+    tables: [
+      { id: 1, schema: 'public', name: 'orders', columns: null, row_filter: null },
+      { id: 2, schema: 'public', name: 'customers', columns: null, row_filter: null },
+    ],
+    operations: ['insert', 'update', 'delete', 'truncate'],
+    publish_via_partition_root: false,
+  },
+  tables: [
+    { id: 1, schema: 'public', name: 'orders', kind: 'table', partition_parent_id: null },
+    { id: 2, schema: 'public', name: 'customers', kind: 'table', partition_parent_id: null },
+  ],
+}
+
+const mockPickerQueries = ({
+  isEditing,
+  tables = TABLES,
+  schemas = SCHEMAS,
+}: {
+  isEditing: boolean
+  tables?: TablesData
+  schemas?: SchemasData
+}) => {
+  addAPIMock({ method: 'get', path: '/platform/projects/:ref', response: PROJECT })
+  addAPIMock({
+    method: 'post',
+    path: '/platform/pg-meta/:ref/query',
+    response: async ({ request }) => {
+      const body = (await request.json()) as RunQueryBody
+      const result = body.query.includes("obj_description(n.oid, 'pg_namespace')")
+        ? schemas
+        : tables
+      return HttpResponse.json<SchemasData | TablesData>(result)
+    },
+  })
+
+  if (isEditing) {
+    addAPIMock({
+      method: 'get',
+      path: '/platform/replication/:ref/sources',
+      response: () => HttpResponse.json<ReplicationSourcesResponse>(SOURCES),
+    })
+    addAPIMock({
+      method: 'get',
+      path: '/platform/replication/v2/:ref/sources/:source_id/publications/:publication_name',
+      response: () => HttpResponse.json<PublicationDetailsResponse>(PUBLICATION),
+    })
+  }
+}
+
+const WarehousePickerHarness = ({ isEditing = false }: { isEditing?: boolean }) => {
+  const setupMutation = useWarehouseSetupMutation()
+  return (
+    <WarehouseSchemaTablePicker
+      isEditing={isEditing}
+      isSubmitting={setupMutation.isPending}
+      onSubmit={(targets) => setupMutation.mutate({ projectRef: 'default', body: { targets } })}
+    />
+  )
+}
+
+const mockSetupMutation = () => {
+  const requests: WarehouseSetupBody[] = []
+  addAPIMock({
+    method: 'post',
+    path: '/platform/warehouse/:ref/setup',
+    response: async ({ request }) => {
+      requests.push((await request.json()) as WarehouseSetupBody)
+      return HttpResponse.json<WarehouseSetupResponse>({ pipeline_id: 1, tables: [] })
+    },
+  })
+  return requests
+}
+
+describe('WarehouseSchemaTablePicker', () => {
+  test('starts empty and supports selecting every table in a schema', async () => {
+    mockPickerQueries({ isEditing: false })
+    const requests = mockSetupMutation()
+
+    customRender(<WarehousePickerHarness />)
+
+    expect(await screen.findByText('0 tables selected')).toBeInTheDocument()
+    expect(screen.getByRole('button', { name: 'Enable Warehouse' })).toBeDisabled()
+
+    fireEvent.click(screen.getByRole('combobox', { name: 'Select tables to replicate' }))
+    await userEvent.click(screen.getAllByText('Select all')[1])
+
+    expect(screen.getByText('2 tables selected')).toBeInTheDocument()
+    await userEvent.click(screen.getByRole('button', { name: 'Enable Warehouse' }))
+
+    await waitFor(() =>
+      expect(requests).toEqual([{ targets: [{ type: 'schema', schema: 'public' }] }])
+    )
+  })
+
+  test('preselects replicated tables and submits a replacement after removal', async () => {
+    mockPickerQueries({ isEditing: true })
+    const requests = mockSetupMutation()
+
+    customRender(<WarehousePickerHarness isEditing />)
+
+    expect(await screen.findByText('2 tables selected')).toBeInTheDocument()
+    fireEvent.click(screen.getByRole('combobox', { name: 'Select tables to replicate' }))
+    await userEvent.click(screen.getByText('customers'))
+
+    expect(screen.getByText('1 table selected')).toBeInTheDocument()
+    await userEvent.click(screen.getByRole('button', { name: 'Update replicated tables' }))
+
+    await waitFor(() =>
+      expect(requests).toEqual([{ targets: [{ type: 'table', schema: 'public', name: 'orders' }] }])
+    )
+  })
+
+  test('filters tables by schema name or qualified table name', async () => {
+    mockPickerQueries({ isEditing: false })
+
+    customRender(<WarehousePickerHarness />)
+
+    fireEvent.click(await screen.findByRole('combobox', { name: 'Select tables to replicate' }))
+    await userEvent.type(screen.getByPlaceholderText('Search schemas and tables...'), 'analytics')
+
+    expect(screen.getByText('events')).toBeInTheDocument()
+    expect(screen.queryByText('orders')).not.toBeInTheDocument()
+    expect(screen.queryByText('customers')).not.toBeInTheDocument()
+
+    await userEvent.clear(screen.getByPlaceholderText('Search schemas and tables...'))
+    await userEvent.type(screen.getByPlaceholderText('Search schemas and tables...'), 'public.ord')
+
+    expect(screen.getByText('orders')).toBeInTheDocument()
+    expect(screen.queryByText('customers')).not.toBeInTheDocument()
+    expect(screen.queryByText('events')).not.toBeInTheDocument()
+
+    await userEvent.clear(screen.getByPlaceholderText('Search schemas and tables...'))
+    await userEvent.type(screen.getByPlaceholderText('Search schemas and tables...'), 'orders')
+
+    expect(screen.getByText('orders')).toBeInTheDocument()
+    expect(screen.queryByText('customers')).not.toBeInTheDocument()
+    expect(screen.queryByText('events')).not.toBeInTheDocument()
+  })
+
+  test('wraps selected table badges and shows an overflow count past the limit', async () => {
+    const manyTables: TablesData = Array.from({ length: 12 }, (_, index) => ({
+      id: index + 1,
+      schema: 'public',
+      name: `table_${String(index + 1).padStart(2, '0')}`,
+      rls_enabled: false,
+      rls_forced: false,
+      replica_identity: 'DEFAULT',
+      bytes: 1024,
+      size: '1024 bytes',
+      live_rows_estimate: 10,
+      dead_rows_estimate: 0,
+      comment: null,
+      primary_keys: [],
+      relationships: [],
+    }))
+
+    mockPickerQueries({
+      isEditing: false,
+      schemas: [{ id: 1, name: 'public', owner: 'postgres', comment: null }],
+      tables: manyTables,
+    })
+
+    customRender(<WarehousePickerHarness />)
+
+    const trigger = await screen.findByRole('combobox', { name: 'Select tables to replicate' })
+    fireEvent.click(trigger)
+    await userEvent.click(screen.getByText('Select all'))
+    // Close the list so trigger textContent is only the selected badges.
+    fireEvent.click(trigger)
+
+    expect(screen.getByText('12 tables selected')).toBeInTheDocument()
+    expect(trigger.firstElementChild).toHaveClass('flex-wrap')
+    expect(trigger).toHaveTextContent('public.table_01')
+    expect(trigger).toHaveTextContent('public.table_10')
+    expect(trigger).toHaveTextContent('+2')
+    expect(trigger).not.toHaveTextContent('public.table_11')
+    expect(trigger).not.toHaveTextContent('public.table_12')
+  })
+})
@@ -1,16 +1,8 @@
 import { useParams } from 'common'
-import { ChevronRight } from 'lucide-react'
 import { useMemo, useState } from 'react'
-import {
-  Button,
-  Card,
-  CardContent,
-  CardFooter,
-  Checkbox,
-  Collapsible,
-  CollapsibleContent,
-  CollapsibleTrigger,
-} from 'ui'
+import { Button, Card, CardContent, CardFooter, CommandGroup } from 'ui'
+import { FormLayout } from 'ui-patterns/form/Layout/FormLayout'
+import { MultiSelector } from 'ui-patterns/multi-select'
 import {
  PageSection,
  PageSectionContent,
@@ -24,7 +16,6 @@ import { GenericSkeletonLoader } from 'ui-patterns/ShimmeringLoader'
 import {
  buildSelectionFromPublicationTables,
  buildWarehouseSetupTargets,
-  getSchemaCheckedState,
  getSchemaTableKey,
  getSelectedTableCount,
  hasSelectionChanged,
@@ -40,12 +31,15 @@ import { useReplicationSourcesQuery } from '@/data/replication/sources-query'
 import { useTablesQuery } from '@/data/tables/tables-query'
 import { useSelectedProjectQuery } from '@/hooks/misc/useSelectedProject'
 import { WAREHOUSE_PUBLICATION_NAME } from '@/lib/warehouse'
+import type { ResponseError } from '@/types'

 export interface WarehouseSchemaTablePickerProps {
  onSubmit: (targets: WarehouseSetupTarget[]) => void
  isSubmitting: boolean
+  /** Set when editing an already-enabled Warehouse rather than setting one up for the first time. */
  isEditing?: boolean
-  error?: { message: string } | null
+  /** Failure from the submit itself, rendered inline rather than as a toast the user can lose. */
+  error?: ResponseError | null
 }

 export const WarehouseSchemaTablePicker = ({
@@ -60,7 +54,6 @@ export const WarehouseSchemaTablePicker = ({
  // `null` until the user touches a checkbox, so the selection seeded from the existing
  // publication can arrive asynchronously without an effect syncing it into state.
  const [selectionOverride, setSelectionOverride] = useState<SchemaTableSelection | null>(null)
-  const [expandedOverrides, setExpandedOverrides] = useState<Record<string, boolean>>({})

  const {
    data: schemas,
@@ -85,32 +78,34 @@ export const WarehouseSchemaTablePicker = ({
    isLoading: isSourcesLoading,
    isError: isSourcesError,
    error: sourcesError,
-  } = useReplicationSourcesQuery({ projectRef })
+  } = useReplicationSourcesQuery({ projectRef }, { enabled: isEditing })
  const sourceId = sourcesData?.sources.find((source) => source.name === projectRef)?.id

  const {
    data: publication,
    isError: isPublicationsError,
    error: publicationsError,
-  } = useReplicationPublicationQuery({
-    projectRef,
-    sourceId,
-    publicationName: WAREHOUSE_PUBLICATION_NAME,
-  })
+  } = useReplicationPublicationQuery(
+    {
+      projectRef,
+      sourceId,
+      publicationName: WAREHOUSE_PUBLICATION_NAME,
+    },
+    { enabled: isEditing }
+  )

  // Derived from data presence rather than fetch status, so there's no render gap between the
  // publication query becoming enabled and it actually starting to fetch.
  const isSelectionPending =
-    isSourcesLoading ||
-    (sourceId !== undefined && publication === undefined && !isPublicationsError)
+    isEditing &&
+    (isSourcesLoading ||
+      (sourceId !== undefined && publication === undefined && !isPublicationsError))

  const initialSelection = useMemo(
-    () => buildSelectionFromPublicationTables(publication?.tables ?? []),
-    [publication]
+    () => (isEditing ? buildSelectionFromPublicationTables(publication?.tables ?? []) : {}),
+    [isEditing, publication]
  )

-  const selection = selectionOverride ?? initialSelection
-
  const schemasWithTables: SchemaWithTables[] = useMemo(() => {
    if (!schemas || !tables) return []
    return schemas
@@ -122,38 +117,26 @@ export const WarehouseSchemaTablePicker = ({
      .sort((a, b) => a.schema.localeCompare(b.schema))
  }, [schemas, tables])

+  const tableKeys = schemasWithTables.flatMap(({ schema, tables }) =>
+    tables.map((table) => getSchemaTableKey(schema, table))
+  )
+  const selection = selectionOverride ?? initialSelection
  const selectedCount = getSelectedTableCount(selection)
+  const selectedTableKeys = tableKeys.filter((key) => selection[key])
  const hasChanges = hasSelectionChanged(selection, initialSelection)

  const updateSelection = (updater: (current: SchemaTableSelection) => SchemaTableSelection) => {
    setSelectionOverride((prev) => updater(prev ?? initialSelection))
  }

-  const toggleTable = (schema: string, table: string) => {
-    const key = getSchemaTableKey(schema, table)
-    if (isEditing && initialSelection[key]) return
-    updateSelection((current) => ({ ...current, [key]: !current[key] }))
-  }
+  const handleSchemaSelection = (schema: SchemaWithTables) => {
+    const keys = schema.tables.map((table) => getSchemaTableKey(schema.schema, table))
+    const areAllSelected = keys.every((key) => selection[key])

-  const toggleSchema = (schema: SchemaWithTables) => {
-    const selectableTables = schema.tables.filter(
-      (table) => !isEditing || !initialSelection[getSchemaTableKey(schema.schema, table)]
-    )
-    const areAllSelectableTablesSelected =
-      selectableTables.length > 0 &&
-      selectableTables.every((table) => selection[getSchemaTableKey(schema.schema, table)])
-
-    updateSelection((current) => {
-      const next = { ...current }
-      selectableTables.forEach((table) => {
-        next[getSchemaTableKey(schema.schema, table)] = !areAllSelectableTablesSelected
-      })
-      return next
-    })
-  }
-
-  const setExpanded = (schemaName: string, isOpen: boolean) => {
-    setExpandedOverrides((prev) => ({ ...prev, [schemaName]: isOpen }))
+    updateSelection((current) => ({
+      ...current,
+      ...Object.fromEntries(keys.map((key) => [key, !areAllSelected])),
+    }))
  }

  const handleSubmit = () => {
@@ -162,18 +145,18 @@ export const WarehouseSchemaTablePicker = ({
    onSubmit(targets)
  }

-  // Waiting on the publication too, so the pre-checked selection is in place before the user can
-  // start toggling (an early toggle would otherwise pin an override that omits existing tables).
+  // When editing, wait for the publication so its selection is in place before the user can
+  // start toggling. First-time setup always starts empty and does not need to wait for it.
  if (isSchemasPending || isTablesPending || isSelectionPending) {
    return (
-      <PageSection className="pt-0!">
+      <PageSection className="pt-5!">
        <PageSectionMeta>
          <PageSectionSummary>
            <PageSectionTitle>Tables</PageSectionTitle>
            <PageSectionDescription>
              {isEditing
-                ? 'Add more schemas or tables to Warehouse.'
-                : 'Choose which schemas or tables to replicate.'}
+                ? 'Tables currently replicating are selected. Changes apply on save.'
+                : 'Choose which schemas or tables to replicate. You can change this at any time.'}
            </PageSectionDescription>
          </PageSectionSummary>
        </PageSectionMeta>
@@ -187,129 +170,160 @@ export const WarehouseSchemaTablePicker = ({
      </PageSection>
    )
  }
-  if (isSchemasError) return <AlertError subject="Failed to load schemas" error={schemasError} />
-  if (isTablesError) return <AlertError subject="Failed to load tables" error={tablesError} />
+  if (isSchemasError) {
+    return (
+      <AlertError projectRef={projectRef} subject="Failed to load schemas" error={schemasError} />
+    )
+  }
+  if (isTablesError) {
+    return (
+      <AlertError projectRef={projectRef} subject="Failed to load tables" error={tablesError} />
+    )
+  }
  if (isEditing && isSourcesError) {
-    return <AlertError subject="Failed to load replicated tables" error={sourcesError} />
+    return (
+      <AlertError
+        projectRef={projectRef}
+        subject="Failed to load replicated tables"
+        error={sourcesError}
+      />
+    )
  }
  // Only blocking when editing: a first-time setup starts from an empty selection anyway, so a
  // failed publication lookup shouldn't stop the user from enabling Warehouse at all.
  if (isEditing && isPublicationsError) {
-    return <AlertError subject="Failed to load replicated tables" error={publicationsError} />
+    return (
+      <AlertError
+        projectRef={projectRef}
+        subject="Failed to load replicated tables"
+        error={publicationsError}
+      />
+    )
  }

  return (
-    <PageSection className="pt-0!">
+    <PageSection className="pt-5!">
      <PageSectionMeta>
        <PageSectionSummary>
          <PageSectionTitle>Tables</PageSectionTitle>
          <PageSectionDescription>
            {isEditing
-              ? 'Add more schemas or tables to Warehouse. Tables already replicating cannot be removed yet.'
-              : 'Choose which schemas or tables to replicate.'}
+              ? 'Tables currently replicating are selected. Changes apply on save.'
+              : 'Choose which schemas or tables to replicate. You can change this at any time.'}
          </PageSectionDescription>
        </PageSectionSummary>
      </PageSectionMeta>
      <PageSectionContent className="space-y-4">
        {!!error && (
          <AlertError
-            subject={isEditing ? 'Failed to add replicated tables' : 'Failed to enable Warehouse'}
+            subject={
+              isEditing ? 'Failed to update replicated tables' : 'Failed to enable Warehouse'
+            }
            error={error}
+            projectRef={projectRef}
          />
        )}
        <Card>
-          <CardContent className="space-y-3">
-            <p className="text-sm font-medium text-foreground-light">
-              Schemas and tables to replicate
-            </p>
-            <div className="overflow-hidden rounded-md border divide-y">
-              {schemasWithTables.map((schema) => {
-                const keys = schema.tables.map((table) => getSchemaTableKey(schema.schema, table))
-                const checkedCount = keys.filter((key) => selection[key]).length
-                const checkedState = getSchemaCheckedState({
-                  selectedCount: checkedCount,
-                  totalCount: keys.length,
-                })
-                const isOpen = expandedOverrides[schema.schema] ?? checkedCount > 0
-                const selectableTableCount = schema.tables.filter(
-                  (table) =>
-                    !isEditing || !initialSelection[getSchemaTableKey(schema.schema, table)]
-                ).length
+          <CardContent>
+            <FormLayout
+              layout="horizontal"
+              label="Tables to replicate"
+              description={
+                <span aria-live="polite">
+                  {selectedCount} table{selectedCount === 1 ? '' : 's'} selected
+                </span>
+              }
+            >
+              <MultiSelector
+                values={selectedTableKeys}
+                filter={(value, search, keywords) => {
+                  const normalizedSearch = search.toLowerCase()
+                  return [value, ...(keywords ?? [])].some((candidate) =>
+                    candidate.toLowerCase().includes(normalizedSearch)
+                  )
+                    ? 1
+                    : 0
+                }}
+                onValuesChange={(values) => {
+                  const nextValues = new Set(values)
+                  updateSelection(() =>
+                    Object.fromEntries(tableKeys.map((key) => [key, nextValues.has(key)]))
+                  )
+                }}
+                className="w-full"
+              >
+                <MultiSelector.Trigger
+                  aria-label="Select tables to replicate"
+                  label="Select tables..."
+                  // Show plenty of selected tables; wrapBadges keeps the +n count visible when wrapping.
+                  badgeLimit={10}
+                  wrapBadges
+                  renderValue={(value) => <span className="max-w-32 truncate">{value}</span>}
+                  className="w-full"
+                />
+                <MultiSelector.Content>
+                  <MultiSelector.Input placeholder="Search schemas and tables..." showResetIcon />
+                  <MultiSelector.List emptyLabel="No tables available">
+                    {schemasWithTables
+                      .filter((schema) => schema.tables.length > 0)
+                      .map((schema) => {
+                        const keys = schema.tables.map((table) =>
+                          getSchemaTableKey(schema.schema, table)
+                        )
+                        const areAllSelected = keys.every((key) => selection[key])

-                return (
-                  <Collapsible
-                    key={schema.schema}
-                    open={isOpen}
-                    onOpenChange={(open) => setExpanded(schema.schema, open)}
-                  >
-                    <div className="flex items-center gap-2 px-3 py-2 bg-surface-75">
-                      <CollapsibleTrigger
-                        aria-label={
-                          isOpen ? `Collapse ${schema.schema}` : `Expand ${schema.schema}`
-                        }
-                        className="group text-foreground-lighter"
-                      >
-                        <ChevronRight
-                          size={14}
-                          className="transition-transform group-data-[state=open]:rotate-90"
-                        />
-                      </CollapsibleTrigger>
-                      <Checkbox
-                        checked={checkedState}
-                        onCheckedChange={() => toggleSchema(schema)}
-                        disabled={selectableTableCount === 0}
-                        aria-label={`Select all tables in ${schema.schema}`}
-                        // The shared Checkbox only fills itself for `data-state=checked`, so a partial
-                        // selection would otherwise render identically to an empty one. A muted fill
-                        // keeps all three states visually distinct.
-                        className="data-[state=indeterminate]:border-foreground-lighter data-[state=indeterminate]:bg-foreground-lighter"
-                      />
-                      <span className="text-sm font-mono text-foreground">{schema.schema}</span>
-                      <span className="text-xs text-foreground-lighter ml-auto">
-                        {checkedCount}/{keys.length} tables
-                      </span>
-                    </div>
-                    <CollapsibleContent>
-                      {schema.tables.map((table) => {
-                        const key = getSchemaTableKey(schema.schema, table)
                        return (
-                          <div key={key} className="flex items-center gap-2 pl-10 pr-3 py-2">
-                            <Checkbox
-                              checked={!!selection[key]}
-                              disabled={isEditing && !!initialSelection[key]}
-                              onCheckedChange={() => toggleTable(schema.schema, table)}
-                              aria-label={`Select ${schema.schema}.${table}`}
-                            />
-                            <span className="text-sm font-mono text-foreground-light">{table}</span>
-                          </div>
+                          <CommandGroup
+                            key={schema.schema}
+                            heading={
+                              <div className="flex items-center justify-between">
+                                <span>{schema.schema}</span>
+                                <Button
+                                  type="button"
+                                  variant="text"
+                                  size="tiny"
+                                  className="font-sans! normal-case! tracking-normal!"
+                                  onClick={() => handleSchemaSelection(schema)}
+                                >
+                                  {areAllSelected ? 'Clear' : 'Select all'}
+                                </Button>
+                              </div>
+                            }
+                          >
+                            {schema.tables.map((table) => {
+                              const key = getSchemaTableKey(schema.schema, table)
+                              return (
+                                <MultiSelector.Item
+                                  key={key}
+                                  value={key}
+                                  keywords={[schema.schema, key]}
+                                >
+                                  {table}
+                                </MultiSelector.Item>
+                              )
+                            })}
+                          </CommandGroup>
                        )
                      })}
-                      {schema.tables.length === 0 && (
-                        <p className="pl-10 pr-3 py-2 text-sm text-foreground-lighter">
-                          No tables in this schema.
-                        </p>
-                      )}
-                    </CollapsibleContent>
-                  </Collapsible>
-                )
-              })}
-            </div>
+                  </MultiSelector.List>
+                </MultiSelector.Content>
+              </MultiSelector>
+            </FormLayout>
          </CardContent>
-          <CardFooter className="justify-between">
-            <span className="text-sm text-foreground-lighter">
-              {selectedCount} table{selectedCount === 1 ? '' : 's'} selected
-            </span>
-            <div className="flex items-center gap-x-2">
-              {hasChanges && <Button onClick={() => setSelectionOverride(null)}>Cancel</Button>}
-              <Button
-                variant="primary"
-                disabled={selectedCount === 0 || !hasChanges}
-                loading={isSubmitting}
-                onClick={handleSubmit}
-              >
-                {isEditing ? 'Add replicated tables' : 'Enable Warehouse'}
-              </Button>
-            </div>
+          <CardFooter className="justify-end">
+            {/*
+              An empty selection is a valid request that tears Warehouse down, so submitting one
+              from here would destroy a project's Warehouse with no confirmation. Disabling keeps
+              teardown on the dedicated action, which asks first.
+            */}
+            <Button
+              variant="primary"
+              disabled={selectedCount === 0 || !hasChanges}
+              loading={isSubmitting}
+              onClick={handleSubmit}
+            >
+              {isEditing ? 'Update replicated tables' : 'Enable Warehouse'}
+            </Button>
          </CardFooter>
        </Card>
      </PageSectionContent>
@@ -9,6 +9,7 @@ import {
  type WarehouseSetupTarget,
 } from './Warehouse.utils'
 import { WarehouseConnectSection } from './WarehouseConnectSection'
+import { WarehouseDisableCard } from './WarehouseDisableCard'
 import { WarehouseSchemaTablePicker } from './WarehouseSchemaTablePicker'
 import {
  WarehouseEnablingProgress,
@@ -18,9 +19,11 @@ import { AlertError } from '@/components/ui/AlertError'
 import { checkLocalETLNotSetUp } from '@/data/replication/utils'
 import { useWarehouseSetupMutation } from '@/data/warehouse/warehouse-setup-mutation'
 import { useWarehouseSetupStatusQuery } from '@/data/warehouse/warehouse-setup-status-query'
+import { useTrack } from '@/lib/telemetry/track'

 export const WarehouseSetupPanel = () => {
  const { ref: projectRef } = useParams()
+  const track = useTrack()

  const { data, isPending, isFetching, isError, error, refetch } = useWarehouseSetupStatusQuery(
    { projectRef },
@@ -36,8 +39,22 @@ export const WarehouseSetupPanel = () => {

  const handleSetup = (targets: WarehouseSetupTarget[]) => {
    if (!projectRef || targets.length === 0) return
+    const isInitialSetup = data?.setup_status !== 'complete'

-    setupMutation.mutate({ projectRef, body: { targets } })
+    setupMutation.mutate(
+      { projectRef, body: { targets } },
+      {
+        onSuccess: () => {
+          if (isInitialSetup) {
+            track('warehouse_enabled', {
+              source: 'integrations_overview',
+              schemaTargetCount: targets.filter((target) => target.type === 'schema').length,
+              tableTargetCount: targets.filter((target) => target.type === 'table').length,
+            })
+          }
+        },
+      }
+    )
  }

  if (isPending) return <GenericSkeletonLoader />
@@ -123,6 +140,7 @@ export const WarehouseSetupPanel = () => {
        error={setupMutation.error}
      />
      <WarehouseConnectSection />
+      <WarehouseDisableCard />
    </>
  )
 }
@@ -94,7 +94,7 @@ const WarehouseTableStatusList = ({ tables }: WarehouseTableStatusListProps) =>
 }

 export const WarehouseReplicatedTablesSection = ({ tables }: WarehouseTableStatusListProps) => (
-  <PageSection className="pt-0!">
+  <PageSection className="pt-5!">
    <PageSectionMeta>
      <PageSectionSummary>
        <PageSectionTitle>Status</PageSectionTitle>
@@ -119,7 +119,7 @@ export const WarehouseEnablingProgress = ({ status }: WarehouseEnablingProgressP
      : `Backfilling selected tables. ${syncedTableCount} of ${status.tables.length} tables synced.`

  return (
-    <PageSection className="pt-0!">
+    <PageSection className="pt-5!">
      <PageSectionMeta>
        <PageSectionSummary>
          <PageSectionTitle>Status</PageSectionTitle>
@@ -4,6 +4,7 @@ import { toast } from 'sonner'

 import { warehouseKeys } from './keys'
 import { handleError, post } from '@/data/fetchers'
+import { replicationKeys } from '@/data/replication/keys'
 import type { ResponseError, UseCustomMutationOptions } from '@/types'

 export type WarehouseSetupBody = components['schemas']['WarehouseSetupBody']
@@ -42,9 +43,14 @@ export const useWarehouseSetupMutation = ({
  return useMutation<WarehouseSetupData, ResponseError, WarehouseSetupVariables>({
    mutationFn: (vars) => setupWarehouse(vars),
    async onSuccess(data, variables, context) {
-      await queryClient.invalidateQueries({
-        queryKey: warehouseKeys.setupStatus(variables.projectRef),
-      })
+      await Promise.all([
+        queryClient.invalidateQueries({
+          queryKey: warehouseKeys.setupStatus(variables.projectRef),
+        }),
+        queryClient.invalidateQueries({
+          queryKey: replicationKeys.sources(variables.projectRef),
+        }),
+      ])
      await onSuccess?.(data, variables, context)
    },
    async onError(error, variables, context) {