mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-12 16:15:14 +00:00
Stop syncing sources of a connector an admin turned off
Scheduled syncs checked the connection's status but not the admin policy, so a disabled connector kept pulling data. Connector sources now return "disabled" and S3/Reddit syncs get no credentials until the connector is turned back on.
This commit is contained in:
1 parent
2411118048
commit
d9f823d601
3 files changed
+61
-4
No files matched your search
@@ -1200,6 +1200,20 @@ def load_policies(conn) -> dict[str, dict]:
|
||||
return policies
|
||||
|
||||
|
||||
def connector_enabled(conn, row: dict) -> bool:
|
||||
"""Whether an admin left on the connector a connection row belongs to.
|
||||
|
||||
Args:
|
||||
conn: Open database connection.
|
||||
row: A ``connector_sessions`` row.
|
||||
|
||||
Returns:
|
||||
False when the connector (or custom MCP servers) is turned off.
|
||||
"""
|
||||
key = catalog.connector_key_for_row(row)
|
||||
return (load_policies(conn).get(key) or {}).get("enabled") is not False
|
||||
|
||||
|
||||
def ensure_connector_allowed(conn, connector_key: Optional[str]) -> None:
|
||||
"""Refuse a new connection to a connector an admin turned off.
|
||||
|
||||
|
||||
+9
-4
@@ -2186,14 +2186,15 @@ def _with_connection_credentials(source_data, connection_id: str):
|
||||
|
||||
S3 and Reddit sources made from a connection keep their keys on the
|
||||
connection only, never in ``sources.remote_data``. Returns None when the
|
||||
connection is gone or needs reconnecting.
|
||||
connection is gone, needs reconnecting, or its connector is turned off.
|
||||
"""
|
||||
from docsgpt.connectors import service
|
||||
from docsgpt.storage.db.repositories.connector_sessions import ConnectorSessionsRepository
|
||||
|
||||
with db_readonly() as conn:
|
||||
row = ConnectorSessionsRepository(conn).get(str(connection_id))
|
||||
if row is None:
|
||||
enabled = row is not None and service.connector_enabled(conn, row)
|
||||
if row is None or not enabled:
|
||||
return None
|
||||
try:
|
||||
credentials = service.get_credentials(row)
|
||||
@@ -2235,9 +2236,10 @@ def sync_connector_source(self, source_id: str) -> Dict[str, Any]:
|
||||
source_id: The source to sync.
|
||||
|
||||
Returns:
|
||||
``{"status": "success" | "paused" | "skipped"}`` plus the ingest result.
|
||||
``{"status": "success" | "paused" | "disabled" | "skipped"}`` plus
|
||||
the ingest result.
|
||||
"""
|
||||
from docsgpt.connectors.service import ConnectionUnavailable, normalize_status
|
||||
from docsgpt.connectors.service import ConnectionUnavailable, connector_enabled, normalize_status
|
||||
from docsgpt.storage.db.repositories.connector_sessions import ConnectorSessionsRepository
|
||||
|
||||
with db_readonly() as conn:
|
||||
@@ -2256,8 +2258,11 @@ def sync_connector_source(self, source_id: str) -> Dict[str, Any]:
|
||||
if source and source.get("connection_id")
|
||||
else None
|
||||
)
|
||||
enabled = connection is not None and connector_enabled(conn, connection)
|
||||
if not source or not connection:
|
||||
return {"status": "skipped"}
|
||||
if not enabled:
|
||||
return {"status": "disabled"}
|
||||
if normalize_status(connection) != "connected":
|
||||
return {"status": "paused"}
|
||||
remote_data = source.get("remote_data") or {}
|
||||
|
||||
@@ -219,6 +219,44 @@ class TestScheduledSync:
|
||||
assert result == {"status": "paused"}
|
||||
ingest.assert_not_called()
|
||||
|
||||
def test_disabled_connector_is_not_synced(self, pg_conn):
|
||||
from docsgpt import worker
|
||||
from docsgpt.storage.db.repositories.connector_policies import ConnectorPoliciesRepository
|
||||
|
||||
cid = _connection(pg_conn, provider="google_drive", auth_kind="oauth",
|
||||
secrets={"token_info": {"access_token": "a"}})
|
||||
source = pg_conn.execute(text(
|
||||
"INSERT INTO sources (user_id, name, type, sync_frequency, connection_id, remote_data) "
|
||||
"VALUES ('alice', 'Drive', 'connector:file', 'weekly', CAST(:c AS uuid), "
|
||||
"'{\"provider\": \"google_drive\"}') RETURNING id"
|
||||
), {"c": cid}).scalar()
|
||||
ConnectorPoliciesRepository(pg_conn).upsert("google_drive", enabled=False)
|
||||
|
||||
@contextmanager
|
||||
def _yield():
|
||||
yield pg_conn
|
||||
|
||||
with patch.object(worker, "db_readonly", _yield), patch.object(worker, "ingest_connector") as ingest:
|
||||
result = worker.sync_connector_source(MagicMock(), str(source))
|
||||
assert result == {"status": "disabled"}
|
||||
ingest.assert_not_called()
|
||||
|
||||
def test_disabled_connector_gives_remote_sync_no_credentials(self, pg_conn):
|
||||
from docsgpt import worker
|
||||
from docsgpt.storage.db.repositories.connector_policies import ConnectorPoliciesRepository
|
||||
|
||||
cid = _connection(pg_conn, provider="s3", auth_kind="api_key",
|
||||
secrets={"credentials": {"aws_access_key_id": "AKIA", "aws_secret_access_key": "s"}})
|
||||
|
||||
@contextmanager
|
||||
def _yield():
|
||||
yield pg_conn
|
||||
|
||||
with patch.object(worker, "db_readonly", _yield):
|
||||
assert worker._with_connection_credentials({"bucket": "b"}, cid)["aws_access_key_id"] == "AKIA"
|
||||
ConnectorPoliciesRepository(pg_conn).upsert("s3", enabled=False)
|
||||
assert worker._with_connection_credentials({"bucket": "b"}, cid) is None
|
||||
|
||||
def test_sync_runs_as_the_connection_without_a_browser(self, pg_conn):
|
||||
from docsgpt import worker
|
||||
|
||||
|
||||
Reference in new issue
Block a user