mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-13 18:15:02 +00:00
Add the connector catalog and read-only connections API
connector_sessions already holds one row per signed-in account or MCP server. Migration 0038 names each row (connector_key, auth_kind, display_name, account_label) and lets sources and tools point at the row they use through connection_id, backfilling the links it can prove. docsgpt/connectors/catalog.py declares every service DocsGPT connects to, with the server settings it needs and what a connection provides. New endpoints: GET /api/connectors/catalog, GET /api/connections, GET /api/connections/<id> and POST /api/connections/<id>/disconnect. /api/available_tools now groups tools into built_in, service and custom, and source listings carry connectionId.
This commit is contained in:
arc53-machine committed
1 parent
335f2af70c
commit
5321a94388
17 files changed
+1714
-2
No files matched your search
@@ -0,0 +1,136 @@
|
||||
"""0038 connections — connector_sessions becomes the connections table.
|
||||
|
||||
``connector_sessions`` already holds one row per signed-in account (OAuth
|
||||
ingest providers) or per MCP server. This migration names what each row is
|
||||
and lets sources and tools point at the row they use:
|
||||
|
||||
* ``connector_key`` is the catalog entry (``google_drive``, ``custom_mcp``,
|
||||
``telegram``), ``auth_kind`` how the row signs in, ``display_name`` and
|
||||
``account_label`` what the Connectors page shows.
|
||||
* ``sources.connection_id`` and ``user_tools.connection_id`` link the
|
||||
resources a connection feeds. ``ON DELETE SET NULL`` keeps a source's
|
||||
indexed content when its connection is removed.
|
||||
|
||||
Backfill (idempotent, only fills NULLs):
|
||||
|
||||
1. ``connector_key``, ``auth_kind``, ``display_name`` from ``provider``.
|
||||
2. ``account_label`` from ``user_email`` for OAuth rows.
|
||||
3. ``sources.connection_id`` for ``connector:file`` sources, matched to the
|
||||
owner's only row for ``remote_data->>'provider'``.
|
||||
4. ``user_tools.connection_id`` for OAuth MCP tools, matched to the owner's
|
||||
row for the tool's server base URL.
|
||||
|
||||
Revision ID: 0038_connections
|
||||
Revises: 0037_request_traces
|
||||
"""
|
||||
|
||||
from typing import Sequence, Union
|
||||
|
||||
from alembic import op
|
||||
|
||||
|
||||
revision: str = "0038_connections"
|
||||
down_revision: Union[str, None] = "0037_request_traces"
|
||||
branch_labels: Union[str, Sequence[str], None] = None
|
||||
depends_on: Union[str, Sequence[str], None] = None
|
||||
|
||||
|
||||
_OAUTH_PROVIDERS = {
|
||||
"google_drive": "Google Drive",
|
||||
"share_point": "SharePoint",
|
||||
"confluence": "Confluence",
|
||||
}
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
op.execute(
|
||||
"""
|
||||
ALTER TABLE connector_sessions
|
||||
ADD COLUMN IF NOT EXISTS connector_key TEXT,
|
||||
ADD COLUMN IF NOT EXISTS display_name TEXT,
|
||||
ADD COLUMN IF NOT EXISTS account_label TEXT,
|
||||
ADD COLUMN IF NOT EXISTS auth_kind TEXT,
|
||||
ADD COLUMN IF NOT EXISTS updated_at TIMESTAMPTZ NOT NULL DEFAULT now();
|
||||
"""
|
||||
)
|
||||
op.execute(
|
||||
"ALTER TABLE sources ADD COLUMN IF NOT EXISTS connection_id UUID "
|
||||
"REFERENCES connector_sessions(id) ON DELETE SET NULL;"
|
||||
)
|
||||
op.execute(
|
||||
"ALTER TABLE user_tools ADD COLUMN IF NOT EXISTS connection_id UUID "
|
||||
"REFERENCES connector_sessions(id) ON DELETE SET NULL;"
|
||||
)
|
||||
op.execute("CREATE INDEX IF NOT EXISTS sources_connection_idx ON sources (connection_id);")
|
||||
op.execute("CREATE INDEX IF NOT EXISTS user_tools_connection_idx ON user_tools (connection_id);")
|
||||
|
||||
# 1 + 2: name the existing rows.
|
||||
for provider, name in _OAUTH_PROVIDERS.items():
|
||||
op.execute(
|
||||
f"""
|
||||
UPDATE connector_sessions SET
|
||||
connector_key = COALESCE(connector_key, '{provider}'),
|
||||
auth_kind = COALESCE(auth_kind, 'oauth'),
|
||||
display_name = COALESCE(display_name, '{name}'),
|
||||
account_label = COALESCE(account_label, user_email)
|
||||
WHERE provider = '{provider}';
|
||||
"""
|
||||
)
|
||||
op.execute(
|
||||
"""
|
||||
UPDATE connector_sessions SET
|
||||
connector_key = COALESCE(connector_key, 'custom_mcp'),
|
||||
auth_kind = COALESCE(auth_kind, 'mcp_oauth'),
|
||||
display_name = COALESCE(
|
||||
display_name,
|
||||
regexp_replace(COALESCE(server_url, substr(provider, 5)), '^https?://', '')
|
||||
)
|
||||
WHERE provider LIKE 'mcp:%';
|
||||
"""
|
||||
)
|
||||
|
||||
# 3: connector sources point at the owner's only session for the provider.
|
||||
op.execute(
|
||||
"""
|
||||
UPDATE sources s SET connection_id = cs.id
|
||||
FROM connector_sessions cs
|
||||
WHERE s.connection_id IS NULL
|
||||
AND s.type = 'connector:file'
|
||||
AND cs.user_id = s.user_id
|
||||
AND cs.provider = s.remote_data->>'provider'
|
||||
AND (
|
||||
SELECT count(*) FROM connector_sessions c2
|
||||
WHERE c2.user_id = s.user_id AND c2.provider = cs.provider
|
||||
) = 1;
|
||||
"""
|
||||
)
|
||||
|
||||
# 4: OAuth MCP tools point at the owner's session for the server's base URL.
|
||||
op.execute(
|
||||
"""
|
||||
UPDATE user_tools t SET connection_id = cs.id
|
||||
FROM connector_sessions cs
|
||||
WHERE t.connection_id IS NULL
|
||||
AND t.name = 'mcp_tool'
|
||||
AND t.config->>'auth_type' = 'oauth'
|
||||
AND cs.user_id = t.user_id
|
||||
AND cs.provider = 'mcp:' || substring(t.config->>'server_url' from '^(https?://[^/]+)');
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
op.execute("DROP INDEX IF EXISTS user_tools_connection_idx;")
|
||||
op.execute("DROP INDEX IF EXISTS sources_connection_idx;")
|
||||
op.execute("ALTER TABLE user_tools DROP COLUMN IF EXISTS connection_id;")
|
||||
op.execute("ALTER TABLE sources DROP COLUMN IF EXISTS connection_id;")
|
||||
op.execute(
|
||||
"""
|
||||
ALTER TABLE connector_sessions
|
||||
DROP COLUMN IF EXISTS updated_at,
|
||||
DROP COLUMN IF EXISTS auth_kind,
|
||||
DROP COLUMN IF EXISTS account_label,
|
||||
DROP COLUMN IF EXISTS display_name,
|
||||
DROP COLUMN IF EXISTS connector_key;
|
||||
"""
|
||||
)
|
||||
@@ -0,0 +1,106 @@
|
||||
"""Connectors catalog and connections API.
|
||||
|
||||
``/api/connectors/catalog`` lists every service DocsGPT can connect to, with
|
||||
whether the server is set up for it and the caller's connection summary.
|
||||
``/api/connections`` lists and manages the caller's connections. Responses
|
||||
never include tokens or secrets.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from flask import current_app, jsonify, make_response, request
|
||||
from flask_restx import Namespace, Resource
|
||||
|
||||
from docsgpt.api import api
|
||||
from docsgpt.api.user.authz import ROLE_ADMIN, has_role
|
||||
from docsgpt.connectors import service
|
||||
from docsgpt.storage.db.repositories.connector_sessions import ConnectorSessionsRepository
|
||||
from docsgpt.storage.db.session import db_readonly, db_session
|
||||
|
||||
connections_ns = Namespace("connections", description="Connectors and connections", path="/api")
|
||||
api.add_namespace(connections_ns)
|
||||
|
||||
|
||||
def _user_id() -> str | None:
|
||||
token = getattr(request, "decoded_token", None)
|
||||
return token.get("sub") if isinstance(token, dict) else None
|
||||
|
||||
|
||||
def _unauthorized():
|
||||
return make_response(jsonify({"success": False, "error": "Unauthorized"}), 401)
|
||||
|
||||
|
||||
def _not_found():
|
||||
return make_response(jsonify({"success": False, "error": "Connection not found"}), 404)
|
||||
|
||||
|
||||
@connections_ns.route("/connectors/catalog")
|
||||
class ConnectorCatalog(Resource):
|
||||
@api.doc(description="Every connector with its availability and the caller's connection summary")
|
||||
def get(self):
|
||||
user_id = _user_id()
|
||||
if not user_id:
|
||||
return _unauthorized()
|
||||
try:
|
||||
with db_readonly() as conn:
|
||||
entries = service.catalog_for_user(
|
||||
conn, user_id, is_admin=has_role(request.decoded_token, ROLE_ADMIN),
|
||||
)
|
||||
except Exception as err:
|
||||
current_app.logger.error(f"Error building connector catalog: {err}", exc_info=True)
|
||||
return make_response(jsonify({"success": False, "error": "Failed to load connectors"}), 500)
|
||||
return make_response(jsonify({"success": True, "connectors": entries}), 200)
|
||||
|
||||
|
||||
@connections_ns.route("/connections")
|
||||
class ConnectionsList(Resource):
|
||||
@api.doc(description="The caller's connections with status and linked resource counts")
|
||||
def get(self):
|
||||
user_id = _user_id()
|
||||
if not user_id:
|
||||
return _unauthorized()
|
||||
try:
|
||||
with db_readonly() as conn:
|
||||
connections = service.list_connections(conn, user_id)
|
||||
except Exception as err:
|
||||
current_app.logger.error(f"Error listing connections: {err}", exc_info=True)
|
||||
return make_response(jsonify({"success": False, "error": "Failed to load connections"}), 500)
|
||||
return make_response(jsonify({"success": True, "connections": connections}), 200)
|
||||
|
||||
|
||||
@connections_ns.route("/connections/<string:connection_id>")
|
||||
class ConnectionDetail(Resource):
|
||||
@api.doc(description="One connection with the sources it syncs and the tools it provides")
|
||||
def get(self, connection_id: str):
|
||||
user_id = _user_id()
|
||||
if not user_id:
|
||||
return _unauthorized()
|
||||
try:
|
||||
with db_readonly() as conn:
|
||||
row = ConnectorSessionsRepository(conn).get_for_user(connection_id, user_id)
|
||||
if row is None:
|
||||
return _not_found()
|
||||
detail = service.connection_detail(conn, row)
|
||||
except Exception as err:
|
||||
current_app.logger.error(f"Error loading connection: {err}", exc_info=True)
|
||||
return make_response(jsonify({"success": False, "error": "Failed to load connection"}), 500)
|
||||
return make_response(jsonify({"success": True, "connection": detail}), 200)
|
||||
|
||||
|
||||
@connections_ns.route("/connections/<string:connection_id>/disconnect")
|
||||
class ConnectionDisconnect(Resource):
|
||||
@api.doc(description="Delete a connection's stored credentials; its sources and tools stay")
|
||||
def post(self, connection_id: str):
|
||||
user_id = _user_id()
|
||||
if not user_id:
|
||||
return _unauthorized()
|
||||
try:
|
||||
with db_session() as conn:
|
||||
row = ConnectorSessionsRepository(conn).get_for_user(connection_id, user_id)
|
||||
if row is None:
|
||||
return _not_found()
|
||||
connection = service.disconnect(conn, row)
|
||||
except Exception as err:
|
||||
current_app.logger.error(f"Error disconnecting connection: {err}", exc_info=True)
|
||||
return make_response(jsonify({"success": False, "error": "Failed to disconnect"}), 500)
|
||||
return make_response(jsonify({"success": True, "connection": connection}), 200)
|
||||
@@ -357,6 +357,7 @@ DENIED_PREFIXES = (
|
||||
"/api/admin/",
|
||||
"/api/auth/oidc/",
|
||||
"/api/connectors/",
|
||||
"/api/connections",
|
||||
"/api/devices",
|
||||
"/scim/",
|
||||
"/static/",
|
||||
|
||||
@@ -67,6 +67,12 @@ def _get_provider_from_remote_data(remote_data):
|
||||
return None
|
||||
|
||||
|
||||
def _connection_id(row: dict) -> str | None:
|
||||
"""The connection a source syncs from, as a string id, or None."""
|
||||
value = row.get("connection_id")
|
||||
return str(value) if value else None
|
||||
|
||||
|
||||
@sources_ns.route("/sources")
|
||||
class CombinedJson(Resource):
|
||||
@api.doc(description="Provide JSON file with combined available indexes")
|
||||
@@ -117,6 +123,7 @@ class CombinedJson(Resource):
|
||||
"config": SourceConfig.parse(index.get("config")).model_dump(),
|
||||
"ownership": ownership,
|
||||
"team_access": team_access,
|
||||
"connectionId": _connection_id(index),
|
||||
}
|
||||
|
||||
for index in indexes:
|
||||
@@ -207,6 +214,7 @@ class PaginatedSources(Resource):
|
||||
"team_access": (
|
||||
None if owned else team_shared.get(str(doc["id"]))
|
||||
),
|
||||
"connectionId": _connection_id(doc),
|
||||
}
|
||||
)
|
||||
response = {
|
||||
|
||||
@@ -19,6 +19,7 @@ from docsgpt.api import api
|
||||
from docsgpt.api.pat.rules import filter_listing
|
||||
from docsgpt.api.user.artifacts.authz import Principal, authorize_artifact
|
||||
from docsgpt.api.user.team_sharing import effective_write_owner, visible_with_access
|
||||
from docsgpt.connectors.catalog import definition_for_tool
|
||||
from docsgpt.core.settings import settings
|
||||
from docsgpt.core.url_validation import SSRFError, validate_url
|
||||
from docsgpt.security.encryption import decrypt_credentials, encrypt_credentials
|
||||
@@ -186,6 +187,10 @@ def transform_actions(actions_metadata):
|
||||
|
||||
tools_ns = Namespace("tools", description="Tool management operations", path="/api")
|
||||
|
||||
# Tools the Connectors page adds through "Add custom connector" rather than
|
||||
# the Add Tool modal.
|
||||
_CUSTOM_CONNECTOR_TOOLS = {"mcp_tool": "custom_mcp", "api_tool": "custom_openapi"}
|
||||
|
||||
|
||||
@tools_ns.route("/available_tools")
|
||||
class AvailableTools(Resource):
|
||||
@@ -202,6 +207,13 @@ class AvailableTools(Resource):
|
||||
description = lines[1].strip() if len(lines) > 1 else ""
|
||||
config_req = tool_instance.get_config_requirements()
|
||||
actions = tool_instance.get_actions_metadata()
|
||||
definition = definition_for_tool(tool_name)
|
||||
if definition is not None:
|
||||
group, connector_key = "service", definition.key
|
||||
elif tool_name in _CUSTOM_CONNECTOR_TOOLS:
|
||||
group, connector_key = "custom", _CUSTOM_CONNECTOR_TOOLS[tool_name]
|
||||
else:
|
||||
group, connector_key = "built_in", None
|
||||
tools_metadata.append(
|
||||
{
|
||||
"name": tool_name,
|
||||
@@ -209,6 +221,8 @@ class AvailableTools(Resource):
|
||||
"description": description,
|
||||
"configRequirements": config_req,
|
||||
"actions": actions,
|
||||
"group": group,
|
||||
"connector_key": connector_key,
|
||||
}
|
||||
)
|
||||
except Exception as err:
|
||||
|
||||
@@ -29,6 +29,7 @@ from docsgpt.api.scim import scim_bp # noqa: E402
|
||||
from docsgpt.api.user.authz import ROLE_USER, resolve_roles # noqa: E402
|
||||
from docsgpt.api.user.routes import user # noqa: E402
|
||||
from docsgpt.api.connector.routes import connector # noqa: E402
|
||||
from docsgpt.api.connector import connections as _connections # noqa: E402,F401
|
||||
from docsgpt.api.v1 import v1_bp # noqa: E402
|
||||
from docsgpt.celery_init import celery # noqa: E402
|
||||
from docsgpt.core.secret_key import resolve_jwt_secret_key # noqa: E402
|
||||
@@ -209,6 +210,9 @@ def get_config():
|
||||
"hybrid_available": settings.VECTOR_STORE == "pgvector",
|
||||
"tts_available": TTSCreator.is_enabled(settings.TTS_PROVIDER),
|
||||
"stt_available": STTCreator.is_enabled(settings.STT_PROVIDER),
|
||||
# Lets a frontend built before the Connectors page run against this
|
||||
# backend, and a new frontend hide the page against an older one.
|
||||
"connectors_enabled": True,
|
||||
}
|
||||
if settings.AUTH_TYPE == "oidc":
|
||||
response["oidc"] = {
|
||||
|
||||
@@ -0,0 +1,6 @@
|
||||
"""Connectors: one place to connect accounts that feed Sources and Tools.
|
||||
|
||||
``catalog`` declares the services, ``service`` manages connections (the rows
|
||||
of ``connector_sessions``), ``permissions`` classifies tool actions as read
|
||||
or write.
|
||||
"""
|
||||
@@ -0,0 +1,423 @@
|
||||
"""The connector catalog: every service DocsGPT can connect to.
|
||||
|
||||
A connector is something the user connects once (an OAuth sign-in, an API key
|
||||
or an MCP server). Each connection it produces can feed Sources (content
|
||||
synced into DocsGPT) and Tools (actions an agent can take). The definitions
|
||||
here are declarative: they say how a connector signs in, which server
|
||||
settings it needs, what it can sync and which tools it creates, so the API
|
||||
and the frontend never hard-code a list of services.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Any, Iterable, Optional
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import yaml
|
||||
|
||||
from docsgpt.core.settings import settings
|
||||
|
||||
CATEGORIES = ("files", "knowledge", "dev", "messaging", "database", "search", "custom")
|
||||
AUTH_KINDS = ("oauth", "mcp_oauth", "api_key", "none", "mcp")
|
||||
PUBLISHERS = ("built_in", "preset", "custom")
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class CredentialField:
|
||||
"""One field the user fills in to connect an ``api_key`` connector.
|
||||
|
||||
Attributes:
|
||||
key: Name the credential is stored and passed to the tool or loader
|
||||
under, e.g. ``token`` or ``aws_access_key_id``.
|
||||
label: English label; the frontend shows the translated
|
||||
``settings.connectors.fields.<key>`` when it has one.
|
||||
secret: Masked in the form and never returned by the API.
|
||||
required: Whether connecting fails without it.
|
||||
"""
|
||||
|
||||
key: str
|
||||
label: str
|
||||
secret: bool = True
|
||||
required: bool = True
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
return {"key": self.key, "label": self.label, "secret": self.secret, "required": self.required}
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ConnectorDefinition:
|
||||
"""How one connector signs in and what a connection to it provides.
|
||||
|
||||
Attributes:
|
||||
key: Catalog key, stored on each connection as ``connector_key``.
|
||||
name: Display name.
|
||||
description: One line shown on the catalog card.
|
||||
icon: Frontend asset key (``assets/connectors/<icon>.svg``).
|
||||
category: One of :data:`CATEGORIES`.
|
||||
auth_kind: ``oauth`` (server-side OAuth app), ``mcp_oauth`` (MCP
|
||||
OAuth with dynamic client registration), ``api_key`` (fields the
|
||||
user pastes), ``none`` or ``mcp`` (a custom MCP server whose auth
|
||||
the user picks).
|
||||
capabilities: Any of ``sync``, ``read``, ``write``.
|
||||
credential_fields: Fields asked for by ``api_key`` connectors.
|
||||
required_settings: Server settings that must be set for the
|
||||
connector to be usable, e.g. ``GOOGLE_CLIENT_ID``.
|
||||
sync_ingestor: Ingest loader a synced source uses (``google_drive``,
|
||||
``s3``), or None when the connector cannot sync.
|
||||
default_sync_frequency: Sync frequency preselected in the wizard.
|
||||
setup_fields: Per-source fields asked when choosing what to sync
|
||||
(an S3 bucket, Reddit search queries).
|
||||
tool_templates: ``user_tools`` names created on connect.
|
||||
setup: What the wizard does after sign-in: ``tools`` is ``auto``
|
||||
(created and enabled), ``ask`` or ``off``; ``sync`` likewise.
|
||||
mcp_url: MCP endpoint for presets.
|
||||
publisher: ``built_in``, ``preset`` or ``custom``.
|
||||
docs_url: Setup guide for admins.
|
||||
oauth_scopes: Scopes an MCP preset requests.
|
||||
"""
|
||||
|
||||
key: str
|
||||
name: str
|
||||
description: str
|
||||
icon: str
|
||||
category: str
|
||||
auth_kind: str
|
||||
capabilities: tuple[str, ...] = ()
|
||||
credential_fields: tuple[CredentialField, ...] = ()
|
||||
required_settings: tuple[str, ...] = ()
|
||||
sync_ingestor: Optional[str] = None
|
||||
default_sync_frequency: str = "weekly"
|
||||
setup_fields: tuple[CredentialField, ...] = ()
|
||||
tool_templates: tuple[str, ...] = ()
|
||||
setup: dict = field(default_factory=lambda: {"tools": "auto", "sync": "ask"})
|
||||
mcp_url: Optional[str] = None
|
||||
publisher: str = "built_in"
|
||||
docs_url: Optional[str] = None
|
||||
oauth_scopes: tuple[str, ...] = ()
|
||||
|
||||
@property
|
||||
def missing_settings(self) -> list[str]:
|
||||
"""Server settings this connector needs that are not set."""
|
||||
return [name for name in self.required_settings if not getattr(settings, name, None)]
|
||||
|
||||
@property
|
||||
def configured(self) -> bool:
|
||||
"""Whether every required server setting is set."""
|
||||
return not self.missing_settings
|
||||
|
||||
@property
|
||||
def mcp_base_url(self) -> Optional[str]:
|
||||
"""``scheme://host`` of the preset's MCP endpoint (the MCP session key)."""
|
||||
return base_url(self.mcp_url) if self.mcp_url else None
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
"""Serialise the parts the frontend needs (never server secrets)."""
|
||||
return {
|
||||
"key": self.key,
|
||||
"name": self.name,
|
||||
"description": self.description,
|
||||
"icon": self.icon,
|
||||
"category": self.category,
|
||||
"auth_kind": self.auth_kind,
|
||||
"capabilities": list(self.capabilities),
|
||||
"credential_fields": [f.to_dict() for f in self.credential_fields],
|
||||
"setup_fields": [f.to_dict() for f in self.setup_fields],
|
||||
"sync_ingestor": self.sync_ingestor,
|
||||
"default_sync_frequency": self.default_sync_frequency,
|
||||
"tool_templates": list(self.tool_templates),
|
||||
"setup": dict(self.setup),
|
||||
"mcp_url": self.mcp_url,
|
||||
"publisher": self.publisher,
|
||||
"docs_url": self.docs_url,
|
||||
"oauth_scopes": list(self.oauth_scopes),
|
||||
}
|
||||
|
||||
|
||||
def base_url(url: Optional[str]) -> str:
|
||||
"""``scheme://netloc`` of ``url``, the key MCP sessions are stored under."""
|
||||
parsed = urlparse(url or "")
|
||||
if not parsed.scheme or not parsed.netloc:
|
||||
return ""
|
||||
return f"{parsed.scheme}://{parsed.netloc}"
|
||||
|
||||
|
||||
_DOCS = "https://docs.docsgpt.cloud/Guides/Connectors"
|
||||
|
||||
_BUILT_IN: tuple[ConnectorDefinition, ...] = (
|
||||
ConnectorDefinition(
|
||||
key="google_drive",
|
||||
name="Google Drive",
|
||||
description="Sync Docs, Sheets and PDFs into Sources.",
|
||||
icon="drive",
|
||||
category="files",
|
||||
auth_kind="oauth",
|
||||
capabilities=("sync",),
|
||||
required_settings=("GOOGLE_CLIENT_ID", "GOOGLE_CLIENT_SECRET"),
|
||||
sync_ingestor="google_drive",
|
||||
setup={"tools": "off", "sync": "ask"},
|
||||
docs_url=f"{_DOCS}#google-drive",
|
||||
),
|
||||
ConnectorDefinition(
|
||||
key="share_point",
|
||||
name="SharePoint",
|
||||
description="Sync files from SharePoint sites and OneDrive into Sources.",
|
||||
icon="sharepoint",
|
||||
category="files",
|
||||
auth_kind="oauth",
|
||||
capabilities=("sync",),
|
||||
required_settings=("MICROSOFT_CLIENT_ID", "MICROSOFT_CLIENT_SECRET"),
|
||||
sync_ingestor="share_point",
|
||||
setup={"tools": "off", "sync": "ask"},
|
||||
docs_url=f"{_DOCS}#sharepoint-and-onedrive",
|
||||
),
|
||||
ConnectorDefinition(
|
||||
key="confluence",
|
||||
name="Confluence",
|
||||
description="Sync Confluence spaces and pages into Sources.",
|
||||
icon="confluence",
|
||||
category="knowledge",
|
||||
auth_kind="oauth",
|
||||
capabilities=("sync",),
|
||||
required_settings=("CONFLUENCE_CLIENT_ID", "CONFLUENCE_CLIENT_SECRET"),
|
||||
sync_ingestor="confluence",
|
||||
setup={"tools": "off", "sync": "ask"},
|
||||
docs_url=f"{_DOCS}#confluence",
|
||||
),
|
||||
ConnectorDefinition(
|
||||
key="s3",
|
||||
name="Amazon S3",
|
||||
description="Sync documents from an S3 bucket into Sources.",
|
||||
icon="s3",
|
||||
category="files",
|
||||
auth_kind="api_key",
|
||||
capabilities=("sync",),
|
||||
credential_fields=(
|
||||
CredentialField("aws_access_key_id", "Access key ID", secret=False),
|
||||
CredentialField("aws_secret_access_key", "Secret access key"),
|
||||
CredentialField("region", "Region", secret=False, required=False),
|
||||
CredentialField("endpoint_url", "Custom endpoint URL", secret=False, required=False),
|
||||
),
|
||||
sync_ingestor="s3",
|
||||
setup_fields=(
|
||||
CredentialField("bucket", "Bucket", secret=False),
|
||||
CredentialField("prefix", "Path prefix", secret=False, required=False),
|
||||
),
|
||||
setup={"tools": "off", "sync": "ask"},
|
||||
docs_url=f"{_DOCS}#amazon-s3",
|
||||
),
|
||||
ConnectorDefinition(
|
||||
key="reddit",
|
||||
name="Reddit",
|
||||
description="Sync Reddit posts that match your searches into Sources.",
|
||||
icon="reddit",
|
||||
category="knowledge",
|
||||
auth_kind="api_key",
|
||||
capabilities=("sync",),
|
||||
credential_fields=(
|
||||
CredentialField("client_id", "Client ID", secret=False),
|
||||
CredentialField("client_secret", "Client secret"),
|
||||
CredentialField("user_agent", "User agent", secret=False),
|
||||
),
|
||||
sync_ingestor="reddit",
|
||||
setup_fields=(
|
||||
CredentialField("search_queries", "Search queries", secret=False),
|
||||
CredentialField("number_posts", "Number of posts", secret=False),
|
||||
),
|
||||
setup={"tools": "off", "sync": "ask"},
|
||||
),
|
||||
ConnectorDefinition(
|
||||
key="brave",
|
||||
name="Brave Search",
|
||||
description="Search the web and images with the Brave Search API.",
|
||||
icon="tool_brave",
|
||||
category="search",
|
||||
auth_kind="api_key",
|
||||
capabilities=("read",),
|
||||
credential_fields=(CredentialField("token", "API key"),),
|
||||
tool_templates=("brave",),
|
||||
setup={"tools": "auto", "sync": "off"},
|
||||
),
|
||||
ConnectorDefinition(
|
||||
key="telegram",
|
||||
name="Telegram",
|
||||
description="Send messages and images to a Telegram chat.",
|
||||
icon="tool_telegram",
|
||||
category="messaging",
|
||||
auth_kind="api_key",
|
||||
capabilities=("write",),
|
||||
credential_fields=(CredentialField("token", "Bot token"),),
|
||||
tool_templates=("telegram",),
|
||||
setup={"tools": "auto", "sync": "off"},
|
||||
),
|
||||
ConnectorDefinition(
|
||||
key="ntfy",
|
||||
name="ntfy",
|
||||
description="Send push notifications through an ntfy server.",
|
||||
icon="tool_ntfy",
|
||||
category="messaging",
|
||||
auth_kind="api_key",
|
||||
capabilities=("write",),
|
||||
credential_fields=(CredentialField("token", "Access token"),),
|
||||
tool_templates=("ntfy",),
|
||||
setup={"tools": "auto", "sync": "off"},
|
||||
),
|
||||
ConnectorDefinition(
|
||||
key="postgres",
|
||||
name="PostgreSQL",
|
||||
description="Read the schema and run SQL against a Postgres database.",
|
||||
icon="tool_postgres",
|
||||
category="database",
|
||||
auth_kind="api_key",
|
||||
capabilities=("read", "write"),
|
||||
credential_fields=(CredentialField("token", "Connection string"),),
|
||||
tool_templates=("postgres",),
|
||||
setup={"tools": "auto", "sync": "off"},
|
||||
),
|
||||
ConnectorDefinition(
|
||||
key="custom_mcp",
|
||||
name="MCP server",
|
||||
description="Connect any remote Model Context Protocol server.",
|
||||
icon="tool_mcp_tool",
|
||||
category="custom",
|
||||
auth_kind="mcp",
|
||||
capabilities=("read", "write"),
|
||||
tool_templates=("mcp_tool",),
|
||||
setup={"tools": "auto", "sync": "off"},
|
||||
publisher="custom",
|
||||
),
|
||||
ConnectorDefinition(
|
||||
key="custom_openapi",
|
||||
name="OpenAPI / REST",
|
||||
description="Import an OpenAPI spec and call its endpoints as tools.",
|
||||
icon="tool_api_tool",
|
||||
category="custom",
|
||||
auth_kind="none",
|
||||
capabilities=("read", "write"),
|
||||
tool_templates=("api_tool",),
|
||||
setup={"tools": "ask", "sync": "off"},
|
||||
publisher="custom",
|
||||
),
|
||||
)
|
||||
|
||||
_PRESETS_FILE = Path(__file__).parent / "presets" / "mcp.yaml"
|
||||
|
||||
|
||||
def _load_presets(path: Optional[Path] = None) -> tuple[ConnectorDefinition, ...]:
|
||||
"""Read the curated MCP server presets shipped with the repository."""
|
||||
path = path or _PRESETS_FILE
|
||||
if not path.exists():
|
||||
return ()
|
||||
with path.open(encoding="utf-8") as fh:
|
||||
entries = yaml.safe_load(fh) or []
|
||||
presets = []
|
||||
for entry in entries:
|
||||
presets.append(
|
||||
ConnectorDefinition(
|
||||
key=entry["key"],
|
||||
name=entry["name"],
|
||||
description=entry["description"],
|
||||
icon=entry.get("icon") or "tool_mcp_tool",
|
||||
category=entry.get("category", "knowledge"),
|
||||
auth_kind=entry.get("auth_kind", "mcp_oauth"),
|
||||
capabilities=tuple(entry.get("capabilities") or ("read", "write")),
|
||||
tool_templates=("mcp_tool",),
|
||||
setup={"tools": "auto", "sync": "off"},
|
||||
mcp_url=entry["mcp_url"],
|
||||
publisher="preset",
|
||||
docs_url=entry.get("docs_url"),
|
||||
oauth_scopes=tuple(entry.get("oauth_scopes") or ()),
|
||||
)
|
||||
)
|
||||
return tuple(presets)
|
||||
|
||||
|
||||
_REGISTRY: dict[str, ConnectorDefinition] = {}
|
||||
|
||||
|
||||
def _registry() -> dict[str, ConnectorDefinition]:
|
||||
if not _REGISTRY:
|
||||
for definition in (*_BUILT_IN, *_load_presets()):
|
||||
_REGISTRY[definition.key] = definition
|
||||
return _REGISTRY
|
||||
|
||||
|
||||
def all_definitions() -> list[ConnectorDefinition]:
|
||||
"""Every catalog entry, built-ins first, then presets."""
|
||||
return list(_registry().values())
|
||||
|
||||
|
||||
def get_definition(key: Optional[str]) -> Optional[ConnectorDefinition]:
|
||||
"""The definition for ``key``, or None when it is not in the catalog."""
|
||||
if not key:
|
||||
return None
|
||||
return _registry().get(key)
|
||||
|
||||
|
||||
def preset_for_url(url: Optional[str]) -> Optional[ConnectorDefinition]:
|
||||
"""The MCP preset whose server shares ``url``'s base URL, if any."""
|
||||
target = base_url(url)
|
||||
if not target:
|
||||
return None
|
||||
for definition in _registry().values():
|
||||
if definition.publisher == "preset" and definition.mcp_base_url == target:
|
||||
return definition
|
||||
return None
|
||||
|
||||
|
||||
def definition_for_tool(tool_name: str) -> Optional[ConnectorDefinition]:
|
||||
"""The built-in connector that provides the ``user_tools`` template ``tool_name``."""
|
||||
for definition in _BUILT_IN:
|
||||
if definition.publisher == "built_in" and tool_name in definition.tool_templates:
|
||||
return definition
|
||||
return None
|
||||
|
||||
|
||||
def connector_key_for_row(row: dict) -> Optional[str]:
|
||||
"""Catalog key for a ``connector_sessions`` row.
|
||||
|
||||
Rows written before ``0038_connections`` carry no key; they are named from
|
||||
``provider``. A custom MCP row whose server matches a preset is reported
|
||||
as that preset.
|
||||
"""
|
||||
key = row.get("connector_key")
|
||||
provider = row.get("provider") or ""
|
||||
if not key:
|
||||
if provider.startswith("mcp:"):
|
||||
key = "custom_mcp"
|
||||
elif get_definition(provider):
|
||||
key = provider
|
||||
if key == "custom_mcp":
|
||||
preset = preset_for_url(row.get("server_url") or provider[4:])
|
||||
if preset:
|
||||
return preset.key
|
||||
return key
|
||||
|
||||
|
||||
def tool_connector_keys() -> set[str]:
|
||||
"""``user_tools`` names that belong to a built-in service connector."""
|
||||
return {
|
||||
name
|
||||
for definition in _BUILT_IN
|
||||
if definition.publisher == "built_in"
|
||||
for name in definition.tool_templates
|
||||
}
|
||||
|
||||
|
||||
def iter_by_category(definitions: Iterable[ConnectorDefinition]) -> dict[str, list[ConnectorDefinition]]:
|
||||
"""Group definitions by category, keeping :data:`CATEGORIES` order."""
|
||||
grouped: dict[str, list[ConnectorDefinition]] = {c: [] for c in CATEGORIES}
|
||||
for definition in definitions:
|
||||
grouped.setdefault(definition.category, []).append(definition)
|
||||
return grouped
|
||||
|
||||
|
||||
def reset_registry_for_tests() -> None:
|
||||
"""Drop the cached registry so a test can load a different presets file."""
|
||||
_REGISTRY.clear()
|
||||
|
||||
|
||||
def to_public(definition: ConnectorDefinition, **extra: Any) -> dict:
|
||||
"""Definition dict plus request-specific fields (availability, counts)."""
|
||||
return {**definition.to_dict(), **extra}
|
||||
@@ -0,0 +1,83 @@
|
||||
"""Read / write classification and permissions for connection-backed tools.
|
||||
|
||||
Every action a connection's tool offers is either a *read* (it only looks
|
||||
something up) or a *write* (it changes or sends something). Reads default to
|
||||
"Always allow"; writes default to "Needs approval". The permission is stored
|
||||
on the action itself: ``active`` off means "Off", ``require_approval`` means
|
||||
"Needs approval", neither means "Always allow".
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Optional
|
||||
|
||||
ACCESS_READ = "read"
|
||||
ACCESS_WRITE = "write"
|
||||
|
||||
PERMISSION_ALWAYS = "always"
|
||||
PERMISSION_ASK = "ask"
|
||||
PERMISSION_OFF = "off"
|
||||
PERMISSIONS = (PERMISSION_ALWAYS, PERMISSION_ASK, PERMISSION_OFF)
|
||||
|
||||
_READ_WORDS = ("search", "query", "find", "list", "get", "read", "fetch", "lookup", "describe", "retrieve")
|
||||
|
||||
|
||||
def action_access(tool_name: Optional[str], action: dict) -> str:
|
||||
"""Whether ``action`` reads or writes.
|
||||
|
||||
Order of evidence: an explicit ``access`` on the action metadata, MCP tool
|
||||
annotations (``readOnlyHint`` / ``destructiveHint``), the HTTP method of an
|
||||
API tool action, then the action's name.
|
||||
|
||||
Args:
|
||||
tool_name: The ``user_tools`` name the action belongs to.
|
||||
action: One entry of the tool's ``actions``.
|
||||
|
||||
Returns:
|
||||
``read`` or ``write``.
|
||||
"""
|
||||
access = action.get("access")
|
||||
if access in (ACCESS_READ, ACCESS_WRITE):
|
||||
return access
|
||||
annotations = action.get("annotations") or {}
|
||||
if isinstance(annotations, dict):
|
||||
if annotations.get("readOnlyHint") is True:
|
||||
return ACCESS_READ
|
||||
if annotations.get("destructiveHint") is True or annotations.get("readOnlyHint") is False:
|
||||
return ACCESS_WRITE
|
||||
method = (action.get("method") or "").upper()
|
||||
if tool_name == "api_tool" and method:
|
||||
return ACCESS_READ if method in ("GET", "HEAD", "OPTIONS") else ACCESS_WRITE
|
||||
name = (action.get("name") or "").lower()
|
||||
return ACCESS_READ if any(word in name for word in _READ_WORDS) else ACCESS_WRITE
|
||||
|
||||
|
||||
def action_permission(action: dict) -> str:
|
||||
"""The action's permission: ``always``, ``ask`` or ``off``."""
|
||||
if action.get("active") is False:
|
||||
return PERMISSION_OFF
|
||||
if action.get("require_approval"):
|
||||
return PERMISSION_ASK
|
||||
return PERMISSION_ALWAYS
|
||||
|
||||
|
||||
def apply_permission(action: dict, permission: str) -> dict:
|
||||
"""Return ``action`` with ``permission`` written onto its flags."""
|
||||
if permission not in PERMISSIONS:
|
||||
raise ValueError(f"Unknown permission: {permission}")
|
||||
updated = dict(action)
|
||||
updated["active"] = permission != PERMISSION_OFF
|
||||
updated["require_approval"] = permission == PERMISSION_ASK
|
||||
return updated
|
||||
|
||||
|
||||
def apply_default_permissions(tool_name: Optional[str], actions: list[dict]) -> list[dict]:
|
||||
"""Stamp ``access`` on each action and default writes to needing approval."""
|
||||
stamped = []
|
||||
for action in actions:
|
||||
access = action_access(tool_name, action)
|
||||
updated = {**action, "access": access}
|
||||
if access == ACCESS_WRITE:
|
||||
updated["require_approval"] = True
|
||||
stamped.append(updated)
|
||||
return stamped
|
||||
@@ -0,0 +1,264 @@
|
||||
"""Connection lifecycle: listing, status, and what each connection provides.
|
||||
|
||||
A connection is one ``connector_sessions`` row: a signed-in OAuth account, an
|
||||
MCP server or a set of API credentials. Sources and tools point at it through
|
||||
``connection_id``. Everything the API returns about a connection goes through
|
||||
:func:`serialize_connection`, which never includes credentials.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from typing import Any, Iterable, Optional
|
||||
|
||||
from docsgpt.connectors import catalog
|
||||
from docsgpt.connectors.catalog import ConnectorDefinition
|
||||
from docsgpt.storage.db.repositories.connector_sessions import ConnectorSessionsRepository
|
||||
|
||||
STATUS_CONNECTED = "connected"
|
||||
STATUS_RECONNECT = "reconnect_needed"
|
||||
STATUS_DISCONNECTED = "disconnected"
|
||||
STATUS_ERROR = "error"
|
||||
STATUS_PENDING = "pending"
|
||||
|
||||
# Worst first: a card shows the worst status among the user's accounts.
|
||||
_STATUS_SEVERITY = {
|
||||
STATUS_RECONNECT: 4,
|
||||
STATUS_ERROR: 3,
|
||||
STATUS_DISCONNECTED: 2,
|
||||
STATUS_CONNECTED: 1,
|
||||
STATUS_PENDING: 0,
|
||||
}
|
||||
|
||||
|
||||
def _json(value: Any) -> Any:
|
||||
if isinstance(value, str):
|
||||
try:
|
||||
return json.loads(value)
|
||||
except ValueError:
|
||||
return None
|
||||
return value
|
||||
|
||||
|
||||
def has_credentials(row: dict) -> bool:
|
||||
"""Whether a connection holds usable credentials (without decrypting them)."""
|
||||
if row.get("encrypted_credentials"):
|
||||
return True
|
||||
token_info = _json(row.get("token_info")) or {}
|
||||
if isinstance(token_info, dict) and token_info.get("access_token"):
|
||||
return True
|
||||
session_data = _json(row.get("session_data")) or {}
|
||||
tokens = session_data.get("tokens") if isinstance(session_data, dict) else None
|
||||
return bool(isinstance(tokens, dict) and tokens.get("access_token"))
|
||||
|
||||
|
||||
def normalize_status(row: dict) -> str:
|
||||
"""Map a row's stored status onto the connection status set.
|
||||
|
||||
Rows written before ``0038_connections`` use ``authorized`` for a
|
||||
finished OAuth sign-in, and MCP rows carry no status at all; both are
|
||||
read from whether the row holds credentials.
|
||||
"""
|
||||
raw = (row.get("status") or "").lower()
|
||||
if raw in (STATUS_RECONNECT, STATUS_DISCONNECTED, STATUS_ERROR):
|
||||
return raw
|
||||
if raw == STATUS_PENDING and not has_credentials(row):
|
||||
return STATUS_PENDING
|
||||
if raw in ("authorized", STATUS_CONNECTED, "active") or has_credentials(row):
|
||||
return STATUS_CONNECTED
|
||||
return STATUS_PENDING
|
||||
|
||||
|
||||
def worst_status(statuses: Iterable[str]) -> Optional[str]:
|
||||
"""The status that needs the most attention, or None for an empty list."""
|
||||
ranked = sorted(statuses, key=lambda s: _STATUS_SEVERITY.get(s, 0), reverse=True)
|
||||
return ranked[0] if ranked else None
|
||||
|
||||
|
||||
def account_label(row: dict) -> str:
|
||||
"""What identifies the account to its owner: an email, a workspace or a host."""
|
||||
return (
|
||||
row.get("account_label")
|
||||
or row.get("user_email")
|
||||
or row.get("display_name")
|
||||
or catalog.base_url(row.get("server_url")).split("://")[-1]
|
||||
or ""
|
||||
)
|
||||
|
||||
|
||||
def _iso(value: Any) -> Optional[str]:
|
||||
return value.isoformat() if hasattr(value, "isoformat") else value
|
||||
|
||||
|
||||
def serialize_connection(row: dict, counts: Optional[dict] = None) -> dict:
|
||||
"""Public shape of a connection. Never includes tokens or secrets."""
|
||||
key = catalog.connector_key_for_row(row)
|
||||
definition = catalog.get_definition(key)
|
||||
counts = counts or {}
|
||||
return {
|
||||
"id": str(row["id"]),
|
||||
"connector_key": key,
|
||||
"name": (definition.name if definition and definition.publisher != "custom" else None)
|
||||
or row.get("display_name")
|
||||
or (definition.name if definition else key),
|
||||
"display_name": row.get("display_name"),
|
||||
"icon": definition.icon if definition else "tool_mcp_tool",
|
||||
"account_label": account_label(row),
|
||||
"auth_kind": row.get("auth_kind") or (definition.auth_kind if definition else None),
|
||||
"status": normalize_status(row),
|
||||
"server_url": row.get("server_url"),
|
||||
"last_error": row.get("last_error"),
|
||||
"created_at": _iso(row.get("created_at")),
|
||||
"updated_at": _iso(row.get("updated_at")),
|
||||
"last_used_at": _iso(row.get("last_used_at")),
|
||||
"source_count": counts.get("sources", 0),
|
||||
"tool_count": counts.get("tools", 0),
|
||||
"credential_mode": row.get("credential_mode"),
|
||||
}
|
||||
|
||||
|
||||
def list_connections(conn, user_id: str) -> list[dict]:
|
||||
"""The user's connections, finished sign-ins only, with resource counts."""
|
||||
repo = ConnectorSessionsRepository(conn)
|
||||
rows = [r for r in repo.list_for_user(user_id) if normalize_status(r) != STATUS_PENDING]
|
||||
counts = repo.resource_counts([str(r["id"]) for r in rows])
|
||||
return [serialize_connection(r, counts.get(str(r["id"]))) for r in rows]
|
||||
|
||||
|
||||
def serialize_source(row: dict, connection_status: Optional[str] = None) -> dict:
|
||||
"""A source linked to a connection, as the connection drawer lists it."""
|
||||
metadata = _json(row.get("metadata")) or {}
|
||||
sync_state = metadata.get("sync_state") if isinstance(metadata, dict) else None
|
||||
if connection_status in (STATUS_RECONNECT, STATUS_DISCONNECTED) and not sync_state:
|
||||
sync_state = "paused_reconnect"
|
||||
return {
|
||||
"id": str(row["id"]),
|
||||
"name": row.get("name"),
|
||||
"type": row.get("type"),
|
||||
"last_sync": _iso(row.get("date")),
|
||||
"sync_frequency": row.get("sync_frequency") or "never",
|
||||
"sync_state": sync_state or "active",
|
||||
}
|
||||
|
||||
|
||||
def serialize_tool(row: dict) -> dict:
|
||||
"""A tool linked to a connection, with its actions and their permissions."""
|
||||
from docsgpt.connectors.permissions import action_access, action_permission
|
||||
|
||||
actions = []
|
||||
for action in _json(row.get("actions")) or []:
|
||||
if not isinstance(action, dict):
|
||||
continue
|
||||
access = action_access(row.get("name"), action)
|
||||
actions.append(
|
||||
{
|
||||
"name": action.get("name"),
|
||||
"description": action.get("description", ""),
|
||||
"access": access,
|
||||
"permission": action_permission(action),
|
||||
}
|
||||
)
|
||||
return {
|
||||
"id": str(row["id"]),
|
||||
"name": row.get("name"),
|
||||
"display_name": row.get("custom_name") or row.get("display_name") or row.get("name"),
|
||||
"status": bool(row.get("status")),
|
||||
"credential_mode": row.get("credential_mode") or "owner",
|
||||
"actions": actions,
|
||||
}
|
||||
|
||||
|
||||
def connection_detail(conn, row: dict) -> dict:
|
||||
"""A connection with the sources and tools it feeds."""
|
||||
repo = ConnectorSessionsRepository(conn)
|
||||
status = normalize_status(row)
|
||||
sources = [serialize_source(s, status) for s in repo.list_sources(str(row["id"]))]
|
||||
tools = [serialize_tool(t) for t in repo.list_tools(str(row["id"]))]
|
||||
detail = serialize_connection(row, {"sources": len(sources), "tools": len(tools)})
|
||||
detail["sources"] = sources
|
||||
detail["tools"] = tools
|
||||
return detail
|
||||
|
||||
|
||||
def _card_state(definition: ConnectorDefinition, available: bool, disabled: bool, status: Optional[str]) -> str:
|
||||
if disabled:
|
||||
return "disabled"
|
||||
if status in (STATUS_RECONNECT, STATUS_ERROR):
|
||||
return "reconnect"
|
||||
if status == STATUS_CONNECTED:
|
||||
return "connected"
|
||||
if not available:
|
||||
return "needs_setup"
|
||||
if definition.publisher == "custom":
|
||||
return "custom"
|
||||
return "available"
|
||||
|
||||
|
||||
def catalog_for_user(conn, user_id: str, *, is_admin: bool, policies: Optional[dict] = None) -> list[dict]:
|
||||
"""Catalog entries with availability and the caller's connection summary.
|
||||
|
||||
Args:
|
||||
conn: Open database connection.
|
||||
user_id: The caller.
|
||||
is_admin: Admins see the names of missing server settings; everyone
|
||||
else only learns that setup is needed.
|
||||
policies: ``connector_key`` to policy row, when admin policies exist.
|
||||
"""
|
||||
policies = policies or {}
|
||||
by_key: dict[str, list[str]] = {}
|
||||
for connection in list_connections(conn, user_id):
|
||||
by_key.setdefault(connection["connector_key"], []).append(connection["status"])
|
||||
|
||||
entries = []
|
||||
for definition in catalog.all_definitions():
|
||||
policy = policies.get(definition.key) or {}
|
||||
disabled = policy.get("enabled") is False
|
||||
missing = definition.missing_settings
|
||||
available = not missing and not disabled
|
||||
statuses = by_key.get(definition.key, [])
|
||||
status = worst_status(statuses)
|
||||
entries.append(
|
||||
{
|
||||
**definition.to_dict(),
|
||||
"available": available,
|
||||
"disabled": disabled,
|
||||
"needs_setup": bool(missing),
|
||||
"missing_settings": missing if is_admin else [],
|
||||
"connected_count": sum(1 for s in statuses if s == STATUS_CONNECTED),
|
||||
"connection_count": len(statuses),
|
||||
"status": status,
|
||||
"state": _card_state(definition, available, disabled, status),
|
||||
"credential_policy": policy.get("credential_mode") or "choose",
|
||||
}
|
||||
)
|
||||
return entries
|
||||
|
||||
|
||||
def disconnect(conn, row: dict) -> dict:
|
||||
"""Forget a connection's credentials but keep the row and what it feeds.
|
||||
|
||||
Sources keep their indexed content and stop syncing; tools stop working
|
||||
until the account is reconnected. An MCP server's client registration is
|
||||
kept so reconnecting skips dynamic client registration.
|
||||
|
||||
Args:
|
||||
conn: Open database connection inside a transaction.
|
||||
row: The connection row, already authorised for the caller.
|
||||
|
||||
Returns:
|
||||
The connection's public shape after the change.
|
||||
"""
|
||||
repo = ConnectorSessionsRepository(conn)
|
||||
session_data = _json(row.get("session_data")) or {}
|
||||
if isinstance(session_data, dict):
|
||||
session_data = {k: v for k, v in session_data.items() if k != "tokens"}
|
||||
repo.update(
|
||||
str(row["id"]),
|
||||
{
|
||||
"status": STATUS_DISCONNECTED,
|
||||
"token_info": None,
|
||||
"session_token": None,
|
||||
"session_data": session_data or {},
|
||||
},
|
||||
)
|
||||
return serialize_connection(repo.get(str(row["id"])))
|
||||
@@ -216,6 +216,9 @@ user_tools_table = Table(
|
||||
Column("created_at", DateTime(timezone=True), nullable=False, server_default=func.now()),
|
||||
Column("updated_at", DateTime(timezone=True), nullable=False, server_default=func.now()),
|
||||
Column("legacy_mongo_id", Text),
|
||||
Column(
|
||||
"connection_id", UUID(as_uuid=True), ForeignKey("connector_sessions.id", ondelete="SET NULL"),
|
||||
),
|
||||
)
|
||||
|
||||
token_usage_table = Table(
|
||||
@@ -338,6 +341,9 @@ sources_table = Table(
|
||||
Column("created_at", DateTime(timezone=True), nullable=False, server_default=func.now()),
|
||||
Column("updated_at", DateTime(timezone=True), nullable=False, server_default=func.now()),
|
||||
Column("legacy_mongo_id", Text),
|
||||
Column(
|
||||
"connection_id", UUID(as_uuid=True), ForeignKey("connector_sessions.id", ondelete="SET NULL"),
|
||||
),
|
||||
)
|
||||
|
||||
agents_table = Table(
|
||||
@@ -585,6 +591,13 @@ connector_sessions_table = Table(
|
||||
Column("expires_at", DateTime(timezone=True)),
|
||||
Column("created_at", DateTime(timezone=True), nullable=False, server_default=func.now()),
|
||||
Column("legacy_mongo_id", Text),
|
||||
# Added in ``0038_connections``: each row is a connection (one signed-in
|
||||
# account, one MCP server or one set of API credentials).
|
||||
Column("connector_key", Text),
|
||||
Column("display_name", Text),
|
||||
Column("account_label", Text),
|
||||
Column("auth_kind", Text),
|
||||
Column("updated_at", DateTime(timezone=True), nullable=False, server_default=func.now()),
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -24,12 +24,13 @@ from typing import Any, Optional
|
||||
|
||||
from sqlalchemy import Connection, text
|
||||
|
||||
from docsgpt.storage.db.base_repository import row_to_dict
|
||||
from docsgpt.storage.db.base_repository import looks_like_uuid, row_to_dict
|
||||
from docsgpt.storage.db.serialization import PGNativeJSONEncoder
|
||||
|
||||
|
||||
_UPDATABLE_SCALARS = {
|
||||
"server_url", "session_token", "user_email", "status", "expires_at",
|
||||
"connector_key", "display_name", "account_label", "auth_kind",
|
||||
}
|
||||
_UPDATABLE_JSONB = {"session_data", "token_info"}
|
||||
|
||||
@@ -176,11 +177,83 @@ class ConnectorSessionsRepository:
|
||||
|
||||
def list_for_user(self, user_id: str) -> list[dict]:
|
||||
result = self._conn.execute(
|
||||
text("SELECT * FROM connector_sessions WHERE user_id = :user_id"),
|
||||
text("SELECT * FROM connector_sessions WHERE user_id = :user_id ORDER BY created_at"),
|
||||
{"user_id": user_id},
|
||||
)
|
||||
return [row_to_dict(r) for r in result.fetchall()]
|
||||
|
||||
def get(self, connection_id: str) -> Optional[dict]:
|
||||
"""Fetch a connection by id, whoever owns it. Callers authorise."""
|
||||
if not looks_like_uuid(connection_id):
|
||||
return None
|
||||
result = self._conn.execute(
|
||||
text("SELECT * FROM connector_sessions WHERE id = CAST(:id AS uuid)"),
|
||||
{"id": str(connection_id)},
|
||||
)
|
||||
row = result.fetchone()
|
||||
return row_to_dict(row) if row is not None else None
|
||||
|
||||
def get_for_user(self, connection_id: str, user_id: str) -> Optional[dict]:
|
||||
"""Fetch a connection only when ``user_id`` owns it."""
|
||||
row = self.get(connection_id)
|
||||
if row is None or row.get("user_id") != user_id:
|
||||
return None
|
||||
return row
|
||||
|
||||
def get_for_update(self, connection_id: str) -> Optional[dict]:
|
||||
"""Fetch and row-lock a connection until the transaction ends.
|
||||
|
||||
Token refresh holds this lock so two workers refreshing a rotating
|
||||
refresh token (Microsoft, Atlassian) cannot both spend it.
|
||||
"""
|
||||
if not looks_like_uuid(connection_id):
|
||||
return None
|
||||
result = self._conn.execute(
|
||||
text("SELECT * FROM connector_sessions WHERE id = CAST(:id AS uuid) FOR UPDATE"),
|
||||
{"id": str(connection_id)},
|
||||
)
|
||||
row = result.fetchone()
|
||||
return row_to_dict(row) if row is not None else None
|
||||
|
||||
def resource_counts(self, connection_ids: list[str]) -> dict[str, dict[str, int]]:
|
||||
"""Number of sources and tools linked to each connection id."""
|
||||
ids = [str(i) for i in connection_ids if looks_like_uuid(str(i))]
|
||||
counts: dict[str, dict[str, int]] = {i: {"sources": 0, "tools": 0} for i in ids}
|
||||
if not ids:
|
||||
return counts
|
||||
for table, key in (("sources", "sources"), ("user_tools", "tools")):
|
||||
result = self._conn.execute(
|
||||
text(
|
||||
f"SELECT connection_id, count(*) FROM {table} "
|
||||
"WHERE connection_id = ANY(CAST(:ids AS uuid[])) GROUP BY connection_id"
|
||||
),
|
||||
{"ids": ids},
|
||||
)
|
||||
for connection_id, count in result.fetchall():
|
||||
counts[str(connection_id)][key] = int(count)
|
||||
return counts
|
||||
|
||||
def list_sources(self, connection_id: str) -> list[dict]:
|
||||
"""Sources synced from a connection, newest first."""
|
||||
result = self._conn.execute(
|
||||
text(
|
||||
"SELECT id, name, type, date, sync_frequency, metadata, remote_data, user_id "
|
||||
"FROM sources WHERE connection_id = CAST(:id AS uuid) ORDER BY date DESC"
|
||||
),
|
||||
{"id": str(connection_id)},
|
||||
)
|
||||
return [row_to_dict(r) for r in result.fetchall()]
|
||||
|
||||
def list_tools(self, connection_id: str) -> list[dict]:
|
||||
"""Tools a connection provides, oldest first."""
|
||||
result = self._conn.execute(
|
||||
text(
|
||||
"SELECT * FROM user_tools WHERE connection_id = CAST(:id AS uuid) ORDER BY created_at"
|
||||
),
|
||||
{"id": str(connection_id)},
|
||||
)
|
||||
return [row_to_dict(r) for r in result.fetchall()]
|
||||
|
||||
def update(self, session_id: str, fields: dict) -> bool:
|
||||
"""Partial update by PG UUID."""
|
||||
filtered = {
|
||||
@@ -198,6 +271,7 @@ class ConnectorSessionsRepository:
|
||||
else:
|
||||
set_clauses.append(f"{col} = :{col}")
|
||||
params[col] = val
|
||||
set_clauses.append("updated_at = now()")
|
||||
result = self._conn.execute(
|
||||
text(
|
||||
f"UPDATE connector_sessions SET {', '.join(set_clauses)} "
|
||||
|
||||
Whitespace-only changes.
@@ -0,0 +1,116 @@
|
||||
"""Tests for the connector catalog."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from docsgpt.connectors import catalog
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _fresh_registry():
|
||||
catalog.reset_registry_for_tests()
|
||||
yield
|
||||
catalog.reset_registry_for_tests()
|
||||
|
||||
|
||||
class TestDefinitions:
|
||||
def test_built_ins_are_registered(self):
|
||||
keys = {d.key for d in catalog.all_definitions()}
|
||||
assert {
|
||||
"google_drive", "share_point", "confluence", "s3", "reddit",
|
||||
"brave", "telegram", "ntfy", "postgres", "custom_mcp", "custom_openapi",
|
||||
} <= keys
|
||||
|
||||
def test_every_definition_uses_known_values(self):
|
||||
for definition in catalog.all_definitions():
|
||||
assert definition.category in catalog.CATEGORIES, definition.key
|
||||
assert definition.auth_kind in catalog.AUTH_KINDS, definition.key
|
||||
assert definition.publisher in catalog.PUBLISHERS, definition.key
|
||||
assert set(definition.capabilities) <= {"sync", "read", "write"}, definition.key
|
||||
|
||||
def test_tool_connectors_match_tool_config_requirements(self):
|
||||
"""Credential fields must be the keys the tool reads its secrets from."""
|
||||
import docsgpt.api.user # noqa: F401 (mcp_tool imports it; load it first)
|
||||
from docsgpt.agents.tools.tool_manager import ToolManager
|
||||
|
||||
tools = ToolManager(config={}).tools
|
||||
for definition in catalog.all_definitions():
|
||||
if definition.publisher != "built_in":
|
||||
continue
|
||||
for tool_name in definition.tool_templates:
|
||||
requirements = tools[tool_name].get_config_requirements()
|
||||
secret_keys = {k for k, spec in requirements.items() if spec.get("secret")}
|
||||
field_keys = {f.key for f in definition.credential_fields}
|
||||
assert secret_keys <= field_keys, definition.key
|
||||
|
||||
def test_to_dict_has_no_server_secrets(self):
|
||||
payload = catalog.get_definition("google_drive").to_dict()
|
||||
assert "required_settings" not in payload
|
||||
assert payload["capabilities"] == ["sync"]
|
||||
|
||||
|
||||
class TestAvailability:
|
||||
def test_missing_settings_listed(self, monkeypatch):
|
||||
from docsgpt.core.settings import settings
|
||||
|
||||
monkeypatch.setattr(settings, "GOOGLE_CLIENT_ID", None)
|
||||
monkeypatch.setattr(settings, "GOOGLE_CLIENT_SECRET", "secret")
|
||||
definition = catalog.get_definition("google_drive")
|
||||
assert definition.missing_settings == ["GOOGLE_CLIENT_ID"]
|
||||
assert not definition.configured
|
||||
|
||||
def test_configured_when_all_set(self, monkeypatch):
|
||||
from docsgpt.core.settings import settings
|
||||
|
||||
monkeypatch.setattr(settings, "GOOGLE_CLIENT_ID", "id")
|
||||
monkeypatch.setattr(settings, "GOOGLE_CLIENT_SECRET", "secret")
|
||||
assert catalog.get_definition("google_drive").configured
|
||||
|
||||
def test_api_key_connectors_need_no_settings(self):
|
||||
assert catalog.get_definition("telegram").configured
|
||||
|
||||
|
||||
class TestRowMapping:
|
||||
def test_oauth_provider_maps_to_key(self):
|
||||
assert catalog.connector_key_for_row({"provider": "share_point"}) == "share_point"
|
||||
|
||||
def test_legacy_mcp_row_is_custom(self):
|
||||
row = {"provider": "mcp:https://mcp.unknown.dev", "server_url": "https://mcp.unknown.dev"}
|
||||
assert catalog.connector_key_for_row(row) == "custom_mcp"
|
||||
|
||||
def test_stored_key_wins(self):
|
||||
assert catalog.connector_key_for_row({"provider": "telegram", "connector_key": "telegram"}) == "telegram"
|
||||
|
||||
def test_unknown_provider(self):
|
||||
assert catalog.connector_key_for_row({"provider": "nope"}) is None
|
||||
|
||||
def test_definition_for_tool(self):
|
||||
assert catalog.definition_for_tool("telegram").key == "telegram"
|
||||
assert catalog.definition_for_tool("memory") is None
|
||||
assert catalog.definition_for_tool("mcp_tool") is None
|
||||
|
||||
|
||||
class TestPresets:
|
||||
def test_presets_load_from_yaml(self, tmp_path, monkeypatch):
|
||||
presets = tmp_path / "mcp.yaml"
|
||||
presets.write_text(
|
||||
"- key: mcp:example\n"
|
||||
" name: Example\n"
|
||||
" description: Example records.\n"
|
||||
" icon: example\n"
|
||||
" category: knowledge\n"
|
||||
" mcp_url: https://mcp.example.com/mcp\n"
|
||||
" auth_kind: mcp_oauth\n"
|
||||
)
|
||||
monkeypatch.setattr(catalog, "_PRESETS_FILE", presets)
|
||||
definition = catalog.get_definition("mcp:example")
|
||||
assert definition.publisher == "preset"
|
||||
assert definition.mcp_base_url == "https://mcp.example.com"
|
||||
assert catalog.preset_for_url("https://mcp.example.com/other").key == "mcp:example"
|
||||
row = {"provider": "mcp:https://mcp.example.com", "server_url": "https://mcp.example.com"}
|
||||
assert catalog.connector_key_for_row(row) == "mcp:example"
|
||||
|
||||
def test_base_url(self):
|
||||
assert catalog.base_url("https://a.example.com:8443/x/y") == "https://a.example.com:8443"
|
||||
assert catalog.base_url("not a url") == ""
|
||||
@@ -0,0 +1,47 @@
|
||||
"""Tests for read / write classification and permissions."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from docsgpt.connectors import permissions as p
|
||||
|
||||
|
||||
class TestActionAccess:
|
||||
def test_explicit_access_wins(self):
|
||||
assert p.action_access("telegram", {"name": "search", "access": "write"}) == "write"
|
||||
|
||||
def test_mcp_annotations(self):
|
||||
assert p.action_access("mcp_tool", {"name": "create_page", "annotations": {"readOnlyHint": True}}) == "read"
|
||||
assert p.action_access("mcp_tool", {"name": "search", "annotations": {"destructiveHint": True}}) == "write"
|
||||
|
||||
def test_api_tool_method(self):
|
||||
assert p.action_access("api_tool", {"name": "x", "method": "get"}) == "read"
|
||||
assert p.action_access("api_tool", {"name": "x", "method": "POST"}) == "write"
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"name,expected",
|
||||
[("search_pages", "read"), ("list_issues", "read"), ("create_issue", "write"), ("send_message", "write")],
|
||||
)
|
||||
def test_name_heuristic(self, name, expected):
|
||||
assert p.action_access("mcp_tool", {"name": name}) == expected
|
||||
|
||||
|
||||
class TestPermissions:
|
||||
def test_permission_from_flags(self):
|
||||
assert p.action_permission({"active": False, "require_approval": True}) == "off"
|
||||
assert p.action_permission({"active": True, "require_approval": True}) == "ask"
|
||||
assert p.action_permission({"active": True}) == "always"
|
||||
|
||||
def test_apply_permission(self):
|
||||
assert p.apply_permission({"name": "a"}, "ask") == {"name": "a", "active": True, "require_approval": True}
|
||||
assert p.apply_permission({"name": "a"}, "off")["active"] is False
|
||||
with pytest.raises(ValueError):
|
||||
p.apply_permission({}, "sometimes")
|
||||
|
||||
def test_defaults_writes_to_approval(self):
|
||||
stamped = p.apply_default_permissions(
|
||||
"mcp_tool", [{"name": "search"}, {"name": "delete_page"}]
|
||||
)
|
||||
assert stamped[0] == {"name": "search", "access": "read"}
|
||||
assert stamped[1]["access"] == "write" and stamped[1]["require_approval"] is True
|
||||
@@ -0,0 +1,263 @@
|
||||
"""Tests for the read side of the connection service and its API."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from contextlib import contextmanager
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
from flask import Flask
|
||||
from sqlalchemy import text
|
||||
|
||||
from docsgpt.connectors import service
|
||||
|
||||
|
||||
def _session(conn, user="alice", provider="google_drive", **cols) -> str:
|
||||
cols = {"status": "authorized", "user_email": f"{user}@example.com", **cols}
|
||||
casts = {"token_info": "jsonb", "session_data": "jsonb"}
|
||||
names = ", ".join(["user_id", "provider", *cols])
|
||||
values = ", ".join(
|
||||
[":user_id", ":provider", *[f"CAST(:{k} AS {casts[k]})" if k in casts else f":{k}" for k in cols]]
|
||||
)
|
||||
return str(
|
||||
conn.execute(
|
||||
text(f"INSERT INTO connector_sessions ({names}) VALUES ({values}) RETURNING id"),
|
||||
{"user_id": user, "provider": provider, **cols},
|
||||
).scalar()
|
||||
)
|
||||
|
||||
|
||||
def _source(conn, connection_id, user="alice", name="Handbook") -> str:
|
||||
return str(
|
||||
conn.execute(
|
||||
text(
|
||||
"INSERT INTO sources (user_id, name, type, sync_frequency, connection_id) "
|
||||
"VALUES (:u, :n, 'connector:file', 'weekly', CAST(:c AS uuid)) RETURNING id"
|
||||
),
|
||||
{"u": user, "n": name, "c": connection_id},
|
||||
).scalar()
|
||||
)
|
||||
|
||||
|
||||
def _tool(conn, connection_id, user="alice") -> str:
|
||||
actions = [
|
||||
{"name": "search_pages", "description": "Search", "active": True},
|
||||
{"name": "create_page", "description": "Create", "active": True, "require_approval": True},
|
||||
]
|
||||
return str(
|
||||
conn.execute(
|
||||
text(
|
||||
"INSERT INTO user_tools (user_id, name, display_name, actions, connection_id) "
|
||||
"VALUES (:u, 'mcp_tool', 'Notion', CAST(:a AS jsonb), CAST(:c AS uuid)) RETURNING id"
|
||||
),
|
||||
{"u": user, "a": json.dumps(actions), "c": connection_id},
|
||||
).scalar()
|
||||
)
|
||||
|
||||
|
||||
class TestNormalizeStatus:
|
||||
@pytest.mark.parametrize(
|
||||
"row,expected",
|
||||
[
|
||||
({"status": "authorized"}, "connected"),
|
||||
({"status": "reconnect_needed"}, "reconnect_needed"),
|
||||
({"status": "pending"}, "pending"),
|
||||
({"status": "pending", "token_info": {"access_token": "x"}}, "connected"),
|
||||
({"status": None, "session_data": {"tokens": {"access_token": "x"}}}, "connected"),
|
||||
({"status": None, "session_data": {"client_info": {}}}, "pending"),
|
||||
({"status": None, "encrypted_credentials": "v2:..."}, "connected"),
|
||||
],
|
||||
)
|
||||
def test_statuses(self, row, expected):
|
||||
assert service.normalize_status(row) == expected
|
||||
|
||||
def test_worst_status(self):
|
||||
assert service.worst_status(["connected", "reconnect_needed"]) == "reconnect_needed"
|
||||
assert service.worst_status([]) is None
|
||||
|
||||
|
||||
class TestListing:
|
||||
def test_lists_finished_connections_with_counts(self, pg_conn):
|
||||
drive = _session(pg_conn)
|
||||
_session(pg_conn, provider="confluence", status="pending", user_email=None)
|
||||
_source(pg_conn, drive)
|
||||
_source(pg_conn, drive, name="Wiki")
|
||||
connections = service.list_connections(pg_conn, "alice")
|
||||
assert [c["connector_key"] for c in connections] == ["google_drive"]
|
||||
only = connections[0]
|
||||
assert only["name"] == "Google Drive"
|
||||
assert only["account_label"] == "alice@example.com"
|
||||
assert only["status"] == "connected"
|
||||
assert only["source_count"] == 2
|
||||
assert "token_info" not in only and "session_token" not in only
|
||||
|
||||
def test_other_users_rows_hidden(self, pg_conn):
|
||||
_session(pg_conn, user="bob")
|
||||
assert service.list_connections(pg_conn, "alice") == []
|
||||
|
||||
def test_detail_lists_sources_and_tools(self, pg_conn):
|
||||
from docsgpt.storage.db.repositories.connector_sessions import ConnectorSessionsRepository
|
||||
|
||||
drive = _session(pg_conn, status="reconnect_needed")
|
||||
_source(pg_conn, drive)
|
||||
_tool(pg_conn, drive)
|
||||
row = ConnectorSessionsRepository(pg_conn).get(drive)
|
||||
detail = service.connection_detail(pg_conn, row)
|
||||
assert detail["sources"][0]["sync_state"] == "paused_reconnect"
|
||||
actions = {a["name"]: a for a in detail["tools"][0]["actions"]}
|
||||
assert actions["search_pages"] == {
|
||||
"name": "search_pages", "description": "Search", "access": "read", "permission": "always",
|
||||
}
|
||||
assert actions["create_page"]["access"] == "write"
|
||||
assert actions["create_page"]["permission"] == "ask"
|
||||
|
||||
|
||||
class TestCatalogForUser:
|
||||
def test_states(self, pg_conn, monkeypatch):
|
||||
from docsgpt.core.settings import settings
|
||||
|
||||
monkeypatch.setattr(settings, "GOOGLE_CLIENT_ID", "id")
|
||||
monkeypatch.setattr(settings, "GOOGLE_CLIENT_SECRET", "secret")
|
||||
monkeypatch.setattr(settings, "MICROSOFT_CLIENT_ID", None)
|
||||
_session(pg_conn)
|
||||
_session(pg_conn, provider="confluence", status="reconnect_needed")
|
||||
entries = {e["key"]: e for e in service.catalog_for_user(pg_conn, "alice", is_admin=False)}
|
||||
assert entries["google_drive"]["state"] == "connected"
|
||||
assert entries["google_drive"]["connected_count"] == 1
|
||||
assert entries["confluence"]["state"] == "reconnect"
|
||||
assert entries["share_point"]["state"] == "needs_setup"
|
||||
assert entries["share_point"]["missing_settings"] == []
|
||||
assert entries["telegram"]["state"] == "available"
|
||||
assert entries["custom_mcp"]["state"] == "custom"
|
||||
|
||||
def test_admin_sees_missing_setting_names(self, pg_conn, monkeypatch):
|
||||
from docsgpt.core.settings import settings
|
||||
|
||||
monkeypatch.setattr(settings, "MICROSOFT_CLIENT_ID", None)
|
||||
entries = {e["key"]: e for e in service.catalog_for_user(pg_conn, "alice", is_admin=True)}
|
||||
assert "MICROSOFT_CLIENT_ID" in entries["share_point"]["missing_settings"]
|
||||
|
||||
def test_policy_disables(self, pg_conn):
|
||||
entries = {
|
||||
e["key"]: e
|
||||
for e in service.catalog_for_user(
|
||||
pg_conn, "alice", is_admin=False, policies={"telegram": {"enabled": False}}
|
||||
)
|
||||
}
|
||||
assert entries["telegram"]["state"] == "disabled"
|
||||
assert entries["telegram"]["available"] is False
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _patched_db(conn):
|
||||
@contextmanager
|
||||
def _yield():
|
||||
yield conn
|
||||
|
||||
with patch("docsgpt.api.connector.connections.db_readonly", _yield):
|
||||
yield
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def app():
|
||||
return Flask(__name__)
|
||||
|
||||
|
||||
def _call(app, resource, path, token, *args):
|
||||
with app.test_request_context(path):
|
||||
from flask import request
|
||||
|
||||
request.decoded_token = token
|
||||
return resource().get(*args)
|
||||
|
||||
|
||||
class TestRoutes:
|
||||
def test_catalog_requires_auth(self, app):
|
||||
from docsgpt.api.connector.connections import ConnectorCatalog
|
||||
|
||||
assert _call(app, ConnectorCatalog, "/api/connectors/catalog", None).status_code == 401
|
||||
|
||||
def test_catalog(self, app, pg_conn):
|
||||
from docsgpt.api.connector.connections import ConnectorCatalog
|
||||
|
||||
with _patched_db(pg_conn):
|
||||
resp = _call(app, ConnectorCatalog, "/api/connectors/catalog", {"sub": "alice"})
|
||||
assert resp.status_code == 200
|
||||
keys = {c["key"] for c in resp.get_json()["connectors"]}
|
||||
assert "telegram" in keys
|
||||
|
||||
def test_list_and_detail(self, app, pg_conn):
|
||||
from docsgpt.api.connector.connections import ConnectionDetail, ConnectionsList
|
||||
|
||||
drive = _session(pg_conn)
|
||||
with _patched_db(pg_conn):
|
||||
listed = _call(app, ConnectionsList, "/api/connections", {"sub": "alice"})
|
||||
detail = _call(app, ConnectionDetail, f"/api/connections/{drive}", {"sub": "alice"}, drive)
|
||||
assert listed.get_json()["connections"][0]["id"] == drive
|
||||
assert detail.get_json()["connection"]["id"] == drive
|
||||
|
||||
def test_detail_of_another_users_connection_is_404(self, app, pg_conn):
|
||||
from docsgpt.api.connector.connections import ConnectionDetail
|
||||
|
||||
drive = _session(pg_conn, user="bob")
|
||||
with _patched_db(pg_conn):
|
||||
resp = _call(app, ConnectionDetail, f"/api/connections/{drive}", {"sub": "alice"}, drive)
|
||||
assert resp.status_code == 404
|
||||
|
||||
def test_detail_of_bad_id_is_404(self, app, pg_conn):
|
||||
from docsgpt.api.connector.connections import ConnectionDetail
|
||||
|
||||
with _patched_db(pg_conn):
|
||||
resp = _call(app, ConnectionDetail, "/api/connections/nope", {"sub": "alice"}, "nope")
|
||||
assert resp.status_code == 404
|
||||
|
||||
|
||||
class TestDisconnect:
|
||||
def test_clears_credentials_keeps_resources(self, pg_conn):
|
||||
from docsgpt.storage.db.repositories.connector_sessions import ConnectorSessionsRepository
|
||||
|
||||
drive = _session(
|
||||
pg_conn, session_token="tok", token_info=json.dumps({"access_token": "at", "refresh_token": "rt"}),
|
||||
)
|
||||
source = _source(pg_conn, drive)
|
||||
repo = ConnectorSessionsRepository(pg_conn)
|
||||
result = service.disconnect(pg_conn, repo.get(drive))
|
||||
assert result["status"] == "disconnected"
|
||||
row = repo.get(drive)
|
||||
assert row["token_info"] is None and row["session_token"] is None
|
||||
linked = pg_conn.execute(
|
||||
text("SELECT connection_id FROM sources WHERE id = CAST(:id AS uuid)"), {"id": source}
|
||||
).scalar()
|
||||
assert str(linked) == drive
|
||||
|
||||
def test_mcp_keeps_client_registration(self, pg_conn):
|
||||
from docsgpt.storage.db.repositories.connector_sessions import ConnectorSessionsRepository
|
||||
|
||||
mcp = _session(
|
||||
pg_conn,
|
||||
provider="mcp:https://mcp.example.com",
|
||||
status=None,
|
||||
session_data=json.dumps({"tokens": {"access_token": "x"}, "client_info": {"client_id": "c"}}),
|
||||
)
|
||||
repo = ConnectorSessionsRepository(pg_conn)
|
||||
service.disconnect(pg_conn, repo.get(mcp))
|
||||
assert repo.get(mcp)["session_data"] == {"client_info": {"client_id": "c"}}
|
||||
|
||||
def test_route_rejects_other_users(self, app, pg_conn):
|
||||
from docsgpt.api.connector.connections import ConnectionDisconnect
|
||||
|
||||
drive = _session(pg_conn, user="bob")
|
||||
|
||||
@contextmanager
|
||||
def _yield():
|
||||
yield pg_conn
|
||||
|
||||
with patch("docsgpt.api.connector.connections.db_session", _yield), app.test_request_context(
|
||||
f"/api/connections/{drive}/disconnect", method="POST"
|
||||
):
|
||||
from flask import request
|
||||
|
||||
request.decoded_token = {"sub": "alice"}
|
||||
resp = ConnectionDisconnect().post(drive)
|
||||
assert resp.status_code == 404
|
||||
@@ -0,0 +1,154 @@
|
||||
"""Migration tests for 0038_connections: columns, backfill and round trip."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
from sqlalchemy import text
|
||||
|
||||
|
||||
pytestmark = pytest.mark.integration
|
||||
|
||||
_0037 = "0037_request_traces"
|
||||
|
||||
|
||||
def _alembic_ini() -> Path:
|
||||
return Path(__file__).resolve().parents[3] / "docsgpt" / "alembic.ini"
|
||||
|
||||
|
||||
def _run_alembic(url: str, *args: str) -> None:
|
||||
subprocess.check_call(
|
||||
[sys.executable, "-m", "alembic", "-c", str(_alembic_ini()), *args],
|
||||
timeout=120,
|
||||
env={**os.environ, "POSTGRES_URI": url},
|
||||
)
|
||||
|
||||
|
||||
def _columns(conn, table: str) -> set[str]:
|
||||
rows = conn.execute(
|
||||
text("SELECT column_name FROM information_schema.columns WHERE table_name = :t"),
|
||||
{"t": table},
|
||||
).fetchall()
|
||||
return {r[0] for r in rows}
|
||||
|
||||
|
||||
def _seed_legacy(conn) -> dict:
|
||||
"""Rows as a pre-0038 install left them."""
|
||||
ids = {}
|
||||
ids["drive"] = conn.execute(
|
||||
text(
|
||||
"INSERT INTO connector_sessions (user_id, provider, session_token, user_email, status, token_info) "
|
||||
"VALUES ('alice', 'google_drive', 'tok-a', 'alice@example.com', 'authorized', "
|
||||
"CAST(:ti AS jsonb)) RETURNING id"
|
||||
),
|
||||
{"ti": json.dumps({"access_token": "at", "refresh_token": "rt"})},
|
||||
).scalar()
|
||||
ids["mcp"] = conn.execute(
|
||||
text(
|
||||
"INSERT INTO connector_sessions (user_id, provider, server_url, session_data) "
|
||||
"VALUES ('alice', 'mcp:https://mcp.example.com', 'https://mcp.example.com', CAST(:sd AS jsonb)) "
|
||||
"RETURNING id"
|
||||
),
|
||||
{"sd": json.dumps({"tokens": {"access_token": "m"}})},
|
||||
).scalar()
|
||||
ids["source"] = conn.execute(
|
||||
text(
|
||||
"INSERT INTO sources (user_id, name, type, remote_data) "
|
||||
"VALUES ('alice', 'Handbook', 'connector:file', CAST(:rd AS jsonb)) RETURNING id"
|
||||
),
|
||||
{"rd": json.dumps({"provider": "google_drive", "file_ids": ["f1"]})},
|
||||
).scalar()
|
||||
ids["other_source"] = conn.execute(
|
||||
text(
|
||||
"INSERT INTO sources (user_id, name, type, remote_data) "
|
||||
"VALUES ('bob', 'Bob files', 'connector:file', CAST(:rd AS jsonb)) RETURNING id"
|
||||
),
|
||||
{"rd": json.dumps({"provider": "google_drive"})},
|
||||
).scalar()
|
||||
ids["mcp_tool"] = conn.execute(
|
||||
text(
|
||||
"INSERT INTO user_tools (user_id, name, config) "
|
||||
"VALUES ('alice', 'mcp_tool', CAST(:c AS jsonb)) RETURNING id"
|
||||
),
|
||||
{"c": json.dumps({"server_url": "https://mcp.example.com/mcp", "auth_type": "oauth"})},
|
||||
).scalar()
|
||||
ids["bearer_tool"] = conn.execute(
|
||||
text(
|
||||
"INSERT INTO user_tools (user_id, name, config) "
|
||||
"VALUES ('alice', 'mcp_tool', CAST(:c AS jsonb)) RETURNING id"
|
||||
),
|
||||
{"c": json.dumps({"server_url": "https://mcp.example.com/mcp", "auth_type": "bearer"})},
|
||||
).scalar()
|
||||
return ids
|
||||
|
||||
|
||||
class TestMigration0038:
|
||||
def test_head_has_connection_columns(self, pg_engine):
|
||||
with pg_engine.connect() as conn:
|
||||
assert {"connector_key", "display_name", "account_label", "auth_kind", "updated_at"} <= _columns(
|
||||
conn, "connector_sessions"
|
||||
)
|
||||
assert "connection_id" in _columns(conn, "sources")
|
||||
assert "connection_id" in _columns(conn, "user_tools")
|
||||
|
||||
def test_backfill_links_legacy_rows(self, pg_engine):
|
||||
url = pg_engine.url.render_as_string(hide_password=False)
|
||||
_run_alembic(url, "downgrade", _0037)
|
||||
with pg_engine.begin() as conn:
|
||||
ids = _seed_legacy(conn)
|
||||
_run_alembic(url, "upgrade", "head")
|
||||
|
||||
with pg_engine.connect() as conn:
|
||||
drive = conn.execute(
|
||||
text(
|
||||
"SELECT connector_key, auth_kind, display_name, account_label "
|
||||
"FROM connector_sessions WHERE id = :id"
|
||||
),
|
||||
{"id": ids["drive"]},
|
||||
).one()
|
||||
assert tuple(drive) == ("google_drive", "oauth", "Google Drive", "alice@example.com")
|
||||
|
||||
mcp = conn.execute(
|
||||
text("SELECT connector_key, auth_kind, display_name FROM connector_sessions WHERE id = :id"),
|
||||
{"id": ids["mcp"]},
|
||||
).one()
|
||||
assert tuple(mcp) == ("custom_mcp", "mcp_oauth", "mcp.example.com")
|
||||
|
||||
linked = conn.execute(
|
||||
text("SELECT connection_id FROM sources WHERE id = :id"), {"id": ids["source"]}
|
||||
).scalar()
|
||||
assert linked == ids["drive"]
|
||||
# Bob has no session for the provider: nothing to link.
|
||||
assert (
|
||||
conn.execute(
|
||||
text("SELECT connection_id FROM sources WHERE id = :id"), {"id": ids["other_source"]}
|
||||
).scalar()
|
||||
is None
|
||||
)
|
||||
tool = conn.execute(
|
||||
text("SELECT connection_id FROM user_tools WHERE id = :id"), {"id": ids["mcp_tool"]}
|
||||
).scalar()
|
||||
assert tool == ids["mcp"]
|
||||
# A bearer MCP tool has no OAuth session to point at.
|
||||
bearer = conn.execute(
|
||||
text("SELECT connection_id FROM user_tools WHERE id = :id"), {"id": ids["bearer_tool"]}
|
||||
).scalar()
|
||||
assert bearer is None
|
||||
|
||||
def test_downgrade_then_upgrade_is_idempotent(self, pg_engine):
|
||||
url = pg_engine.url.render_as_string(hide_password=False)
|
||||
_run_alembic(url, "downgrade", _0037)
|
||||
with pg_engine.connect() as conn:
|
||||
assert "connector_key" not in _columns(conn, "connector_sessions")
|
||||
assert "connection_id" not in _columns(conn, "sources")
|
||||
_run_alembic(url, "upgrade", "head")
|
||||
_run_alembic(url, "downgrade", _0037)
|
||||
_run_alembic(url, "upgrade", "head")
|
||||
with pg_engine.connect() as conn:
|
||||
assert "connector_key" in _columns(conn, "connector_sessions")
|
||||
assert "connection_id" in _columns(conn, "user_tools")
|
||||
Reference in new issue
Block a user