From 2b2e731992503a298615d87561e94ae63900b4f7 Mon Sep 17 00:00:00 2001 From: Edwin Jose Date: Wed, 3 Jun 2026 11:03:45 -0500 Subject: [PATCH 1/3] Add admin-managed global connector toggle Introduce an admin-only workspace toggle for connectors. - Add migration (0007_seed_connectors_manage_global) to create the connectors:manage:global permission and grant it to the admin role. - Persist per-connector enabled state in workspace config under the "connectors" section (default = enabled) and add the section to WorkspaceConfigRepo. - Backend: add helpers (is_connector_enabled, get_connector_enabled_map, user_can_manage_connectors, assert_connector_enabled), list_connectors now returns an "enabled" flag and hides disabled connectors from non-admins, add set_connector_enabled endpoint (PUT /connectors/{connector_type}/enabled), and enforce global-disable checks across relevant connector endpoints and sync flows. - Frontend: include an enabled flag in connector query, add a toggle Switch in connector cards, new useToggleConnectorMutation for optimistic UI updates and API call to flip workspace enablement. - Tests: add unit tests for connector enablement behavior and update existing tests to include session/rbac mocks. This implements a workspace-wide "kill switch" for connectors with admin bypass and UI controls to manage and display the state. --- .../0007_seed_connectors_manage_global.py | 100 +++++++++ .../mutations/useToggleConnectorMutation.ts | 67 ++++++ .../app/api/queries/useGetConnectorsQuery.ts | 4 + .../settings/_components/connector-card.tsx | 61 +++-- .../settings/_components/connector-cards.tsx | 15 ++ src/api/auth.py | 16 +- src/api/connectors.py | 157 ++++++++++++- src/app/routes/internal.py | 6 + src/db/repositories/workspace_config_repo.py | 4 +- src/db/seed.py | 1 + ...st_reconcile_orphans_for_connector_type.py | 25 ++- .../connectors/test_connector_enablement.py | 208 ++++++++++++++++++ .../test_connector_file_type_validation.py | 8 + 13 files changed, 647 insertions(+), 25 deletions(-) create mode 100644 alembic/versions/0007_seed_connectors_manage_global.py create mode 100644 frontend/app/api/mutations/useToggleConnectorMutation.ts create mode 100644 tests/unit/connectors/test_connector_enablement.py diff --git a/alembic/versions/0007_seed_connectors_manage_global.py b/alembic/versions/0007_seed_connectors_manage_global.py new file mode 100644 index 000000000..4a171e040 --- /dev/null +++ b/alembic/versions/0007_seed_connectors_manage_global.py @@ -0,0 +1,100 @@ +"""seed connectors:manage:global permission and grant to admin + +Revision ID: 0007_seed_connectors_manage_global +Revises: 0006_revoke_provider_override_nonadmin +Create Date: 2026-06-03 00:00:00.000000 + +Adds the admin-only ``connectors:manage:global`` permission (workspace-wide +enable/disable of connectors) and grants it to the built-in ``admin`` role. + +Idempotent: skips the permission row if it already exists and skips the +``admin`` join row if it is already present. Mirrors the sync insert pattern +in ``0002_seed_roles_permissions``. + +""" +from collections.abc import Sequence +from typing import Union +import uuid + +import sqlalchemy as sa + +from alembic import op + +revision: str = "0007_seed_connectors_manage_global" +down_revision: Union[str, Sequence[str], None] = "0006_revoke_provider_override_nonadmin" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + +_PERM_NAME = "connectors:manage:global" +_PERM_RESOURCE = "connectors" +_PERM_ACTION = "manage:global" +_PERM_DESCRIPTION = "Enable/disable connectors workspace-wide" +_ROLE_NAME = "admin" + + +def upgrade() -> None: + bind = op.get_bind() + + perm_id = bind.execute( + sa.text("SELECT id FROM permissions WHERE name = :name"), + {"name": _PERM_NAME}, + ).scalar() + + if perm_id is None: + perm_id = str(uuid.uuid4()) + perms_table = sa.table( + "permissions", + sa.column("id", sa.String), + sa.column("name", sa.String), + sa.column("resource", sa.String), + sa.column("action", sa.String), + sa.column("description", sa.String), + ) + op.bulk_insert( + perms_table, + [ + { + "id": perm_id, + "name": _PERM_NAME, + "resource": _PERM_RESOURCE, + "action": _PERM_ACTION, + "description": _PERM_DESCRIPTION, + } + ], + ) + + role_id = bind.execute( + sa.text("SELECT id FROM roles WHERE name = :name"), + {"name": _ROLE_NAME}, + ).scalar() + if role_id is None: + return + + already_granted = bind.execute( + sa.text( + "SELECT 1 FROM role_permissions " + "WHERE role_id = :rid AND permission_id = :pid" + ), + {"rid": role_id, "pid": perm_id}, + ).scalar() + if already_granted is None: + rp_table = sa.table( + "role_permissions", + sa.column("role_id", sa.String), + sa.column("permission_id", sa.String), + ) + op.bulk_insert(rp_table, [{"role_id": role_id, "permission_id": perm_id}]) + + +def downgrade() -> None: + op.execute( + sa.text( + "DELETE FROM role_permissions WHERE permission_id IN " + "(SELECT id FROM permissions WHERE name = :name)" + ).bindparams(sa.bindparam("name", _PERM_NAME)) + ) + op.execute( + sa.text("DELETE FROM permissions WHERE name = :name").bindparams( + sa.bindparam("name", _PERM_NAME) + ) + ) diff --git a/frontend/app/api/mutations/useToggleConnectorMutation.ts b/frontend/app/api/mutations/useToggleConnectorMutation.ts new file mode 100644 index 000000000..4f0250104 --- /dev/null +++ b/frontend/app/api/mutations/useToggleConnectorMutation.ts @@ -0,0 +1,67 @@ +import { useMutation, useQueryClient } from "@tanstack/react-query"; +import { toast } from "sonner"; +import type { Connector } from "../queries/useGetConnectorsQuery"; + +interface ToggleConnectorVariables { + connector: Connector; + enabled: boolean; +} + +export const useToggleConnectorMutation = () => { + const queryClient = useQueryClient(); + + return useMutation({ + mutationFn: async ({ connector, enabled }: ToggleConnectorVariables) => { + const response = await fetch( + `/api/connectors/${connector.type}/enabled`, + { + method: "PUT", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ enabled }), + }, + ); + + if (!response.ok) { + const result = await response.json().catch(() => ({})); + throw new Error( + result?.detail?.error || + result.error || + `Failed to update ${connector.name}`, + ); + } + return response.json(); + }, + onMutate: async ({ connector, enabled }) => { + await queryClient.cancelQueries({ queryKey: ["connectors"] }); + + const previousConnectors = queryClient.getQueryData([ + "connectors", + ]); + + if (previousConnectors) { + queryClient.setQueryData( + ["connectors"], + previousConnectors.map((c) => + c.type === connector.type ? { ...c, enabled } : c, + ), + ); + } + + return { previousConnectors }; + }, + onError: (err, { connector }, context) => { + if (context?.previousConnectors) { + queryClient.setQueryData(["connectors"], context.previousConnectors); + } + toast.error(`Failed to update ${connector.name}: ${err.message}`); + }, + onSuccess: (_, { connector, enabled }) => { + toast.success( + `${connector.name} ${enabled ? "enabled" : "disabled"} for the workspace`, + ); + }, + onSettled: () => { + queryClient.invalidateQueries({ queryKey: ["connectors"] }); + }, + }); +}; diff --git a/frontend/app/api/queries/useGetConnectorsQuery.ts b/frontend/app/api/queries/useGetConnectorsQuery.ts index 84a197249..b382e6378 100644 --- a/frontend/app/api/queries/useGetConnectorsQuery.ts +++ b/frontend/app/api/queries/useGetConnectorsQuery.ts @@ -32,6 +32,8 @@ export interface Connector { access_token?: string; selectedFiles?: GoogleDriveFile[] | OneDriveFile[]; available?: boolean; + /** Admin-managed workspace toggle. Absent is treated as enabled. */ + enabled?: boolean; } interface Connection { @@ -89,6 +91,7 @@ export const useGetConnectorsQuery = ( clientId: activeConnection.client_id, baseUrl: activeConnection.base_url, available: connectorData.available, + enabled: connectorData.enabled, } as Connector; } } @@ -102,6 +105,7 @@ export const useGetConnectorsQuery = ( type, connectionId, available: connectorData.available, + enabled: connectorData.enabled, } as Connector; }), ); diff --git a/frontend/app/settings/_components/connector-card.tsx b/frontend/app/settings/_components/connector-card.tsx index 314284ffd..f5698cc92 100644 --- a/frontend/app/settings/_components/connector-card.tsx +++ b/frontend/app/settings/_components/connector-card.tsx @@ -10,6 +10,7 @@ import { CardHeader, CardTitle, } from "@/components/ui/card"; +import { Switch } from "@/components/ui/switch"; import { useIsCloudBrand } from "@/contexts/brand-context"; import { usePermissions } from "@/hooks/use-permissions"; import { cn } from "@/lib/utils"; @@ -23,6 +24,8 @@ export interface Connector { available?: boolean; status?: string; connectionId?: string; + /** Admin-managed workspace toggle. Absent is treated as enabled. */ + enabled?: boolean; } interface ConnectorCardProps { @@ -34,6 +37,9 @@ interface ConnectorCardProps { onNavigateToKnowledge: (connector: Connector) => void; /** Optional: open a connector-specific settings/edit dialog */ onConfigure?: (connector: Connector) => void; + /** Admin-only: toggle the connector on/off workspace-wide */ + onToggleEnabled?: (connector: Connector, enabled: boolean) => void; + isTogglingEnabled?: boolean; } export default function ConnectorCard({ @@ -44,6 +50,8 @@ export default function ConnectorCard({ onDisconnect, onNavigateToKnowledge, onConfigure, + onToggleEnabled, + isTogglingEnabled, }: ConnectorCardProps) { const isCloudBrand = useIsCloudBrand(); const { can, canAny } = usePermissions(); @@ -53,6 +61,8 @@ export default function ConnectorCard({ "connectors:delete:any", ]); const canUpload = can("knowledge:upload"); + const canManageGlobal = can("connectors:manage:global"); + const isEnabled = connector.enabled !== false; const isConnected = connector.status === "connected" && connector.connectionId; @@ -72,19 +82,36 @@ export default function ConnectorCard({ {connector.icon} - {isConnected ? ( -
- - Active -
- ) : null} +
+ {isConnected ? ( +
+ + Active +
+ ) : null} + {canManageGlobal && onToggleEnabled ? ( + + onToggleEnabled(connector, checked) + } + aria-label={`${isEnabled ? "Disable" : "Enable"} ${connector.name} for the workspace`} + title={ + isEnabled + ? "Enabled for the workspace — toggle off to hide from all users" + : "Disabled for the workspace — toggle on to make available to all users" + } + /> + ) : null} +
- {isConnected || connector?.available - ? `${connector.name} is configured.` - : "Not configured."} + {canManageGlobal && !isEnabled + ? `${connector.name} is disabled for the workspace.` + : isConnected || connector?.available + ? `${connector.name} is configured.` + : "Not configured."}
diff --git a/frontend/app/settings/_components/connector-cards.tsx b/frontend/app/settings/_components/connector-cards.tsx index 5f0880e53..d6db86010 100644 --- a/frontend/app/settings/_components/connector-cards.tsx +++ b/frontend/app/settings/_components/connector-cards.tsx @@ -4,6 +4,7 @@ import { useRouter } from "next/navigation"; import { useCallback, useState } from "react"; import { useConnectConnectorMutation } from "@/app/api/mutations/useConnectConnectorMutation"; import { useDisconnectConnectorMutation } from "@/app/api/mutations/useDisconnectConnectorMutation"; +import { useToggleConnectorMutation } from "@/app/api/mutations/useToggleConnectorMutation"; import { type Connector as QueryConnector, useGetConnectorsQuery, @@ -34,6 +35,7 @@ export default function ConnectorCards() { const connectMutation = useConnectConnectorMutation(); const disconnectMutation = useDisconnectConnectorMutation(); + const toggleEnabledMutation = useToggleConnectorMutation(); const getConnectorIcon = useCallback((iconName: string) => { const iconMap: { [key: string]: React.ReactElement } = { @@ -74,6 +76,13 @@ export default function ConnectorCards() { disconnectMutation.mutate(connector as unknown as QueryConnector); }; + const handleToggleEnabled = (connector: Connector, enabled: boolean) => { + toggleEnabledMutation.mutate({ + connector: connector as unknown as QueryConnector, + enabled, + }); + }; + const navigateToKnowledgePage = (connector: Connector) => { const provider = connector.type.replace(/-/g, "_"); router.push(`/upload/${provider}`); @@ -120,6 +129,12 @@ export default function ConnectorCards() { onDisconnect={handleDisconnect} onNavigateToKnowledge={navigateToKnowledgePage} onConfigure={getConfigureHandler(connector)} + onToggleEnabled={handleToggleEnabled} + isTogglingEnabled={ + toggleEnabledMutation.isPending && + toggleEnabledMutation.variables?.connector?.type === + connector.type + } /> )) )} diff --git a/src/api/auth.py b/src/api/auth.py index 4a9df0e94..163ca6713 100644 --- a/src/api/auth.py +++ b/src/api/auth.py @@ -8,10 +8,14 @@ logger = get_logger(__name__) +from sqlalchemy.ext.asyncio import AsyncSession + from dependencies import ( get_auth_service, - get_optional_user, get_current_user, + get_db_session, + get_optional_user, + get_rbac_service, ) from pydantic import BaseModel from session_manager import User @@ -41,8 +45,18 @@ async def auth_init( request: Request, auth_service=Depends(get_auth_service), user: Optional[User] = Depends(get_optional_user), + session: AsyncSession = Depends(get_db_session), + rbac=Depends(get_rbac_service), ): """Initialize OAuth flow for authentication or data source connection""" + # Block connecting a globally-disabled connector for non-admins. Done before + # the try/except so the 403 is not swallowed into a generic 500. + # App-login OAuth ("app_auth") is never a data-source connect, so skip it. + if body.purpose != "app_auth": + from api.connectors import assert_connector_enabled + + await assert_connector_enabled(body.connector_type, user, rbac, session) + try: connection_name = body.name or f"{body.connector_type}_{body.purpose}" user_id = user.user_id if user else None diff --git a/src/api/connectors.py b/src/api/connectors.py index 665962d87..4cc68bd6f 100644 --- a/src/api/connectors.py +++ b/src/api/connectors.py @@ -1,23 +1,72 @@ from typing import Any -from fastapi import Depends, Request +from fastapi import Depends, HTTPException, Request from fastapi.responses import JSONResponse, PlainTextResponse from pydantic import BaseModel +from sqlalchemy.ext.asyncio import AsyncSession from config.settings import get_index_name from connectors.sharepoint.utils import is_valid_sharepoint_url +from db.repositories.workspace_config_repo import WorkspaceConfigRepo from dependencies import ( get_connector_service, get_current_user, + get_db_session, + get_rbac_service, get_session_manager, require_permission, ) +from services.rbac_service import is_rbac_enforced from session_manager import User from utils.logging_config import get_logger from utils.telemetry import Category, MessageId, TelemetryClient logger = get_logger(__name__) +# Workspace-config section holding the admin-managed per-connector enable/disable +# state. Value shape: ``{"": bool}``. A connector absent from the +# map is treated as enabled (default-enabled), preserving pre-toggle behavior. +CONNECTORS_CONFIG_SECTION = "connectors" +MANAGE_GLOBAL_PERMISSION = "connectors:manage:global" + + +async def get_connector_enabled_map(session: AsyncSession) -> dict[str, bool]: + """Return the admin-set ``{connector_type: enabled}`` map from workspace config.""" + repo = WorkspaceConfigRepo(session) + value = await repo.get_section(CONNECTORS_CONFIG_SECTION) or {} + return {str(k): bool(v) for k, v in value.items()} + + +def is_connector_enabled(enabled_map: dict[str, bool], connector_type: str) -> bool: + """A connector is enabled unless an admin has explicitly disabled it.""" + return enabled_map.get(connector_type, True) + + +async def user_can_manage_connectors(user: User | None, rbac) -> bool: + """True when the user holds the admin-only ``connectors:manage:global`` perm.""" + uid = getattr(user, "db_user_id", None) or (user.user_id if user else None) + if not uid: + return False + return await rbac.has_permission(uid, MANAGE_GLOBAL_PERMISSION) + + +async def assert_connector_enabled( + connector_type: str, user: User | None, rbac, session: AsyncSession +) -> None: + """Raise 403 if ``connector_type`` is globally disabled and the user is not an + admin who can manage connectors. No-op when RBAC enforcement is off.""" + if not is_rbac_enforced(): + return + enabled_map = await get_connector_enabled_map(session) + if is_connector_enabled(enabled_map, connector_type): + return + if await user_can_manage_connectors(user, rbac): + return + raise HTTPException( + status_code=403, + detail={"error": "connector_disabled", "connector_type": connector_type}, + ) + def _connector_sync_should_replace(connector_type: str) -> bool: """Return True for connector types where sync should replace existing indexed files.""" @@ -358,8 +407,11 @@ async def connector_check_duplicates( connector_service=Depends(get_connector_service), session_manager=Depends(get_session_manager), user: User = Depends(require_permission("connectors:use")), + session: AsyncSession = Depends(get_db_session), + rbac=Depends(get_rbac_service), ): """Check if any of the selected files or folders contain files that already exist in the index""" + await assert_connector_enabled(connector_type, user, rbac, session) selected_files_raw = body.selected_files if not selected_files_raw: return JSONResponse({"duplicate_names": []}) @@ -500,16 +552,72 @@ async def connector_check_duplicates( async def list_connectors( connector_service=Depends(get_connector_service), user: User = Depends(get_current_user), + session: AsyncSession = Depends(get_db_session), + rbac=Depends(get_rbac_service), ): - """List available connector types with metadata""" + """List available connector types with metadata. + + Each connector carries an ``enabled`` flag reflecting the admin-managed + workspace toggle. Non-admin users never see globally-disabled connectors; + admins (``connectors:manage:global``) always see every connector so the + toggle can render its current state. + """ try: connector_types = connector_service.connection_manager.get_available_connector_types( user_id=user.user_id ) - return JSONResponse({"connectors": connector_types}) + enabled_map = await get_connector_enabled_map(session) + is_admin = await user_can_manage_connectors(user, rbac) + + result: dict[str, dict[str, Any]] = {} + for ctype, meta in connector_types.items(): + enabled = is_connector_enabled(enabled_map, ctype) + if not enabled and not is_admin: + continue + result[ctype] = {**meta, "enabled": enabled} + + return JSONResponse({"connectors": result}) except Exception as e: logger.error("[CONNECTOR] Error listing connectors", error=str(e)) - return JSONResponse({"connectors": []}) + return JSONResponse({"connectors": {}}) + + +class SetConnectorEnabledBody(BaseModel): + enabled: bool + + +async def set_connector_enabled( + connector_type: str, + body: SetConnectorEnabledBody, + connector_service=Depends(get_connector_service), + session: AsyncSession = Depends(get_db_session), + user: User = Depends(require_permission(MANAGE_GLOBAL_PERMISSION)), +): + """Admin-only: enable or disable a connector workspace-wide. + + Persists the state in the ``connectors`` workspace-config section. + """ + known_types = set(connector_service.connection_manager.get_available_connector_types().keys()) + if connector_type not in known_types: + raise HTTPException( + status_code=404, + detail={"error": "unknown_connector_type", "connector_type": connector_type}, + ) + + repo = WorkspaceConfigRepo(session) + value = dict(await repo.get_section(CONNECTORS_CONFIG_SECTION) or {}) + value[connector_type] = body.enabled + actor_user_id = user.db_user_id or user.user_id + await repo.upsert(CONNECTORS_CONFIG_SECTION, value, actor_user_id=actor_user_id) + await session.commit() + + logger.info( + "[CONNECTOR] Connector global-enabled state updated", + connector_type=connector_type, + enabled=body.enabled, + actor=actor_user_id, + ) + return JSONResponse({"connector_type": connector_type, "enabled": body.enabled}) async def connector_sync( @@ -518,8 +626,11 @@ async def connector_sync( connector_service=Depends(get_connector_service), session_manager=Depends(get_session_manager), user: User = Depends(require_permission("connectors:use")), + session: AsyncSession = Depends(get_db_session), + rbac=Depends(get_rbac_service), ): """Sync files from all active connections of a connector type""" + await assert_connector_enabled(connector_type, user, rbac, session) max_files = body.max_files selected_files_raw = body.selected_files selected_files = None @@ -1097,6 +1208,8 @@ async def sync_all_connectors( connector_service=Depends(get_connector_service), session_manager=Depends(get_session_manager), user: User = Depends(require_permission("connectors:use")), + session: AsyncSession = Depends(get_db_session), + rbac=Depends(get_rbac_service), ): """ Sync files from all active cloud connector connections. @@ -1110,6 +1223,10 @@ async def sync_all_connectors( # Cloud connector types to sync cloud_connector_types = ["google_drive", "onedrive", "sharepoint", "ibm_cos", "aws_s3"] + # Globally-disabled connectors are skipped for non-admins. + enabled_map = await get_connector_enabled_map(session) + is_admin = await user_can_manage_connectors(user, rbac) + all_task_ids = [] synced_connectors = [] skipped_connectors = [] @@ -1117,6 +1234,17 @@ async def sync_all_connectors( errors = [] for connector_type in cloud_connector_types: + if ( + is_rbac_enforced() + and not is_admin + and not is_connector_enabled(enabled_map, connector_type) + ): + logger.debug( + "Connector globally disabled, skipping in sync-all", + connector_type=connector_type, + ) + skipped_connectors.append(connector_type) + continue try: # First, get existing file IDs/filenames from OpenSearch for this connector type ( @@ -1346,11 +1474,14 @@ async def connector_sync_preview( connector_service=Depends(get_connector_service), session_manager=Depends(get_session_manager), user: User = Depends(require_permission("connectors:use")), + session: AsyncSession = Depends(get_db_session), + rbac=Depends(get_rbac_service), ): """Preview the impact of syncing a connector type without performing any deletion or ingest. Returns the list of orphan files (present in OpenSearch but no longer at the source) by filename, plus the total synced count. """ + await assert_connector_enabled(connector_type, user, rbac, session) try: orphans, synced_count = await _preview_orphans_for_connector_type( connector_type=connector_type, @@ -1377,6 +1508,8 @@ async def connectors_sync_all_preview( connector_service=Depends(get_connector_service), session_manager=Depends(get_session_manager), user: User = Depends(require_permission("connectors:use")), + session: AsyncSession = Depends(get_db_session), + rbac=Depends(get_rbac_service), ): """Preview the impact of sync-all-connectors across every cloud connector type. Returns orphan filenames grouped by connector_type plus a per-type @@ -1387,7 +1520,17 @@ async def connectors_sync_all_preview( synced_count_by_type: dict[str, int] = {} orphans_available_by_type: dict[str, bool] = {} + # Globally-disabled connectors are skipped for non-admins. + enabled_map = await get_connector_enabled_map(session) + is_admin = await user_can_manage_connectors(user, rbac) + for connector_type in CLOUD_CONNECTOR_TYPES: + if ( + is_rbac_enforced() + and not is_admin + and not is_connector_enabled(enabled_map, connector_type) + ): + continue try: orphans, synced_count = await _preview_orphans_for_connector_type( connector_type=connector_type, @@ -1431,10 +1574,13 @@ async def connector_token( request: Request, connector_service=Depends(get_connector_service), user: User = Depends(get_current_user), + session: AsyncSession = Depends(get_db_session), + rbac=Depends(get_rbac_service), ): """Get access token for connector API calls (e.g., Pickers).""" url_connector_type = connector_type + await assert_connector_enabled(connector_type, user, rbac, session) try: # 1) Load the connection and verify ownership connection = await connector_service.connection_manager.get_connection(connection_id) @@ -1550,6 +1696,8 @@ async def browse_connection_files( connector_service=Depends(get_connector_service), session_manager=Depends(get_session_manager), user: User = Depends(get_current_user), + session: AsyncSession = Depends(get_db_session), + rbac=Depends(get_rbac_service), bucket: str | None = None, search: str | None = None, page_token: str | None = None, @@ -1561,6 +1709,7 @@ async def browse_connection_files( Lists files from the remote source (e.g., S3 bucket) and marks each as ingested or not by cross-referencing with OpenSearch. """ + await assert_connector_enabled(connector_type, user, rbac, session) try: connector = await connector_service.get_connector(connection_id) if not connector: diff --git a/src/app/routes/internal.py b/src/app/routes/internal.py index 1605f88a1..f1dbdff98 100644 --- a/src/app/routes/internal.py +++ b/src/app/routes/internal.py @@ -302,6 +302,12 @@ def register_internal_routes(app: FastAPI): methods=["GET"], tags=["internal"], ) + app.add_api_route( + "/connectors/{connector_type}/enabled", + connectors.set_connector_enabled, + methods=["PUT"], + tags=["internal"], + ) app.add_api_route( "/connectors/{connector_type}/token", connectors.connector_token, diff --git a/src/db/repositories/workspace_config_repo.py b/src/db/repositories/workspace_config_repo.py index 47ddf734d..124cb7f3e 100644 --- a/src/db/repositories/workspace_config_repo.py +++ b/src/db/repositories/workspace_config_repo.py @@ -14,7 +14,9 @@ from db.models import WorkspaceConfig # Section names recognized by the migration / service. -SECTIONS = ("providers", "knowledge", "agent", "onboarding", "meta") +# 'connectors' holds the admin-managed per-connector enable/disable map and is +# read/written directly by the connectors API (not mapped into OpenRAGConfig). +SECTIONS = ("providers", "knowledge", "agent", "onboarding", "meta", "connectors") class WorkspaceConfigRepo: diff --git a/src/db/seed.py b/src/db/seed.py index 5633a6069..55640099e 100644 --- a/src/db/seed.py +++ b/src/db/seed.py @@ -46,6 +46,7 @@ ("connectors", "delete:own", "Delete own connectors"), ("connectors", "delete:any", "Delete any connector"), ("connectors", "use", "Use connector OAuth and browse"), + ("connectors", "manage:global", "Enable/disable connectors workspace-wide"), # Knowledge ("knowledge", "upload", "Upload documents"), ("knowledge", "delete:own", "Delete own documents"), diff --git a/tests/unit/api/test_reconcile_orphans_for_connector_type.py b/tests/unit/api/test_reconcile_orphans_for_connector_type.py index 9e339ca37..955ad7b40 100644 --- a/tests/unit/api/test_reconcile_orphans_for_connector_type.py +++ b/tests/unit/api/test_reconcile_orphans_for_connector_type.py @@ -27,6 +27,19 @@ def _make_connection(connection_id: str, is_active: bool = True): return SimpleNamespace(connection_id=connection_id, is_active=is_active) +def _enabled_session(): + """DB session whose 'connectors' config row is absent → all enabled.""" + session = MagicMock() + session.get = AsyncMock(return_value=None) + return session + + +def _rbac(is_admin: bool = False): + rbac = MagicMock() + rbac.has_permission = AsyncMock(return_value=is_admin) + return rbac + + def _make_connector(remote_file_ids, *, authenticated=True, raise_on_list=False): connector = MagicMock() connector.is_authenticated = authenticated @@ -461,7 +474,9 @@ async def test_connector_sync_filters_orphan_ids_before_resync(monkeypatch): connectors_api.ConnectorSyncBody(), connector_service=service, session_manager=MagicMock(), - user=SimpleNamespace(user_id="alice", jwt_token="token"), + user=SimpleNamespace(user_id="alice", jwt_token="token", db_user_id="alice"), + session=_enabled_session(), + rbac=_rbac(), ) assert response.status_code == 201 @@ -501,7 +516,9 @@ async def test_connector_sync_returns_no_files_when_all_ids_are_orphans(monkeypa connectors_api.ConnectorSyncBody(), connector_service=service, session_manager=MagicMock(), - user=SimpleNamespace(user_id="alice", jwt_token="token"), + user=SimpleNamespace(user_id="alice", jwt_token="token", db_user_id="alice"), + session=_enabled_session(), + rbac=_rbac(), ) assert response.status_code == 200 @@ -540,7 +557,9 @@ async def fake_synced_ids(connector_type, *args, **kwargs): response = await connectors_api.sync_all_connectors( connector_service=service, session_manager=MagicMock(), - user=SimpleNamespace(user_id="alice", jwt_token="token"), + user=SimpleNamespace(user_id="alice", jwt_token="token", db_user_id="alice"), + session=_enabled_session(), + rbac=_rbac(), ) body = _json(response) diff --git a/tests/unit/connectors/test_connector_enablement.py b/tests/unit/connectors/test_connector_enablement.py new file mode 100644 index 000000000..80399bf76 --- /dev/null +++ b/tests/unit/connectors/test_connector_enablement.py @@ -0,0 +1,208 @@ +"""Unit tests for the admin-managed connector enable/disable toggle. + +Covers the helpers and endpoints added for ``connectors:manage:global``: +- ``is_connector_enabled`` default-enabled semantics +- ``assert_connector_enabled`` gating (kill switch, admin bypass, 403) +- ``set_connector_enabled`` persistence + authz validation +- ``list_connectors`` enabled-flag merge and non-admin hiding +""" + +import json +import sys +from pathlib import Path +from unittest.mock import AsyncMock, MagicMock + +import pytest +from fastapi import HTTPException + +ROOT = Path(__file__).resolve().parent.parent.parent.parent +SRC = ROOT / "src" +if str(SRC) not in sys.path: + sys.path.insert(0, str(SRC)) + + +def _session_returning(value): + """A DB-session mock whose ``get`` resolves to a WorkspaceConfig-like row.""" + session = MagicMock() + if value is None: + session.get = AsyncMock(return_value=None) + else: + row = MagicMock() + row.value = value + session.get = AsyncMock(return_value=row) + return session + + +def _connector_service(types=("google_drive", "sharepoint", "onedrive")): + svc = MagicMock() + cm = MagicMock() + cm.get_available_connector_types = MagicMock( + return_value={ + t: {"name": t, "description": "", "icon": t} for t in types + } + ) + svc.connection_manager = cm + return svc + + +def test_is_connector_enabled_defaults_to_enabled(): + from api.connectors import is_connector_enabled + + assert is_connector_enabled({}, "google_drive") is True + assert is_connector_enabled({"google_drive": False}, "google_drive") is False + assert is_connector_enabled({"sharepoint": False}, "google_drive") is True + + +@pytest.mark.asyncio +async def test_assert_connector_enabled_bypasses_when_rbac_off(monkeypatch): + import api.connectors as connectors + + monkeypatch.setattr(connectors, "is_rbac_enforced", lambda: False) + # session.get would explode if reached — it must not be reached. + session = MagicMock() + session.get = AsyncMock(side_effect=AssertionError("should not query DB")) + await connectors.assert_connector_enabled("google_drive", MagicMock(), MagicMock(), session) + + +@pytest.mark.asyncio +async def test_assert_connector_enabled_blocks_non_admin(monkeypatch): + import api.connectors as connectors + + monkeypatch.setattr(connectors, "is_rbac_enforced", lambda: True) + session = _session_returning({"google_drive": False}) + rbac = MagicMock() + rbac.has_permission = AsyncMock(return_value=False) + user = MagicMock() + user.db_user_id = "u1" + + with pytest.raises(HTTPException) as exc: + await connectors.assert_connector_enabled("google_drive", user, rbac, session) + assert exc.value.status_code == 403 + assert exc.value.detail["error"] == "connector_disabled" + + +@pytest.mark.asyncio +async def test_assert_connector_enabled_allows_admin(monkeypatch): + import api.connectors as connectors + + monkeypatch.setattr(connectors, "is_rbac_enforced", lambda: True) + session = _session_returning({"google_drive": False}) + rbac = MagicMock() + rbac.has_permission = AsyncMock(return_value=True) + user = MagicMock() + user.db_user_id = "admin" + + # Disabled connector, but admin → no raise. + await connectors.assert_connector_enabled("google_drive", user, rbac, session) + + +@pytest.mark.asyncio +async def test_set_connector_enabled_persists(monkeypatch): + import api.connectors as connectors + + svc = _connector_service() + session = _session_returning(None) # no existing row + session.commit = AsyncMock() + + captured = {} + + async def fake_upsert(self, section, value, actor_user_id=None): + captured["section"] = section + captured["value"] = value + captured["actor"] = actor_user_id + + monkeypatch.setattr(connectors.WorkspaceConfigRepo, "upsert", fake_upsert) + + user = MagicMock() + user.db_user_id = "admin" + body = connectors.SetConnectorEnabledBody(enabled=False) + + resp = await connectors.set_connector_enabled( + connector_type="sharepoint", + body=body, + connector_service=svc, + session=session, + user=user, + ) + + data = json.loads(resp.body.decode()) + assert data == {"connector_type": "sharepoint", "enabled": False} + assert captured["section"] == "connectors" + assert captured["value"] == {"sharepoint": False} + assert captured["actor"] == "admin" + session.commit.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_set_connector_enabled_rejects_unknown_type(): + import api.connectors as connectors + + svc = _connector_service() + session = _session_returning(None) + user = MagicMock() + body = connectors.SetConnectorEnabledBody(enabled=True) + + with pytest.raises(HTTPException) as exc: + await connectors.set_connector_enabled( + connector_type="not_a_connector", + body=body, + connector_service=svc, + session=session, + user=user, + ) + assert exc.value.status_code == 404 + + +@pytest.mark.asyncio +async def test_list_connectors_hides_disabled_for_non_admin(monkeypatch): + import api.connectors as connectors + + svc = _connector_service() + # get_available_connector_types is called with user_id=... in the handler. + svc.connection_manager.get_available_connector_types = MagicMock( + return_value={ + "google_drive": {"name": "g", "description": "", "icon": "g", "available": True}, + "sharepoint": {"name": "s", "description": "", "icon": "s", "available": True}, + } + ) + session = _session_returning({"sharepoint": False}) + rbac = MagicMock() + rbac.has_permission = AsyncMock(return_value=False) # non-admin + user = MagicMock() + user.user_id = "u1" + user.db_user_id = "u1" + + resp = await connectors.list_connectors( + connector_service=svc, user=user, session=session, rbac=rbac + ) + data = json.loads(resp.body.decode())["connectors"] + assert "google_drive" in data + assert data["google_drive"]["enabled"] is True + assert "sharepoint" not in data # hidden for non-admin + + +@pytest.mark.asyncio +async def test_list_connectors_shows_disabled_with_flag_for_admin(monkeypatch): + import api.connectors as connectors + + svc = _connector_service() + svc.connection_manager.get_available_connector_types = MagicMock( + return_value={ + "google_drive": {"name": "g", "description": "", "icon": "g", "available": True}, + "sharepoint": {"name": "s", "description": "", "icon": "s", "available": True}, + } + ) + session = _session_returning({"sharepoint": False}) + rbac = MagicMock() + rbac.has_permission = AsyncMock(return_value=True) # admin + user = MagicMock() + user.user_id = "admin" + user.db_user_id = "admin" + + resp = await connectors.list_connectors( + connector_service=svc, user=user, session=session, rbac=rbac + ) + data = json.loads(resp.body.decode())["connectors"] + assert data["google_drive"]["enabled"] is True + assert "sharepoint" in data # visible to admin + assert data["sharepoint"]["enabled"] is False diff --git a/tests/unit/connectors/test_connector_file_type_validation.py b/tests/unit/connectors/test_connector_file_type_validation.py index 12e48a154..65384232b 100644 --- a/tests/unit/connectors/test_connector_file_type_validation.py +++ b/tests/unit/connectors/test_connector_file_type_validation.py @@ -133,6 +133,12 @@ async def test_connector_check_duplicates(): user.user_id = "user-id" user.jwt_token = "jwt-token" + # No 'connectors' workspace-config row → connector treated as enabled. + db_session = MagicMock() + db_session.get = AsyncMock(return_value=None) + rbac = MagicMock() + rbac.has_permission = AsyncMock(return_value=False) + body = ConnectorCheckDuplicatesBody( connection_id="conn-id", selected_files=[{"id": "folder-1", "name": "Folder 1", "isFolder": True}], @@ -144,6 +150,8 @@ async def test_connector_check_duplicates(): connector_service=connector_service, session_manager=session_manager, user=user, + session=db_session, + rbac=rbac, ) assert isinstance(response, JSONResponse) From 743932d953b13739f82ef29d02a2efe2c8f83806 Mon Sep 17 00:00:00 2001 From: "autofix-ci[bot]" <114827586+autofix-ci[bot]@users.noreply.github.com> Date: Wed, 3 Jun 2026 16:43:45 +0000 Subject: [PATCH 2/3] style: ruff autofix (auto) --- .../0007_seed_connectors_manage_global.py | 14 ++++---- src/api/auth.py | 35 +++++++------------ src/db/repositories/workspace_config_repo.py | 4 +-- .../connectors/test_connector_enablement.py | 4 +-- 4 files changed, 22 insertions(+), 35 deletions(-) diff --git a/alembic/versions/0007_seed_connectors_manage_global.py b/alembic/versions/0007_seed_connectors_manage_global.py index 4a171e040..0a3b4cb90 100644 --- a/alembic/versions/0007_seed_connectors_manage_global.py +++ b/alembic/versions/0007_seed_connectors_manage_global.py @@ -12,18 +12,19 @@ in ``0002_seed_roles_permissions``. """ + +import uuid from collections.abc import Sequence from typing import Union -import uuid import sqlalchemy as sa from alembic import op revision: str = "0007_seed_connectors_manage_global" -down_revision: Union[str, Sequence[str], None] = "0006_revoke_provider_override_nonadmin" -branch_labels: Union[str, Sequence[str], None] = None -depends_on: Union[str, Sequence[str], None] = None +down_revision: str | Sequence[str] | None = "0006_revoke_provider_override_nonadmin" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None _PERM_NAME = "connectors:manage:global" _PERM_RESOURCE = "connectors" @@ -71,10 +72,7 @@ def upgrade() -> None: return already_granted = bind.execute( - sa.text( - "SELECT 1 FROM role_permissions " - "WHERE role_id = :rid AND permission_id = :pid" - ), + sa.text("SELECT 1 FROM role_permissions WHERE role_id = :rid AND permission_id = :pid"), {"rid": role_id, "pid": perm_id}, ).scalar() if already_granted is None: diff --git a/src/api/auth.py b/src/api/auth.py index 163ca6713..fa708bcb6 100644 --- a/src/api/auth.py +++ b/src/api/auth.py @@ -2,12 +2,14 @@ from fastapi import Depends, HTTPException, Request from fastapi.responses import JSONResponse -from utils.telemetry import TelemetryClient, Category, MessageId -from utils.version_utils import OPENRAG_VERSION + from utils.logging_config import get_logger +from utils.telemetry import Category, MessageId, TelemetryClient +from utils.version_utils import OPENRAG_VERSION logger = get_logger(__name__) +from pydantic import BaseModel from sqlalchemy.ext.asyncio import AsyncSession from dependencies import ( @@ -17,17 +19,14 @@ get_optional_user, get_rbac_service, ) -from pydantic import BaseModel from session_manager import User class AuthInitBody(BaseModel): connector_type: str purpose: str = "data_source" - name: Optional[str] = None - redirect_uri: Optional[str] = None - - + name: str | None = None + redirect_uri: str | None = None class AuthCallbackBody(BaseModel): @@ -44,7 +43,7 @@ async def auth_init( body: AuthInitBody, request: Request, auth_service=Depends(get_auth_service), - user: Optional[User] = Depends(get_optional_user), + user: User | None = Depends(get_optional_user), session: AsyncSession = Depends(get_db_session), rbac=Depends(get_rbac_service), ): @@ -68,9 +67,7 @@ async def auth_init( except Exception as e: logger.exception("[AUTH] OAuth init failed") - return JSONResponse( - {"error": f"Failed to initialize OAuth: {str(e)}"}, status_code=500 - ) + return JSONResponse({"error": f"Failed to initialize OAuth: {str(e)}"}, status_code=500) async def auth_callback( @@ -89,14 +86,12 @@ async def auth_callback( # If this is app auth, set JWT cookie if result.get("purpose") == "app_auth" and result.get("jwt_token"): await TelemetryClient.send_event(Category.AUTHENTICATION, MessageId.ORB_AUTH_SUCCESS) - response = JSONResponse( - {k: v for k, v in result.items() if k != "jwt_token"} - ) + response = JSONResponse({k: v for k, v in result.items() if k != "jwt_token"}) # Store only the raw JWT (without "Bearer " prefix) in the cookie. # The prefix is added by the OpenSearch client when building the Authorization header. jwt_value = result["jwt_token"] if jwt_value.startswith("Bearer "): - jwt_value = jwt_value[len("Bearer "):] + jwt_value = jwt_value[len("Bearer ") :] response.set_cookie( key="auth_token", value=jwt_value, @@ -118,7 +113,7 @@ async def auth_callback( async def auth_me( request: Request, auth_service=Depends(get_auth_service), - user: Optional[User] = Depends(get_optional_user), + user: User | None = Depends(get_optional_user), ): """Get current user information""" result = await auth_service.get_user_info(request) @@ -152,14 +147,10 @@ async def auth_logout( response.delete_cookie(key="ibm-auth-basic", httponly=True, samesite="lax") return response - response = JSONResponse( - {"status": "logged_out", "message": "Successfully logged out"} - ) + response = JSONResponse({"status": "logged_out", "message": "Successfully logged out"}) # Clear the auth cookie - response.delete_cookie( - key="auth_token", httponly=True, secure=False, samesite="lax" - ) + response.delete_cookie(key="auth_token", httponly=True, secure=False, samesite="lax") return response diff --git a/src/db/repositories/workspace_config_repo.py b/src/db/repositories/workspace_config_repo.py index 124cb7f3e..8377d1d72 100644 --- a/src/db/repositories/workspace_config_repo.py +++ b/src/db/repositories/workspace_config_repo.py @@ -23,7 +23,7 @@ class WorkspaceConfigRepo: def __init__(self, session: AsyncSession): self.session = session - async def get_section(self, section: str) -> Optional[dict[str, Any]]: + async def get_section(self, section: str) -> dict[str, Any] | None: row = await self.session.get(WorkspaceConfig, section) return None if row is None else (row.value or {}) @@ -35,7 +35,7 @@ async def upsert( self, section: str, value: dict[str, Any], - actor_user_id: Optional[str] = None, + actor_user_id: str | None = None, ) -> WorkspaceConfig: existing = await self.session.get(WorkspaceConfig, section) if existing is None: diff --git a/tests/unit/connectors/test_connector_enablement.py b/tests/unit/connectors/test_connector_enablement.py index 80399bf76..c1c36607e 100644 --- a/tests/unit/connectors/test_connector_enablement.py +++ b/tests/unit/connectors/test_connector_enablement.py @@ -37,9 +37,7 @@ def _connector_service(types=("google_drive", "sharepoint", "onedrive")): svc = MagicMock() cm = MagicMock() cm.get_available_connector_types = MagicMock( - return_value={ - t: {"name": t, "description": "", "icon": t} for t in types - } + return_value={t: {"name": t, "description": "", "icon": t} for t in types} ) svc.connection_manager = cm return svc From c669aee1805bfa4cb00fdd56e50e4b3cab784960 Mon Sep 17 00:00:00 2001 From: Edwin Jose Date: Wed, 3 Jun 2026 18:22:36 -0500 Subject: [PATCH 3/3] Remove unused imports and tidy auth imports Remove unused typing imports (Union and Optional) and clean up import ordering/whitespace. In src/api/auth.py move logging/telemetry/version imports and logger initialization to group related imports, and remove an unused Optional import in workspace_config_repo.py. Small refactor to reduce lint warnings and improve import organization. --- .../versions/0007_seed_connectors_manage_global.py | 1 - src/api/auth.py | 14 +++++--------- src/db/repositories/workspace_config_repo.py | 2 +- 3 files changed, 6 insertions(+), 11 deletions(-) diff --git a/alembic/versions/0007_seed_connectors_manage_global.py b/alembic/versions/0007_seed_connectors_manage_global.py index 0a3b4cb90..8f4a268a9 100644 --- a/alembic/versions/0007_seed_connectors_manage_global.py +++ b/alembic/versions/0007_seed_connectors_manage_global.py @@ -15,7 +15,6 @@ import uuid from collections.abc import Sequence -from typing import Union import sqlalchemy as sa diff --git a/src/api/auth.py b/src/api/auth.py index fa708bcb6..b22d1b4c2 100644 --- a/src/api/auth.py +++ b/src/api/auth.py @@ -1,14 +1,5 @@ -from typing import Optional - from fastapi import Depends, HTTPException, Request from fastapi.responses import JSONResponse - -from utils.logging_config import get_logger -from utils.telemetry import Category, MessageId, TelemetryClient -from utils.version_utils import OPENRAG_VERSION - -logger = get_logger(__name__) - from pydantic import BaseModel from sqlalchemy.ext.asyncio import AsyncSession @@ -20,6 +11,11 @@ get_rbac_service, ) from session_manager import User +from utils.logging_config import get_logger +from utils.telemetry import Category, MessageId, TelemetryClient +from utils.version_utils import OPENRAG_VERSION + +logger = get_logger(__name__) class AuthInitBody(BaseModel): diff --git a/src/db/repositories/workspace_config_repo.py b/src/db/repositories/workspace_config_repo.py index 8377d1d72..6c4db3517 100644 --- a/src/db/repositories/workspace_config_repo.py +++ b/src/db/repositories/workspace_config_repo.py @@ -6,7 +6,7 @@ """ from datetime import datetime -from typing import Any, Optional +from typing import Any from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession