mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-04 22:13:08 +00:00
Check a tool's connection before writing its permissions
This commit is contained in:
1 parent
33f29cb408
commit
a260ca44f1
3 files changed
+34
-5
No files matched your search
@@ -426,8 +426,8 @@ class ConnectionToolPermissions(Resource):
|
||||
row = _owned(conn, connection_id, user_id)
|
||||
if row is None:
|
||||
return _not_found()
|
||||
tool = service.set_tool_permissions(conn, user_id, tool_id, permissions)
|
||||
if tool is None or str(tool.get("connection_id")) != connection_id:
|
||||
tool = service.set_tool_permissions(conn, user_id, connection_id, tool_id, permissions)
|
||||
if tool is None:
|
||||
return _not_found()
|
||||
payload = service.serialize_tool(tool)
|
||||
return make_response(jsonify({"success": True, "tool": payload}), 200)
|
||||
|
||||
@@ -1154,14 +1154,28 @@ def remove_connection(conn, row: dict, *, sources: str = "keep", tools: str = "d
|
||||
return linked_sources if sources == "delete" else []
|
||||
|
||||
|
||||
def set_tool_permissions(conn, user_id: str, tool_id: str, permissions: dict) -> Optional[dict]:
|
||||
"""Apply ``{action: always | ask | off}`` to a tool the user owns."""
|
||||
def set_tool_permissions(
|
||||
conn, user_id: str, connection_id: str, tool_id: str, permissions: dict,
|
||||
) -> Optional[dict]:
|
||||
"""Apply ``{action: always | ask | off}`` to a tool the user owns.
|
||||
|
||||
Args:
|
||||
conn: Open database connection.
|
||||
user_id: The caller, who must own the tool.
|
||||
connection_id: The connection the tool must belong to.
|
||||
tool_id: The tool to update.
|
||||
permissions: Action name to ``always``, ``ask`` or ``off``.
|
||||
|
||||
Returns:
|
||||
The updated tool, or None (nothing written) when the tool is not the
|
||||
user's or belongs to another connection.
|
||||
"""
|
||||
from docsgpt.connectors.permissions import apply_permission
|
||||
from docsgpt.storage.db.repositories.user_tools import UserToolsRepository
|
||||
|
||||
tools = UserToolsRepository(conn)
|
||||
tool = tools.get_any(tool_id, user_id)
|
||||
if tool is None or tool.get("user_id") != user_id:
|
||||
if tool is None or tool.get("user_id") != user_id or str(tool.get("connection_id")) != connection_id:
|
||||
return None
|
||||
actions = [
|
||||
apply_permission(action, permissions[action.get("name")]) if action.get("name") in permissions else action
|
||||
|
||||
@@ -128,6 +128,21 @@ class TestToolPermissions:
|
||||
body={"permissions": {"telegram_send_message": "off"}}, args=(cid, tool_id))
|
||||
assert resp.status_code == 404
|
||||
|
||||
def test_tool_of_another_connection_is_left_unchanged(self, app, pg_conn): # noqa: F811
|
||||
from sqlalchemy import text
|
||||
|
||||
from docsgpt.api.connector.connections import ConnectionToolPermissions
|
||||
|
||||
cid = _connection(pg_conn, secrets={"credentials": {"token": "t"}})
|
||||
other = _connection(pg_conn, account_label="…zzzz", secrets={"credentials": {"token": "u"}})
|
||||
tool_id = self._tool(pg_conn, other)
|
||||
with _db(pg_conn):
|
||||
_call(app, ConnectionToolPermissions, "put", "/p",
|
||||
body={"permissions": {"telegram_send_message": "off"}}, args=(cid, tool_id))
|
||||
actions = pg_conn.execute(text("SELECT actions FROM user_tools WHERE id = CAST(:t AS uuid)"),
|
||||
{"t": tool_id}).scalar()
|
||||
assert actions[0]["active"] is True
|
||||
|
||||
|
||||
class TestRefreshTools:
|
||||
def test_returns_the_diff(self, app, pg_conn): # noqa: F811
|
||||
|
||||
Reference in new issue
Block a user