diff --git a/application/storage/db/engine.py b/application/storage/db/engine.py index 9ca0fff4..5102b81e 100644 --- a/application/storage/db/engine.py +++ b/application/storage/db/engine.py @@ -24,38 +24,25 @@ _engine: Optional[Engine] = None def _resolve_uri() -> str: - """Pick the Postgres URI, falling back to the pgvector connection string. + """Return the Postgres URI for user-data tables. - If ``POSTGRES_URI`` is set explicitly, use it. Otherwise, if the app - is already using pgvector with a Postgres-backed vector store - (``PGVECTOR_CONNECTION_STRING``), reuse that same cluster for - user-data tables — they can share a database. The pgvector URI is - normalized from its libpq form to the SQLAlchemy psycopg3 dialect. + Raises: + RuntimeError: If ``settings.POSTGRES_URI`` is unset. Callers that + reach this path without a configured URI have a setup bug — the + error message points them at the right setting. """ - if settings.POSTGRES_URI: - return settings.POSTGRES_URI - - if settings.PGVECTOR_CONNECTION_STRING: - from application.core.db_uri import normalize_postgres_uri - - uri = normalize_postgres_uri(settings.PGVECTOR_CONNECTION_STRING) - if uri: - return uri - - raise RuntimeError( - "POSTGRES_URI is not configured and no PGVECTOR_CONNECTION_STRING " - "to fall back to. Set POSTGRES_URI in your .env to a URI such as " - "'postgresql+psycopg://user:pass@host:5432/docsgpt'." - ) + if not settings.POSTGRES_URI: + raise RuntimeError( + "POSTGRES_URI is not configured. Set it in your .env to a " + "psycopg3 URI such as " + "'postgresql+psycopg://user:pass@host:5432/docsgpt'." + ) + return settings.POSTGRES_URI def get_engine() -> Engine: """Return the process-wide SQLAlchemy Engine, creating it if needed. - Falls back to ``PGVECTOR_CONNECTION_STRING`` when ``POSTGRES_URI`` - is not set, so operators using pgvector on the same cluster don't - need a second env var. - Returns: A SQLAlchemy ``Engine`` configured with a pooled connection to Postgres via psycopg3. diff --git a/application/storage/db/repositories/user_tools.py b/application/storage/db/repositories/user_tools.py index 7d8bfab8..945347f6 100644 --- a/application/storage/db/repositories/user_tools.py +++ b/application/storage/db/repositories/user_tools.py @@ -78,18 +78,33 @@ class UserToolsRepository: filtered = {k: v for k, v in fields.items() if k in allowed} if not filtered: return - set_clauses = [] - params: dict = {"id": tool_id, "user_id": user_id} - for col, val in filtered.items(): - if col == "config": - set_clauses.append(f"{col} = CAST(:val_{col} AS jsonb)") - params[f"val_{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 user_tools SET {', '.join(set_clauses)} WHERE id = CAST(:id AS uuid) AND user_id = :user_id" - self._conn.execute(text(sql), params) + params: dict = { + "id": tool_id, + "user_id": user_id, + "name": filtered.get("name"), + "custom_name": filtered.get("custom_name"), + "display_name": filtered.get("display_name"), + "config": ( + json.dumps(filtered["config"]) + if "config" in filtered and isinstance(filtered["config"], dict) + else filtered.get("config") + ), + } + self._conn.execute( + text( + """ + UPDATE user_tools + SET + name = COALESCE(:name, name), + custom_name = COALESCE(:custom_name, custom_name), + display_name = COALESCE(:display_name, display_name), + config = COALESCE(CAST(:config AS jsonb), config), + updated_at = now() + WHERE id = CAST(:id AS uuid) AND user_id = :user_id + """ + ), + params, + ) def delete(self, tool_id: str, user_id: str) -> bool: result = self._conn.execute( diff --git a/application/vectorstore/pgvector.py b/application/vectorstore/pgvector.py index eb4e7178..3066d8cd 100644 --- a/application/vectorstore/pgvector.py +++ b/application/vectorstore/pgvector.py @@ -27,13 +27,20 @@ class PGVectorStore(BaseVectorStore): self._metadata_column = metadata_column self._embedding = self._get_embeddings(settings.EMBEDDINGS_NAME, embeddings_key) - # Use provided connection string or fall back to settings + # Use provided connection string or fall back to settings. + # If PGVECTOR_CONNECTION_STRING is not set but POSTGRES_URI is, + # reuse the same cluster — normalize from SQLAlchemy dialect to libpq form. self._connection_string = connection_string or getattr(settings, 'PGVECTOR_CONNECTION_STRING', None) - + + if not self._connection_string and getattr(settings, 'POSTGRES_URI', None): + from application.core.db_uri import normalize_pgvector_connection_string + self._connection_string = normalize_pgvector_connection_string(settings.POSTGRES_URI) + if not self._connection_string: raise ValueError( "PostgreSQL connection string is required. " - "Set PGVECTOR_CONNECTION_STRING in settings or pass connection_string parameter." + "Set PGVECTOR_CONNECTION_STRING or POSTGRES_URI in settings, " + "or pass connection_string parameter." ) try: diff --git a/tests/storage/db/repositories/test_user_tools.py b/tests/storage/db/repositories/test_user_tools.py index d3e36059..9c830328 100644 --- a/tests/storage/db/repositories/test_user_tools.py +++ b/tests/storage/db/repositories/test_user_tools.py @@ -88,11 +88,13 @@ class TestDelete: def test_deletes_tool(self, pg_conn): repo = _repo(pg_conn) created = repo.create("user-1", "t") - assert repo.delete(created["id"], "user-1") is True + deleted = repo.delete(created["id"], "user-1") + assert deleted is True assert repo.get(created["id"]) is None def test_delete_wrong_user_returns_false(self, pg_conn): repo = _repo(pg_conn) created = repo.create("user-1", "t") - assert repo.delete(created["id"], "user-other") is False + deleted = repo.delete(created["id"], "user-other") + assert deleted is False assert repo.get(created["id"]) is not None