From cb30a24e05742cf55c6bca6ce040abd59595a9dc Mon Sep 17 00:00:00 2001 From: Alex Date: Sun, 12 Apr 2026 13:51:29 +0100 Subject: [PATCH] feat: fixes on pg2 --- .../versions/0002_add_unique_constraints.py | 57 +++++++++++ application/storage/db/models.py | 4 + .../storage/db/repositories/agent_folders.py | 43 ++++++-- application/storage/db/repositories/agents.py | 99 ++++++++----------- .../storage/db/repositories/attachments.py | 8 +- .../db/repositories/connector_sessions.py | 16 +-- application/storage/db/repositories/notes.py | 22 ++--- .../storage/db/repositories/prompts.py | 29 +++--- .../storage/db/repositories/sources.py | 44 ++++----- application/storage/db/repositories/todos.py | 36 ++++--- .../storage/db/repositories/user_tools.py | 6 +- scripts/db/backfill.py | 5 +- tests/storage/db/repositories/test_agents.py | 34 +++---- .../db/repositories/test_attachments.py | 9 +- tests/storage/db/repositories/test_notes.py | 10 +- tests/storage/db/repositories/test_prompts.py | 33 ++++--- tests/storage/db/repositories/test_sources.py | 58 ++++++----- tests/storage/db/repositories/test_todos.py | 58 ++++++++--- .../db/repositories/test_user_tools.py | 21 ++-- 19 files changed, 355 insertions(+), 237 deletions(-) create mode 100644 application/alembic/versions/0002_add_unique_constraints.py diff --git a/application/alembic/versions/0002_add_unique_constraints.py b/application/alembic/versions/0002_add_unique_constraints.py new file mode 100644 index 00000000..dd767d35 --- /dev/null +++ b/application/alembic/versions/0002_add_unique_constraints.py @@ -0,0 +1,57 @@ +"""0002 add unique constraints for notes and connector_sessions. + +The memories table already has ``memories_user_tool_path_uidx`` from the +0001 baseline. Notes and connector_sessions were missing unique constraints +that their repository upsert logic depends on. + +Before creating the indexes, duplicate rows are cleaned up — keeping only +the row with the latest ``id`` (UUID, lexicographic max) per group. + +Revision ID: 0002_add_unique_constraints +Revises: 0001_initial +Create Date: 2026-04-12 +""" + +from typing import Sequence, Union + +from alembic import op + +revision: str = "0002_add_unique_constraints" +down_revision: Union[str, None] = "0001_initial" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + # Deduplicate notes: keep one row per (user_id, tool_id) + op.execute(""" + DELETE FROM notes + WHERE id NOT IN ( + SELECT DISTINCT ON (user_id, tool_id) id + FROM notes + ORDER BY user_id, tool_id, created_at DESC + ); + """) + op.execute( + "CREATE UNIQUE INDEX IF NOT EXISTS notes_user_tool_uidx " + "ON notes (user_id, tool_id);" + ) + + # Deduplicate connector_sessions: keep one row per (user_id, provider) + op.execute(""" + DELETE FROM connector_sessions + WHERE id NOT IN ( + SELECT DISTINCT ON (user_id, provider) id + FROM connector_sessions + ORDER BY user_id, provider, created_at DESC + ); + """) + op.execute( + "CREATE UNIQUE INDEX IF NOT EXISTS connector_sessions_user_provider_uidx " + "ON connector_sessions (user_id, provider);" + ) + + +def downgrade() -> None: + op.execute("DROP INDEX IF EXISTS connector_sessions_user_provider_uidx;") + op.execute("DROP INDEX IF EXISTS notes_user_tool_uidx;") diff --git a/application/storage/db/models.py b/application/storage/db/models.py index a98a4ccf..475e153c 100644 --- a/application/storage/db/models.py +++ b/application/storage/db/models.py @@ -15,6 +15,7 @@ from sqlalchemy import ( ForeignKey, Integer, MetaData, + UniqueConstraint, Table, Text, func, @@ -190,6 +191,7 @@ memories_table = Table( Column("path", Text, nullable=False), Column("content", Text, nullable=False), Column("updated_at", DateTime(timezone=True), nullable=False, server_default=func.now()), + UniqueConstraint("user_id", "tool_id", "path", name="memories_user_tool_path_uidx"), ) todos_table = Table( @@ -214,6 +216,7 @@ notes_table = Table( Column("content", Text, nullable=False), Column("created_at", DateTime(timezone=True), nullable=False, server_default=func.now()), Column("updated_at", DateTime(timezone=True), nullable=False, server_default=func.now()), + UniqueConstraint("user_id", "tool_id", name="notes_user_tool_uidx"), ) connector_sessions_table = Table( @@ -225,4 +228,5 @@ connector_sessions_table = Table( Column("session_data", JSONB, nullable=False), Column("expires_at", DateTime(timezone=True)), Column("created_at", DateTime(timezone=True), nullable=False, server_default=func.now()), + UniqueConstraint("user_id", "provider", name="connector_sessions_user_provider_uidx"), ) diff --git a/application/storage/db/repositories/agent_folders.py b/application/storage/db/repositories/agent_folders.py index 67c62af4..46fb0e41 100644 --- a/application/storage/db/repositories/agent_folders.py +++ b/application/storage/db/repositories/agent_folders.py @@ -13,16 +13,16 @@ class AgentFoldersRepository: def __init__(self, conn: Connection) -> None: self._conn = conn - def create(self, user_id: str, name: str, *, parent_id: Optional[str] = None) -> dict: + def create(self, user_id: str, name: str, *, description: Optional[str] = None) -> dict: result = self._conn.execute( text( """ INSERT INTO agent_folders (user_id, name, description) - VALUES (:user_id, :name, :parent_id) + VALUES (:user_id, :name, :description) RETURNING * """ ), - {"user_id": user_id, "name": name, "parent_id": parent_id}, + {"user_id": user_id, "name": name, "description": description}, ) return row_to_dict(result.fetchone()) @@ -46,13 +46,38 @@ class AgentFoldersRepository: filtered = {k: v for k, v in fields.items() if k in allowed} if not filtered: return False - set_clauses = [f"{col} = :val_{col}" for col in filtered] - set_clauses.append("updated_at = now()") params: dict = {"id": folder_id, "user_id": user_id} - for col, val in filtered.items(): - params[f"val_{col}"] = val - sql = f"UPDATE agent_folders SET {', '.join(set_clauses)} WHERE id = CAST(:id AS uuid) AND user_id = :user_id" - result = self._conn.execute(text(sql), params) + if "name" in filtered and "description" in filtered: + params["name"] = filtered["name"] + params["description"] = filtered["description"] + result = self._conn.execute( + text( + "UPDATE agent_folders " + "SET name = :name, description = :description, updated_at = now() " + "WHERE id = CAST(:id AS uuid) AND user_id = :user_id" + ), + params, + ) + elif "name" in filtered: + params["name"] = filtered["name"] + result = self._conn.execute( + text( + "UPDATE agent_folders " + "SET name = :name, updated_at = now() " + "WHERE id = CAST(:id AS uuid) AND user_id = :user_id" + ), + params, + ) + else: + params["description"] = filtered["description"] + result = self._conn.execute( + text( + "UPDATE agent_folders " + "SET description = :description, updated_at = now() " + "WHERE id = CAST(:id AS uuid) AND user_id = :user_id" + ), + params, + ) return result.rowcount > 0 def delete(self, folder_id: str, user_id: str) -> bool: diff --git a/application/storage/db/repositories/agents.py b/application/storage/db/repositories/agents.py index 198c94b1..014bd07b 100644 --- a/application/storage/db/repositories/agents.py +++ b/application/storage/db/repositories/agents.py @@ -15,9 +15,11 @@ from __future__ import annotations import json from typing import Optional -from sqlalchemy import Connection, text +from sqlalchemy import Connection, func, text +from sqlalchemy.dialects.postgresql import insert as pg_insert from application.storage.db.base_repository import row_to_dict +from application.storage.db.models import agents_table class AgentsRepository: @@ -25,56 +27,36 @@ class AgentsRepository: self._conn = conn def create(self, user_id: str, name: str, status: str, **kwargs) -> dict: - cols = ["user_id", "name", "status"] - vals = [":user_id", ":name", ":status"] - params: dict = {"user_id": user_id, "name": name, "status": status} + values: dict = {"user_id": user_id, "name": name, "status": status} - _TEXT_COLS = {"description", "agent_type", "key", "retriever", - "default_model_id", "incoming_webhook_token"} - _UUID_COLS = {"source_id", "prompt_id", "folder_id"} - _INT_COLS = {"chunks", "token_limit", "request_limit"} - _BOOL_COLS = {"limited_token_mode", "limited_request_mode", "shared"} - _JSONB_COLS = {"tools", "json_schema", "models"} + _ALLOWED = { + "description", "agent_type", "key", "retriever", + "default_model_id", "incoming_webhook_token", + "source_id", "prompt_id", "folder_id", + "chunks", "token_limit", "request_limit", + "limited_token_mode", "limited_request_mode", "shared", + "tools", "json_schema", "models", + } - for col in _TEXT_COLS: - if col in kwargs and kwargs[col] is not None: - cols.append(col) - vals.append(f":{col}") - params[col] = kwargs[col] - for col in _UUID_COLS: - if col in kwargs and kwargs[col] is not None: - cols.append(col) - vals.append(f"CAST(:{col} AS uuid)") - params[col] = str(kwargs[col]) - for col in _INT_COLS: - if col in kwargs and kwargs[col] is not None: - cols.append(col) - vals.append(f":{col}") - params[col] = int(kwargs[col]) - for col in _BOOL_COLS: - if col in kwargs: - cols.append(col) - vals.append(f":{col}") - params[col] = bool(kwargs[col]) - for col in _JSONB_COLS: - if col in kwargs and kwargs[col] is not None: - cols.append(col) - vals.append(f"CAST(:{col} AS jsonb)") - params[col] = json.dumps(kwargs[col]) + for col, val in kwargs.items(): + if col not in _ALLOWED or val is None: + continue + if col in ("tools", "json_schema", "models"): + values[col] = json.dumps(val) + elif col in ("chunks", "token_limit", "request_limit"): + values[col] = int(val) + elif col in ("limited_token_mode", "limited_request_mode", "shared"): + values[col] = bool(val) + elif col in ("source_id", "prompt_id", "folder_id"): + values[col] = str(val) + else: + values[col] = val - sql = f"INSERT INTO agents ({', '.join(cols)}) VALUES ({', '.join(vals)}) RETURNING *" - result = self._conn.execute(text(sql), params) + stmt = pg_insert(agents_table).values(**values).returning(agents_table) + result = self._conn.execute(stmt) return row_to_dict(result.fetchone()) - def get(self, agent_id: str) -> Optional[dict]: - result = self._conn.execute( - text("SELECT * FROM agents WHERE id = CAST(:id AS uuid)"), - {"id": agent_id}, - ) - row = result.fetchone() - return row_to_dict(row) if row is not None else None - - def get_for_user(self, agent_id: str, user_id: str) -> Optional[dict]: + def get(self, agent_id: str, user_id: str) -> Optional[dict]: result = self._conn.execute( text("SELECT * FROM agents WHERE id = CAST(:id AS uuid) AND user_id = :user_id"), {"id": agent_id, "user_id": user_id}, @@ -123,21 +105,24 @@ class AgentsRepository: if not filtered: return False - set_clauses = [] - params: dict = {"id": agent_id, "user_id": user_id} + values: dict = {} for col, val in filtered.items(): if col in ("tools", "json_schema", "models"): - set_clauses.append(f"{col} = CAST(:val_{col} AS jsonb)") - params[f"val_{col}"] = json.dumps(val) if not isinstance(val, str) else val + values[col] = json.dumps(val) if not isinstance(val, str) else val elif col in ("source_id", "prompt_id", "folder_id"): - set_clauses.append(f"{col} = CAST(:val_{col} AS uuid)") - params[f"val_{col}"] = str(val) if val else None + values[col] = str(val) if val else None else: - set_clauses.append(f"{col} = :val_{col}") - params[f"val_{col}"] = val - set_clauses.append("updated_at = now()") - sql = f"UPDATE agents SET {', '.join(set_clauses)} WHERE id = CAST(:id AS uuid) AND user_id = :user_id" - result = self._conn.execute(text(sql), params) + values[col] = val + values["updated_at"] = func.now() + + t = agents_table + stmt = ( + t.update() + .where(t.c.id == agent_id) + .where(t.c.user_id == user_id) + .values(**values) + ) + result = self._conn.execute(stmt) return result.rowcount > 0 def delete(self, agent_id: str, user_id: str) -> bool: diff --git a/application/storage/db/repositories/attachments.py b/application/storage/db/repositories/attachments.py index bdc47229..7d139561 100644 --- a/application/storage/db/repositories/attachments.py +++ b/application/storage/db/repositories/attachments.py @@ -33,10 +33,12 @@ class AttachmentsRepository: ) return row_to_dict(result.fetchone()) - def get(self, attachment_id: str) -> Optional[dict]: + def get(self, attachment_id: str, user_id: str) -> Optional[dict]: result = self._conn.execute( - text("SELECT * FROM attachments WHERE id = CAST(:id AS uuid)"), - {"id": attachment_id}, + text( + "SELECT * FROM attachments WHERE id = CAST(:id AS uuid) AND user_id = :user_id" + ), + {"id": attachment_id, "user_id": user_id}, ) row = result.fetchone() return row_to_dict(row) if row is not None else None diff --git a/application/storage/db/repositories/connector_sessions.py b/application/storage/db/repositories/connector_sessions.py index 5bc50119..5d5934a1 100644 --- a/application/storage/db/repositories/connector_sessions.py +++ b/application/storage/db/repositories/connector_sessions.py @@ -27,7 +27,8 @@ class ConnectorSessionsRepository: """ INSERT INTO connector_sessions (user_id, provider, session_data) VALUES (:user_id, :provider, CAST(:session_data AS jsonb)) - ON CONFLICT DO NOTHING + ON CONFLICT (user_id, provider) + DO UPDATE SET session_data = EXCLUDED.session_data RETURNING * """ ), @@ -37,18 +38,7 @@ class ConnectorSessionsRepository: "session_data": json.dumps(session_data), }, ) - row = result.fetchone() - if row is not None: - return row_to_dict(row) - # Conflict — update existing row. - self._conn.execute( - text( - "UPDATE connector_sessions SET session_data = CAST(:session_data AS jsonb) " - "WHERE user_id = :user_id AND provider = :provider" - ), - {"user_id": user_id, "provider": provider, "session_data": json.dumps(session_data)}, - ) - return self.get_by_user_provider(user_id, provider) or {} + return row_to_dict(result.fetchone()) def get_by_user_provider(self, user_id: str, provider: str) -> Optional[dict]: result = self._conn.execute( diff --git a/application/storage/db/repositories/notes.py b/application/storage/db/repositories/notes.py index 09c3d9d1..8c68f375 100644 --- a/application/storage/db/repositories/notes.py +++ b/application/storage/db/repositories/notes.py @@ -25,24 +25,14 @@ class NotesRepository: """ INSERT INTO notes (user_id, tool_id, title, content) VALUES (:user_id, CAST(:tool_id AS uuid), :title, :content) - ON CONFLICT DO NOTHING + ON CONFLICT (user_id, tool_id) + DO UPDATE SET content = EXCLUDED.content, title = EXCLUDED.title, updated_at = now() RETURNING * """ ), {"user_id": user_id, "tool_id": tool_id, "title": title, "content": content}, ) - row = result.fetchone() - if row is not None: - return row_to_dict(row) - # Row already existed — update instead. - self._conn.execute( - text( - "UPDATE notes SET content = :content, updated_at = now() " - "WHERE user_id = :user_id AND tool_id = CAST(:tool_id AS uuid)" - ), - {"user_id": user_id, "tool_id": tool_id, "content": content}, - ) - return self.get_for_user_tool(user_id, tool_id) or {} + return row_to_dict(result.fetchone()) def get_for_user_tool(self, user_id: str, tool_id: str) -> Optional[dict]: result = self._conn.execute( @@ -54,10 +44,10 @@ class NotesRepository: row = result.fetchone() return row_to_dict(row) if row is not None else None - def get(self, note_id: str) -> Optional[dict]: + def get(self, note_id: str, user_id: str) -> Optional[dict]: result = self._conn.execute( - text("SELECT * FROM notes WHERE id = CAST(:id AS uuid)"), - {"id": note_id}, + text("SELECT * FROM notes WHERE id = CAST(:id AS uuid) AND user_id = :user_id"), + {"id": note_id, "user_id": user_id}, ) row = result.fetchone() return row_to_dict(row) if row is not None else None diff --git a/application/storage/db/repositories/prompts.py b/application/storage/db/repositories/prompts.py index e959f33f..12a40d0b 100644 --- a/application/storage/db/repositories/prompts.py +++ b/application/storage/db/repositories/prompts.py @@ -40,17 +40,24 @@ class PromptsRepository: ) return row_to_dict(result.fetchone()) - def get(self, prompt_id: str, user_id: Optional[str] = None) -> Optional[dict]: - if user_id is not None: - result = self._conn.execute( - text("SELECT * FROM prompts WHERE id = CAST(:id AS uuid) AND user_id = :user_id"), - {"id": prompt_id, "user_id": user_id}, - ) - else: - result = self._conn.execute( - text("SELECT * FROM prompts WHERE id = CAST(:id AS uuid)"), - {"id": prompt_id}, - ) + def get(self, prompt_id: str, user_id: str) -> Optional[dict]: + result = self._conn.execute( + text("SELECT * FROM prompts WHERE id = CAST(:id AS uuid) AND user_id = :user_id"), + {"id": prompt_id, "user_id": user_id}, + ) + row = result.fetchone() + return row_to_dict(row) if row is not None else None + + def get_for_rendering(self, prompt_id: str) -> Optional[dict]: + """Fetch prompt content by ID without user scoping. + + Used only by stream_processor to render a prompt whose owner is + not known at call time. Do NOT use in user-facing routes. + """ + result = self._conn.execute( + text("SELECT * FROM prompts WHERE id = CAST(:id AS uuid)"), + {"id": prompt_id}, + ) row = result.fetchone() return row_to_dict(row) if row is not None else None diff --git a/application/storage/db/repositories/sources.py b/application/storage/db/repositories/sources.py index add860f8..054f0553 100644 --- a/application/storage/db/repositories/sources.py +++ b/application/storage/db/repositories/sources.py @@ -5,9 +5,10 @@ from __future__ import annotations import json from typing import Optional -from sqlalchemy import Connection, text +from sqlalchemy import Connection, func, text from application.storage.db.base_repository import row_to_dict +from application.storage.db.models import sources_table class SourcesRepository: @@ -33,15 +34,7 @@ class SourcesRepository: ) return row_to_dict(result.fetchone()) - def get(self, source_id: str) -> Optional[dict]: - result = self._conn.execute( - text("SELECT * FROM sources WHERE id = CAST(:id AS uuid)"), - {"id": source_id}, - ) - row = result.fetchone() - return row_to_dict(row) if row is not None else None - - def get_for_user(self, source_id: str, user_id: str) -> Optional[dict]: + def get(self, source_id: str, user_id: str) -> Optional[dict]: result = self._conn.execute( text("SELECT * FROM sources WHERE id = CAST(:id AS uuid) AND user_id = :user_id"), {"id": source_id, "user_id": user_id}, @@ -56,27 +49,32 @@ class SourcesRepository: ) return [row_to_dict(r) for r in result.fetchall()] - def update(self, source_id: str, fields: dict) -> None: + def update(self, source_id: str, user_id: str, fields: dict) -> None: allowed = {"name", "type", "metadata"} filtered = {k: v for k, v in fields.items() if k in allowed} if not filtered: return - set_clauses = [] - params: dict = {"id": source_id} + + values: dict = {} for col, val in filtered.items(): if col == "metadata": - set_clauses.append(f"{col} = CAST(:val_{col} AS jsonb)") - params[f"val_{col}"] = json.dumps(val) if isinstance(val, dict) else val + values[col] = json.dumps(val) if isinstance(val, dict) else val else: - set_clauses.append(f"{col} = :val_{col}") - params[f"val_{col}"] = val - set_clauses.append("updated_at = now()") - sql = f"UPDATE sources SET {', '.join(set_clauses)} WHERE id = CAST(:id AS uuid)" - self._conn.execute(text(sql), params) + values[col] = val + values["updated_at"] = func.now() - def delete(self, source_id: str) -> bool: + t = sources_table + stmt = ( + t.update() + .where(t.c.id == source_id) + .where(t.c.user_id == user_id) + .values(**values) + ) + self._conn.execute(stmt) + + def delete(self, source_id: str, user_id: str) -> bool: result = self._conn.execute( - text("DELETE FROM sources WHERE id = CAST(:id AS uuid)"), - {"id": source_id}, + text("DELETE FROM sources WHERE id = CAST(:id AS uuid) AND user_id = :user_id"), + {"id": source_id, "user_id": user_id}, ) return result.rowcount > 0 diff --git a/application/storage/db/repositories/todos.py b/application/storage/db/repositories/todos.py index b45d5594..4adfd557 100644 --- a/application/storage/db/repositories/todos.py +++ b/application/storage/db/repositories/todos.py @@ -32,10 +32,10 @@ class TodosRepository: ) return row_to_dict(result.fetchone()) - def get(self, todo_id: str) -> Optional[dict]: + def get(self, todo_id: str, user_id: str) -> Optional[dict]: result = self._conn.execute( - text("SELECT * FROM todos WHERE id = CAST(:id AS uuid)"), - {"id": todo_id}, + text("SELECT * FROM todos WHERE id = CAST(:id AS uuid) AND user_id = :user_id"), + {"id": todo_id, "user_id": user_id}, ) row = result.fetchone() return row_to_dict(row) if row is not None else None @@ -50,25 +50,29 @@ class TodosRepository: ) return [row_to_dict(r) for r in result.fetchall()] - def update_title(self, todo_id: str, title: str) -> bool: - result = self._conn.execute( - text("UPDATE todos SET title = :title, updated_at = now() WHERE id = CAST(:id AS uuid)"), - {"id": todo_id, "title": title}, - ) - return result.rowcount > 0 - - def set_completed(self, todo_id: str, completed: bool = True) -> bool: + def update_title(self, todo_id: str, user_id: str, title: str) -> bool: result = self._conn.execute( text( - "UPDATE todos SET completed = :completed, updated_at = now() WHERE id = CAST(:id AS uuid)" + "UPDATE todos SET title = :title, updated_at = now() " + "WHERE id = CAST(:id AS uuid) AND user_id = :user_id" ), - {"id": todo_id, "completed": completed}, + {"id": todo_id, "user_id": user_id, "title": title}, ) return result.rowcount > 0 - def delete(self, todo_id: str) -> bool: + def set_completed(self, todo_id: str, user_id: str, completed: bool = True) -> bool: result = self._conn.execute( - text("DELETE FROM todos WHERE id = CAST(:id AS uuid)"), - {"id": todo_id}, + text( + "UPDATE todos SET completed = :completed, updated_at = now() " + "WHERE id = CAST(:id AS uuid) AND user_id = :user_id" + ), + {"id": todo_id, "user_id": user_id, "completed": completed}, + ) + return result.rowcount > 0 + + def delete(self, todo_id: str, user_id: str) -> bool: + result = self._conn.execute( + text("DELETE FROM todos WHERE id = CAST(:id AS uuid) AND user_id = :user_id"), + {"id": todo_id, "user_id": user_id}, ) return result.rowcount > 0 diff --git a/application/storage/db/repositories/user_tools.py b/application/storage/db/repositories/user_tools.py index 945347f6..1b5bd845 100644 --- a/application/storage/db/repositories/user_tools.py +++ b/application/storage/db/repositories/user_tools.py @@ -53,10 +53,10 @@ class UserToolsRepository: ) return row_to_dict(result.fetchone()) - def get(self, tool_id: str) -> Optional[dict]: + def get(self, tool_id: str, user_id: str) -> Optional[dict]: result = self._conn.execute( - text("SELECT * FROM user_tools WHERE id = CAST(:id AS uuid)"), - {"id": tool_id}, + text("SELECT * FROM user_tools WHERE id = CAST(:id AS uuid) AND user_id = :user_id"), + {"id": tool_id, "user_id": user_id}, ) row = result.fetchone() return row_to_dict(row) if row is not None else None diff --git a/scripts/db/backfill.py b/scripts/db/backfill.py index 2a712893..ed0a0baa 100644 --- a/scripts/db/backfill.py +++ b/scripts/db/backfill.py @@ -703,6 +703,8 @@ def _backfill_notes( """ INSERT INTO notes (user_id, tool_id, title, content) VALUES (:user_id, CAST(:tool_id AS uuid), :title, :content) + ON CONFLICT (user_id, tool_id) DO UPDATE + SET content = EXCLUDED.content, title = EXCLUDED.title """ ) cursor = mongo_db["notes"].find({}, no_cursor_timeout=True).batch_size(batch_size) @@ -743,7 +745,8 @@ def _backfill_connector_sessions( """ INSERT INTO connector_sessions (user_id, provider, session_data) VALUES (:user_id, :provider, CAST(:session_data AS jsonb)) - ON CONFLICT DO NOTHING + ON CONFLICT (user_id, provider) DO UPDATE + SET session_data = EXCLUDED.session_data """ ) cursor = mongo_db["connector_sessions"].find({}, no_cursor_timeout=True).batch_size(batch_size) diff --git a/tests/storage/db/repositories/test_agents.py b/tests/storage/db/repositories/test_agents.py index 86718ceb..0e37712a 100644 --- a/tests/storage/db/repositories/test_agents.py +++ b/tests/storage/db/repositories/test_agents.py @@ -48,26 +48,18 @@ class TestCreate: class TestGet: def test_get_existing(self, pg_conn): repo = _repo(pg_conn) - created = repo.create("u", "a", "draft") - fetched = repo.get(created["id"]) + created = repo.create("user-1", "a", "draft") + fetched = repo.get(created["id"], "user-1") assert fetched["id"] == created["id"] def test_get_nonexistent_returns_none(self, pg_conn): repo = _repo(pg_conn) - assert repo.get("00000000-0000-0000-0000-000000000000") is None + assert repo.get("00000000-0000-0000-0000-000000000000", "user-1") is None - -class TestGetForUser: - def test_get_for_correct_user(self, pg_conn): + def test_get_wrong_user_returns_none(self, pg_conn): repo = _repo(pg_conn) created = repo.create("user-1", "a", "draft") - fetched = repo.get_for_user(created["id"], "user-1") - assert fetched["id"] == created["id"] - - def test_get_for_wrong_user_returns_none(self, pg_conn): - repo = _repo(pg_conn) - created = repo.create("user-1", "a", "draft") - assert repo.get_for_user(created["id"], "user-other") is None + assert repo.get(created["id"], "user-other") is None class TestFindByKey: @@ -99,7 +91,7 @@ class TestUpdate: created = repo.create("user-1", "old", "draft") updated = repo.update(created["id"], "user-1", {"name": "new"}) assert updated is True - fetched = repo.get(created["id"]) + fetched = repo.get(created["id"], "user-1") assert fetched["name"] == "new" def test_update_wrong_user_returns_false(self, pg_conn): @@ -107,7 +99,7 @@ class TestUpdate: created = repo.create("user-1", "old", "draft") updated = repo.update(created["id"], "user-other", {"name": "new"}) assert updated is False - fetched = repo.get(created["id"]) + fetched = repo.get(created["id"], "user-1") assert fetched["name"] == "old" def test_update_disallowed_field_returns_false(self, pg_conn): @@ -123,14 +115,14 @@ class TestDelete: created = repo.create("user-1", "a", "draft") deleted = repo.delete(created["id"], "user-1") assert deleted is True - assert repo.get(created["id"]) is None + assert repo.get(created["id"], "user-1") is None def test_delete_wrong_user_returns_false(self, pg_conn): repo = _repo(pg_conn) created = repo.create("user-1", "a", "draft") deleted = repo.delete(created["id"], "user-other") assert deleted is False - assert repo.get(created["id"]) is not None + assert repo.get(created["id"], "user-1") is not None class TestSetFolder: @@ -142,7 +134,7 @@ class TestSetFolder: repo = _repo(pg_conn) agent = repo.create("user-1", "a", "draft") repo.set_folder(agent["id"], "user-1", folder["id"]) - fetched = repo.get(agent["id"]) + fetched = repo.get(agent["id"], "user-1") assert str(fetched["folder_id"]) == str(folder["id"]) def test_clear_folder(self, pg_conn): @@ -153,7 +145,7 @@ class TestSetFolder: repo = _repo(pg_conn) agent = repo.create("user-1", "a", "draft", folder_id=folder["id"]) repo.set_folder(agent["id"], "user-1", None) - fetched = repo.get(agent["id"]) + fetched = repo.get(agent["id"], "user-1") assert fetched["folder_id"] is None @@ -167,5 +159,5 @@ class TestClearFolderForAll: a1 = repo.create("user-1", "a1", "draft", folder_id=folder["id"]) a2 = repo.create("user-1", "a2", "draft", folder_id=folder["id"]) repo.clear_folder_for_all(folder["id"], "user-1") - assert repo.get(a1["id"])["folder_id"] is None - assert repo.get(a2["id"])["folder_id"] is None + assert repo.get(a1["id"], "user-1")["folder_id"] is None + assert repo.get(a2["id"], "user-1")["folder_id"] is None diff --git a/tests/storage/db/repositories/test_attachments.py b/tests/storage/db/repositories/test_attachments.py index c716dcec..0b24e8b1 100644 --- a/tests/storage/db/repositories/test_attachments.py +++ b/tests/storage/db/repositories/test_attachments.py @@ -42,12 +42,17 @@ class TestGet: def test_get_existing(self, pg_conn): repo = _repo(pg_conn) created = repo.create("u", "f", "/p") - fetched = repo.get(created["id"]) + fetched = repo.get(created["id"], "u") assert fetched["id"] == created["id"] def test_get_nonexistent_returns_none(self, pg_conn): repo = _repo(pg_conn) - assert repo.get("00000000-0000-0000-0000-000000000000") is None + assert repo.get("00000000-0000-0000-0000-000000000000", "u") is None + + def test_get_wrong_user_returns_none(self, pg_conn): + repo = _repo(pg_conn) + created = repo.create("u", "f", "/p") + assert repo.get(created["id"], "other") is None class TestListForUser: diff --git a/tests/storage/db/repositories/test_notes.py b/tests/storage/db/repositories/test_notes.py index 34930979..c082c7c8 100644 --- a/tests/storage/db/repositories/test_notes.py +++ b/tests/storage/db/repositories/test_notes.py @@ -70,12 +70,18 @@ class TestGetById: repo = _repo(pg_conn) tool_id = _make_tool(pg_conn) created = repo.upsert("u", tool_id, "t", "c") - fetched = repo.get(created["id"]) + fetched = repo.get(created["id"], "u") assert fetched["id"] == created["id"] def test_get_nonexistent_returns_none(self, pg_conn): repo = _repo(pg_conn) - assert repo.get("00000000-0000-0000-0000-000000000000") is None + assert repo.get("00000000-0000-0000-0000-000000000000", "u") is None + + def test_get_wrong_user_returns_none(self, pg_conn): + repo = _repo(pg_conn) + tool_id = _make_tool(pg_conn) + created = repo.upsert("u", tool_id, "t", "c") + assert repo.get(created["id"], "other") is None class TestDelete: diff --git a/tests/storage/db/repositories/test_prompts.py b/tests/storage/db/repositories/test_prompts.py index aee9dbd9..c126ab96 100644 --- a/tests/storage/db/repositories/test_prompts.py +++ b/tests/storage/db/repositories/test_prompts.py @@ -35,23 +35,30 @@ class TestGet: def test_get_by_id_and_user(self, pg_conn): repo = _repo(pg_conn) created = repo.create("user-1", "p", "c") - fetched = repo.get(created["id"], user_id="user-1") + fetched = repo.get(created["id"], "user-1") assert fetched["id"] == created["id"] - def test_get_by_id_only(self, pg_conn): - repo = _repo(pg_conn) - created = repo.create("user-1", "p", "c") - fetched = repo.get(created["id"]) - assert fetched is not None - def test_get_wrong_user_returns_none(self, pg_conn): repo = _repo(pg_conn) created = repo.create("user-1", "p", "c") - assert repo.get(created["id"], user_id="user-other") is None + assert repo.get(created["id"], "user-other") is None def test_get_nonexistent_returns_none(self, pg_conn): repo = _repo(pg_conn) - assert repo.get("00000000-0000-0000-0000-000000000000") is None + assert repo.get("00000000-0000-0000-0000-000000000000", "user-1") is None + + +class TestGetForRendering: + def test_returns_prompt_without_user_scoping(self, pg_conn): + repo = _repo(pg_conn) + created = repo.create("user-1", "p", "c") + fetched = repo.get_for_rendering(created["id"]) + assert fetched is not None + assert fetched["id"] == created["id"] + + def test_nonexistent_returns_none(self, pg_conn): + repo = _repo(pg_conn) + assert repo.get_for_rendering("00000000-0000-0000-0000-000000000000") is None class TestListForUser: @@ -70,7 +77,7 @@ class TestUpdate: repo = _repo(pg_conn) created = repo.create("user-1", "old", "old-content") repo.update(created["id"], "user-1", "new", "new-content") - fetched = repo.get(created["id"]) + fetched = repo.get(created["id"], "user-1") assert fetched["name"] == "new" assert fetched["content"] == "new-content" @@ -78,7 +85,7 @@ class TestUpdate: repo = _repo(pg_conn) created = repo.create("user-1", "old", "old-content") repo.update(created["id"], "user-other", "new", "new-content") - fetched = repo.get(created["id"]) + fetched = repo.get(created["id"], "user-1") assert fetched["name"] == "old" @@ -87,13 +94,13 @@ class TestDelete: repo = _repo(pg_conn) created = repo.create("user-1", "p", "c") repo.delete(created["id"], "user-1") - assert repo.get(created["id"]) is None + assert repo.get(created["id"], "user-1") is None def test_delete_wrong_user_is_noop(self, pg_conn): repo = _repo(pg_conn) created = repo.create("user-1", "p", "c") repo.delete(created["id"], "user-other") - assert repo.get(created["id"]) is not None + assert repo.get(created["id"], "user-1") is not None class TestFindOrCreate: diff --git a/tests/storage/db/repositories/test_sources.py b/tests/storage/db/repositories/test_sources.py index f54f3fd7..c4ea7501 100644 --- a/tests/storage/db/repositories/test_sources.py +++ b/tests/storage/db/repositories/test_sources.py @@ -43,28 +43,20 @@ class TestCreate: class TestGet: - def test_get_by_id(self, pg_conn): + def test_get_existing(self, pg_conn): repo = _repo(pg_conn) - created = repo.create("s") - fetched = repo.get(created["id"]) + created = repo.create("s", user_id="user-1") + fetched = repo.get(created["id"], "user-1") assert fetched["id"] == created["id"] def test_get_nonexistent_returns_none(self, pg_conn): repo = _repo(pg_conn) - assert repo.get("00000000-0000-0000-0000-000000000000") is None + assert repo.get("00000000-0000-0000-0000-000000000000", "user-1") is None - -class TestGetForUser: - def test_get_for_correct_user(self, pg_conn): + def test_get_wrong_user_returns_none(self, pg_conn): repo = _repo(pg_conn) created = repo.create("s", user_id="user-1") - fetched = repo.get_for_user(created["id"], "user-1") - assert fetched["id"] == created["id"] - - def test_get_for_wrong_user_returns_none(self, pg_conn): - repo = _repo(pg_conn) - created = repo.create("s", user_id="user-1") - assert repo.get_for_user(created["id"], "user-other") is None + assert repo.get(created["id"], "user-other") is None class TestListForUser: @@ -82,34 +74,48 @@ class TestUpdate: def test_updates_name(self, pg_conn): repo = _repo(pg_conn) created = repo.create("old", user_id="u") - repo.update(created["id"], {"name": "new"}) - fetched = repo.get(created["id"]) + repo.update(created["id"], "u", {"name": "new"}) + fetched = repo.get(created["id"], "u") assert fetched["name"] == "new" def test_updates_metadata(self, pg_conn): repo = _repo(pg_conn) - created = repo.create("s", metadata={"a": 1}) - repo.update(created["id"], {"metadata": {"a": 2, "b": 3}}) - fetched = repo.get(created["id"]) + created = repo.create("s", user_id="u", metadata={"a": 1}) + repo.update(created["id"], "u", {"metadata": {"a": 2, "b": 3}}) + fetched = repo.get(created["id"], "u") assert fetched["metadata"] == {"a": 2, "b": 3} def test_update_disallowed_field_is_noop(self, pg_conn): repo = _repo(pg_conn) - created = repo.create("s") - repo.update(created["id"], {"id": "00000000-0000-0000-0000-000000000000"}) - fetched = repo.get(created["id"]) + created = repo.create("s", user_id="u") + repo.update(created["id"], "u", {"id": "00000000-0000-0000-0000-000000000000"}) + fetched = repo.get(created["id"], "u") assert fetched["id"] == created["id"] + def test_update_wrong_user_is_noop(self, pg_conn): + repo = _repo(pg_conn) + created = repo.create("old", user_id="u") + repo.update(created["id"], "other-user", {"name": "new"}) + fetched = repo.get(created["id"], "u") + assert fetched["name"] == "old" + class TestDelete: def test_deletes_source(self, pg_conn): repo = _repo(pg_conn) - created = repo.create("s") - deleted = repo.delete(created["id"]) + created = repo.create("s", user_id="u") + deleted = repo.delete(created["id"], "u") assert deleted is True - assert repo.get(created["id"]) is None + assert repo.get(created["id"], "u") is None def test_delete_nonexistent_returns_false(self, pg_conn): repo = _repo(pg_conn) - deleted = repo.delete("00000000-0000-0000-0000-000000000000") + deleted = repo.delete("00000000-0000-0000-0000-000000000000", "u") assert deleted is False + + def test_delete_wrong_user_returns_false(self, pg_conn): + repo = _repo(pg_conn) + created = repo.create("s", user_id="u") + deleted = repo.delete(created["id"], "other-user") + assert deleted is False + assert repo.get(created["id"], "u") is not None diff --git a/tests/storage/db/repositories/test_todos.py b/tests/storage/db/repositories/test_todos.py index b573763e..a4d7f479 100644 --- a/tests/storage/db/repositories/test_todos.py +++ b/tests/storage/db/repositories/test_todos.py @@ -51,12 +51,18 @@ class TestGet: repo = _repo(pg_conn) tool_id = _make_tool(pg_conn) created = repo.create("u", tool_id, "t") - fetched = repo.get(created["id"]) + fetched = repo.get(created["id"], "u") assert fetched["id"] == created["id"] def test_get_nonexistent_returns_none(self, pg_conn): repo = _repo(pg_conn) - assert repo.get("00000000-0000-0000-0000-000000000000") is None + assert repo.get("00000000-0000-0000-0000-000000000000", "u") is None + + def test_get_wrong_user_returns_none(self, pg_conn): + repo = _repo(pg_conn) + tool_id = _make_tool(pg_conn) + created = repo.create("u", tool_id, "t") + assert repo.get(created["id"], "other") is None class TestListForUserTool: @@ -83,14 +89,23 @@ class TestUpdateTitle: repo = _repo(pg_conn) tool_id = _make_tool(pg_conn) created = repo.create("u", tool_id, "old") - updated = repo.update_title(created["id"], "new") + updated = repo.update_title(created["id"], "u", "new") assert updated is True - fetched = repo.get(created["id"]) + fetched = repo.get(created["id"], "u") assert fetched["title"] == "new" def test_update_nonexistent_returns_false(self, pg_conn): repo = _repo(pg_conn) - assert repo.update_title("00000000-0000-0000-0000-000000000000", "x") is False + assert repo.update_title("00000000-0000-0000-0000-000000000000", "u", "x") is False + + def test_update_wrong_user_returns_false(self, pg_conn): + repo = _repo(pg_conn) + tool_id = _make_tool(pg_conn) + created = repo.create("u", tool_id, "old") + updated = repo.update_title(created["id"], "other", "new") + assert updated is False + fetched = repo.get(created["id"], "u") + assert fetched["title"] == "old" class TestSetCompleted: @@ -98,17 +113,26 @@ class TestSetCompleted: repo = _repo(pg_conn) tool_id = _make_tool(pg_conn) created = repo.create("u", tool_id, "t") - repo.set_completed(created["id"], True) - fetched = repo.get(created["id"]) + repo.set_completed(created["id"], "u", True) + fetched = repo.get(created["id"], "u") assert fetched["completed"] is True def test_unmarks_completed(self, pg_conn): repo = _repo(pg_conn) tool_id = _make_tool(pg_conn) created = repo.create("u", tool_id, "t") - repo.set_completed(created["id"], True) - repo.set_completed(created["id"], False) - fetched = repo.get(created["id"]) + repo.set_completed(created["id"], "u", True) + repo.set_completed(created["id"], "u", False) + fetched = repo.get(created["id"], "u") + assert fetched["completed"] is False + + def test_set_completed_wrong_user_returns_false(self, pg_conn): + repo = _repo(pg_conn) + tool_id = _make_tool(pg_conn) + created = repo.create("u", tool_id, "t") + result = repo.set_completed(created["id"], "other", True) + assert result is False + fetched = repo.get(created["id"], "u") assert fetched["completed"] is False @@ -117,10 +141,18 @@ class TestDelete: repo = _repo(pg_conn) tool_id = _make_tool(pg_conn) created = repo.create("u", tool_id, "t") - deleted = repo.delete(created["id"]) + deleted = repo.delete(created["id"], "u") assert deleted is True - assert repo.get(created["id"]) is None + assert repo.get(created["id"], "u") is None def test_delete_nonexistent_returns_false(self, pg_conn): repo = _repo(pg_conn) - assert repo.delete("00000000-0000-0000-0000-000000000000") is False + assert repo.delete("00000000-0000-0000-0000-000000000000", "u") is False + + def test_delete_wrong_user_returns_false(self, pg_conn): + repo = _repo(pg_conn) + tool_id = _make_tool(pg_conn) + created = repo.create("u", tool_id, "t") + deleted = repo.delete(created["id"], "other") + assert deleted is False + assert repo.get(created["id"], "u") is not None diff --git a/tests/storage/db/repositories/test_user_tools.py b/tests/storage/db/repositories/test_user_tools.py index 9c830328..eb06e124 100644 --- a/tests/storage/db/repositories/test_user_tools.py +++ b/tests/storage/db/repositories/test_user_tools.py @@ -36,12 +36,17 @@ class TestGet: def test_get_existing(self, pg_conn): repo = _repo(pg_conn) created = repo.create("user-1", "t") - fetched = repo.get(created["id"]) + fetched = repo.get(created["id"], "user-1") assert fetched["id"] == created["id"] def test_get_nonexistent(self, pg_conn): repo = _repo(pg_conn) - assert repo.get("00000000-0000-0000-0000-000000000000") is None + assert repo.get("00000000-0000-0000-0000-000000000000", "user-1") is None + + def test_get_wrong_user_returns_none(self, pg_conn): + repo = _repo(pg_conn) + created = repo.create("user-1", "t") + assert repo.get(created["id"], "user-other") is None class TestListForUser: @@ -59,28 +64,28 @@ class TestUpdate: repo = _repo(pg_conn) created = repo.create("user-1", "old_name") repo.update(created["id"], "user-1", {"name": "new_name"}) - fetched = repo.get(created["id"]) + fetched = repo.get(created["id"], "user-1") assert fetched["name"] == "new_name" def test_updates_config(self, pg_conn): repo = _repo(pg_conn) created = repo.create("user-1", "t", config={"a": 1}) repo.update(created["id"], "user-1", {"config": {"a": 2, "b": 3}}) - fetched = repo.get(created["id"]) + fetched = repo.get(created["id"], "user-1") assert fetched["config"] == {"a": 2, "b": 3} def test_update_wrong_user_is_noop(self, pg_conn): repo = _repo(pg_conn) created = repo.create("user-1", "old") repo.update(created["id"], "user-other", {"name": "new"}) - fetched = repo.get(created["id"]) + fetched = repo.get(created["id"], "user-1") assert fetched["name"] == "old" def test_ignores_disallowed_fields(self, pg_conn): repo = _repo(pg_conn) created = repo.create("user-1", "t") repo.update(created["id"], "user-1", {"id": "00000000-0000-0000-0000-000000000000"}) - fetched = repo.get(created["id"]) + fetched = repo.get(created["id"], "user-1") assert fetched["id"] == created["id"] @@ -90,11 +95,11 @@ class TestDelete: created = repo.create("user-1", "t") deleted = repo.delete(created["id"], "user-1") assert deleted is True - assert repo.get(created["id"]) is None + assert repo.get(created["id"], "user-1") is None def test_delete_wrong_user_returns_false(self, pg_conn): repo = _repo(pg_conn) created = repo.create("user-1", "t") deleted = repo.delete(created["id"], "user-other") assert deleted is False - assert repo.get(created["id"]) is not None + assert repo.get(created["id"], "user-1") is not None