refactor(phase-7): decompose MilvusDB god object into VectorStore + CatalogStore adapters - #393
Merged
Ahmath-Gadji merged 24 commits intoMay 18, 2026
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughThis PR introduces an async PostgreSQL ConnectionManager with Alembic-compatible metadata, SQLAlchemy schema definitions, and a Milvus 2.6-backed MilvusVectorStore supporting dense/hybrid search, upsert, delete, collection administration, and related unit + integration tests. Config, env mappings, migration docs, and decision-log entries were also added. ChangesPostgreSQL persistence, migrations, and configuration
Milvus vector store implementation
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
openrag/services/persistence/migrations/env.py (1)
34-41:⚠️ Potential issue | 🔴 Critical | 🏗️ Heavy liftenv.py always derives the database name from the collection name, but
ConnectionManagerrequires an explicitly setPOSTGRES_DATABASEenvironment variable, risking a mismatch between migration and runtime databases.Lines 27–32 of
env.pyuserag_config.rdb.user,.password,.port,.hostbut explicitly skip.database, instead deriving the database name asf"partitions_for_collection_{rag_config.vectordb.collection_name}". Meanwhile,ConnectionManager.__init__()(connection.py:46–50) raisesValueErrorifconfig.databaseis falsy, requiringPOSTGRES_DATABASEto be set via environment.The infrastructure.py comments (37–40) claim the database name is "wired by the caller," but no such wiring code exists. If
POSTGRES_DATABASEis not set or differs from the formula-derived name, migrations and runtime operations will target different databases, breaking the catalog.Either:
env.pyshould respectconfig.rdb.databaseif set, falling back to the collection-derived name only if None, or- The initialization code must explicitly wire the collection name to
rdb.databasebeforeConnectionManagerconstruction.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@openrag/services/persistence/migrations/env.py` around lines 34 - 41, env.py currently always builds database_url using f"partitions_for_collection_{rag_config.vectordb.collection_name}" which can mismatch ConnectionManager (ConnectionManager.__init__) that requires config.rdb.database/POSTGRES_DATABASE; change env.py to first check for an explicit database value (config.rdb.database or rag_config.rdb.database) and use it if present, otherwise fall back to the collection-derived name (f"partitions_for_collection_{rag_config.vectordb.collection_name}"); ensure the resultant database name is assigned to the config used by ConnectionManager so migrations and runtime share the same config.
🧹 Nitpick comments (6)
openrag/services/persistence/connection.py (3)
79-84: ⚡ Quick winConsider adding
server_settingsparameter for application_name.Setting
application_namein the connection pool helps with database monitoring and identifying connections inpg_stat_activity. This is especially useful in multi-service deployments.Add application_name
self._pool = await asyncpg.create_pool( self._dsn, min_size=self._min_size, max_size=self._max_size, command_timeout=self._command_timeout, + server_settings={"application_name": "openrag-persistence"}, )🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@openrag/services/persistence/connection.py` around lines 79 - 84, The connection pool creation using asyncpg.create_pool (where self._pool is assigned) should include server_settings to set application_name for better DB observability; update the asyncpg.create_pool call to pass server_settings={'application_name': <name>} (preferably using an injected/config value such as a new self._application_name or a parameter added to the class initializer) and ensure any callers/config loading are updated to provide that value so the pool creation uses the configured application name.
114-118: 💤 Low valuePool shutdown should log closure.
For operational observability, log when the pool is closed, similar to the log on successful connection (line 85-90). This helps trace connection lifecycle in production logs.
Add shutdown logging
async def shutdown(self) -> None: if self._pool is None: return await self._pool.close() self._pool = None + logger.info("Closed Postgres connection pool for %s", self._dsn_log)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@openrag/services/persistence/connection.py` around lines 114 - 118, The shutdown method currently closes the pool but does not log the closure; update the async def shutdown(self) to log an informational message when the pool is closed (use the same logger used by the connection setup/logging in this class) before setting self._pool = None, e.g., detect self._pool is not None, await self._pool.close(), call the class logger (e.g., self._logger.info or the module logger used in connect) with a clear message that the connection pool was closed, then set self._pool = None.
67-112: ⚡ Quick winRetry logic should catch only transient connection errors, not permanent configuration errors.
The current retry catches
(OSError, asyncpg.PostgresError), which is too broad.asyncpg.PostgresErrorincludes permanent configuration errors likeInvalidAuthorizationSpecificationError(auth/SSL mismatches) andInvalidCatalogNameError(database doesn't exist) that should fail immediately instead of retrying.Replace with only the transient errors:
OSError(network-level issues)asyncpg.CannotConnectNowError(database not ready, SQLSTATE 57P03)asyncpg.TooManyConnectionsError(temporary resource exhaustion, SQLSTATE 53300)Refined exception handling
- except (OSError, asyncpg.PostgresError) as exc: + except ( + OSError, + asyncpg.CannotConnectNowError, + asyncpg.TooManyConnectionsError, + ) as exc:This prevents unnecessary retry delays during startup when configuration is invalid.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@openrag/services/persistence/connection.py` around lines 67 - 112, The initialize method currently retries on any asyncpg.PostgresError which can hide permanent config issues; update the exception handling in initialize to catch only transient errors: OSError, asyncpg.CannotConnectNowError, and asyncpg.TooManyConnectionsError (instead of asyncpg.PostgresError) so permanent errors like InvalidAuthorizationSpecificationError or InvalidCatalogNameError bubble up immediately; keep the existing retry/backoff logic using _RETRY_ATTEMPTS and _RETRY_BASE_DELAY_SECONDS/_RETRY_MAX_DELAY_SECONDS, preserve setting last_exc and the final RuntimeError raise-from behavior, and only log/retry when the exception is one of those transient types.openrag/core/config/infrastructure.py (1)
42-44: 💤 Low valueConsider documenting connection pool sizing rationale.
The pool defaults (
min=5,max=20,command_timeout=30s) are reasonable for moderate loads, but consider adding inline comments explaining the sizing rationale or linking to performance testing results—especially important for production deployments where these may need tuning.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@openrag/core/config/infrastructure.py` around lines 42 - 44, Add brief inline comments next to the pool configuration variables (pool_min_size, pool_max_size, command_timeout) explaining the rationale for the chosen defaults (e.g., expected concurrent connections, memory/CPU tradeoffs, and typical workload these values target), state units (seconds for command_timeout), and include a pointer or TODO to link performance test results or tuning guidance for production deployments so operators know when and how to adjust them.openrag/services/persistence/schema.py (2)
166-168: 💤 Low valueWorkspace
created_byusesondelete="SET NULL".Similar to files, when a user is deleted, workspaces they created will have
created_byset to NULL. Confirm this matches the intended audit/ownership semantics.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@openrag/services/persistence/schema.py` around lines 166 - 168, The workspace column definition uses ForeignKey("users.id", ondelete="SET NULL") with nullable=True for created_by; confirm desired audit/ownership behavior and update the schema accordingly: if workspaces must retain a creator reference make created_by non-nullable and change ondelete to "RESTRICT" (or remove ondelete), or if losing creator is acceptable keep "SET NULL" and nullable=True; update the created_by Column/ForeignKey in the Workspace model (the created_by column / ForeignKey("users.id", ondelete="SET NULL")) to reflect the chosen semantics and run migrations.
62-64: ⚡ Quick winReview
ondelete="SET NULL"behavior forcreated_byforeign key.The
created_bycolumn usesondelete="SET NULL", which means if a user is deleted, file records will havecreated_byset to NULL. Verify this is the intended behavior for audit trail purposes—you may want to preserve the user ID or use a different cascade strategy if user deletion should be restricted when files exist.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@openrag/services/persistence/schema.py` around lines 62 - 64, The created_by column in openrag/services/persistence/schema.py currently uses ForeignKey("users.id", ondelete="SET NULL"); confirm whether you want deleted users to null out file.created_by or to prevent deletion/preserve the id for auditing. If you need to preserve the user id for audit trails, change the ForeignKey to ondelete="RESTRICT" (or remove ondelete entirely) and make created_by non-nullable if appropriate; if you prefer to keep an immutable audit copy, add a separate created_by_name/created_by_id audit column and stop relying on the FK for historical attribution. Update the ForeignKey definition and any related model constraints (e.g., nullable flag) in the schema accordingly.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@openrag/services/persistence/connection.py`:
- Line 31: Replace the stdlib logger in connection.py (the variable logger
created via logging.getLogger(__name__)) with the project's Loguru wrapper by
importing get_logger from utils.logger and calling get_logger(__name__) to
obtain the logger; update any existing log statements to use that Loguru logger
and where appropriate bind context like file_id and partition (e.g.,
logger.bind(file_id=..., partition=...)) before logging so structured fields are
included.
- Around line 139-145: Wrap the call to command.upgrade(cfg, "head") in a
try/except around the migration execution in connection.py (the block that
builds cfg, sets "script_location" and "sqlalchemy.url" and logs using
self._dsn_log) to catch Alembic and SQL errors; on exception, log a clear error
via the same logger including self._dsn_log and the caught exception (and its
traceback if desired), then re-raise or exit so failure is visible instead of
producing a cryptic stacktrace. Use broad Exception to ensure
alembic.command.upgrade failures are captured but keep the log message specific
(e.g., "Alembic migration failed for %s: %s" with self._dsn_log and exception).
In `@openrag/services/persistence/schema.py`:
- Around line 129-131: The partition FK on the files table is missing an
explicit ondelete strategy, causing asymmetry with partition_memberships and
workspace which use ForeignKey(..., ondelete="CASCADE"); update the files table
ForeignKey("partitions.partition") to include ondelete="CASCADE" (or
alternatively document and enforce that files must be removed before deleting a
partition) so partition deletions behave consistently; locate and modify the
ForeignKey usage in the files table definition and any other partition FKs
(e.g., in partition_memberships and workspace) to ensure consistent cascade
semantics.
- Line 170: The workspaces table's created_at Column declaration
(Column("created_at", DateTime, default=datetime.now)) is missing the nullable
constraint; update that Column in openrag/services/persistence/schema.py to
include nullable=False (i.e., Column("created_at", DateTime,
default=datetime.now, nullable=False)) so it matches the other created_at
definitions and prevents NULL workspace creation timestamps.
---
Outside diff comments:
In `@openrag/services/persistence/migrations/env.py`:
- Around line 34-41: env.py currently always builds database_url using
f"partitions_for_collection_{rag_config.vectordb.collection_name}" which can
mismatch ConnectionManager (ConnectionManager.__init__) that requires
config.rdb.database/POSTGRES_DATABASE; change env.py to first check for an
explicit database value (config.rdb.database or rag_config.rdb.database) and use
it if present, otherwise fall back to the collection-derived name
(f"partitions_for_collection_{rag_config.vectordb.collection_name}"); ensure the
resultant database name is assigned to the config used by ConnectionManager so
migrations and runtime share the same config.
---
Nitpick comments:
In `@openrag/core/config/infrastructure.py`:
- Around line 42-44: Add brief inline comments next to the pool configuration
variables (pool_min_size, pool_max_size, command_timeout) explaining the
rationale for the chosen defaults (e.g., expected concurrent connections,
memory/CPU tradeoffs, and typical workload these values target), state units
(seconds for command_timeout), and include a pointer or TODO to link performance
test results or tuning guidance for production deployments so operators know
when and how to adjust them.
In `@openrag/services/persistence/connection.py`:
- Around line 79-84: The connection pool creation using asyncpg.create_pool
(where self._pool is assigned) should include server_settings to set
application_name for better DB observability; update the asyncpg.create_pool
call to pass server_settings={'application_name': <name>} (preferably using an
injected/config value such as a new self._application_name or a parameter added
to the class initializer) and ensure any callers/config loading are updated to
provide that value so the pool creation uses the configured application name.
- Around line 114-118: The shutdown method currently closes the pool but does
not log the closure; update the async def shutdown(self) to log an informational
message when the pool is closed (use the same logger used by the connection
setup/logging in this class) before setting self._pool = None, e.g., detect
self._pool is not None, await self._pool.close(), call the class logger (e.g.,
self._logger.info or the module logger used in connect) with a clear message
that the connection pool was closed, then set self._pool = None.
- Around line 67-112: The initialize method currently retries on any
asyncpg.PostgresError which can hide permanent config issues; update the
exception handling in initialize to catch only transient errors: OSError,
asyncpg.CannotConnectNowError, and asyncpg.TooManyConnectionsError (instead of
asyncpg.PostgresError) so permanent errors like
InvalidAuthorizationSpecificationError or InvalidCatalogNameError bubble up
immediately; keep the existing retry/backoff logic using _RETRY_ATTEMPTS and
_RETRY_BASE_DELAY_SECONDS/_RETRY_MAX_DELAY_SECONDS, preserve setting last_exc
and the final RuntimeError raise-from behavior, and only log/retry when the
exception is one of those transient types.
In `@openrag/services/persistence/schema.py`:
- Around line 166-168: The workspace column definition uses
ForeignKey("users.id", ondelete="SET NULL") with nullable=True for created_by;
confirm desired audit/ownership behavior and update the schema accordingly: if
workspaces must retain a creator reference make created_by non-nullable and
change ondelete to "RESTRICT" (or remove ondelete), or if losing creator is
acceptable keep "SET NULL" and nullable=True; update the created_by
Column/ForeignKey in the Workspace model (the created_by column /
ForeignKey("users.id", ondelete="SET NULL")) to reflect the chosen semantics and
run migrations.
- Around line 62-64: The created_by column in
openrag/services/persistence/schema.py currently uses ForeignKey("users.id",
ondelete="SET NULL"); confirm whether you want deleted users to null out
file.created_by or to prevent deletion/preserve the id for auditing. If you need
to preserve the user id for audit trails, change the ForeignKey to
ondelete="RESTRICT" (or remove ondelete entirely) and make created_by
non-nullable if appropriate; if you prefer to keep an immutable audit copy, add
a separate created_by_name/created_by_id audit column and stop relying on the FK
for historical attribution. Update the ForeignKey definition and any related
model constraints (e.g., nullable flag) in the schema accordingly.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: b5b97fc2-60fd-4553-8fed-7db772fd278e
📒 Files selected for processing (20)
REFACTORING_DECISION_LOG.mddocs/content/docs/documentation/sql_migration.mdxopenrag/core/config/infrastructure.pyopenrag/core/config/loader.pyopenrag/services/persistence/__init__.pyopenrag/services/persistence/connection.pyopenrag/services/persistence/migrations/READMEopenrag/services/persistence/migrations/alembic.iniopenrag/services/persistence/migrations/env.pyopenrag/services/persistence/migrations/schema_helpers.pyopenrag/services/persistence/migrations/script.py.makoopenrag/services/persistence/migrations/versions/4add4d260575_initial_migration.pyopenrag/services/persistence/migrations/versions/a1b2c3d4e5f6_add_document_relationships.pyopenrag/services/persistence/migrations/versions/c224d4befe71_add_file_count_and_file_quota.pyopenrag/services/persistence/migrations/versions/cd642e4502d8_create_users_memberships_tables.pyopenrag/services/persistence/migrations/versions/cd9b84278028_merge_heads.pyopenrag/services/persistence/migrations/versions/e7f8a9b0c1d2_add_workspaces.pyopenrag/services/persistence/migrations/versions/f1a2b3c4d5e6_add_workspace_files_file_id_fk.pyopenrag/services/persistence/migrations/versions/f5b6c918f741_add_oidc_auth.pyopenrag/services/persistence/schema.py
✅ Files skipped from review due to trivial changes (4)
- openrag/core/config/loader.py
- openrag/services/persistence/init.py
- docs/content/docs/documentation/sql_migration.mdx
- REFACTORING_DECISION_LOG.md
EnjoyBacon7
force-pushed
the
refactor/phase-7-storage-persistence-adapters
branch
from
May 12, 2026 15:11
3c48a9d to
4bb7109
Compare
There was a problem hiding this comment.
Actionable comments posted: 5
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In
`@openrag/services/persistence/migrations/milvus/1.add_created_at_temporal_fields.py`:
- Around line 21-29: Update the example command lines in the migration doc block
to reference the correct script name: replace all occurrences of
"1.add_temporal_fields.py" with "1.add_created_at_temporal_fields.py" so the
dry-run, apply, and downgrade examples point to the actual file; modify the
three example lines in the header of 1.add_created_at_temporal_fields.py that
build the docker compose commands accordingly.
In `@openrag/services/storage/milvus_store.py`:
- Around line 616-620: The inserted entities are being given an indexed_at
timestamp but the schema expects created_at for temporal queries; change the
code that builds entities (uses indexed_at variable and calls _chunk_to_entity)
to set created_at = datetime.now(UTC).isoformat() and pass created_at into
_chunk_to_entity (and into any other entity construction sites at the earlier
block around lines referencing the same pattern, e.g., the other occurrence at
576-577), optionally keeping indexed_at if you still need it but ensure
created_at is populated on each entity so temporal filtering works correctly.
- Around line 283-289: The code adds the "sparse" field in _create_schema() only
when self._hybrid is true but _create_index() always creates an index for
"sparse", causing failures for dense-only configs; update _create_index() to
guard any logic that creates or configures the "sparse" index (index_type
"SPARSE_INVERTED_INDEX" / field_name "sparse") behind the same condition
(self._hybrid or hybrid_search flag) used in _create_schema(), and apply the
same guard to the second index block referenced around the other occurrence so
that indexing of the sparse field only runs when hybrid mode is enabled.
- Around line 872-876: The delete_by_filter guard in delete_by_filter currently
only checks for empty expr but allows tautological raw expressions like "1==1"
or "true" which will delete the whole collection; update delete_by_filter to
detect and reject simple tautologies by inspecting the incoming filters/expr
(e.g., the expr string passed to delete_by_filter) and raising the same
ValueError for expressions that match common tautologies (case-insensitive
"true", numeric tautologies like "1==1", or equivalent whitespace variations) so
callers must explicitly call drop_collection() to remove all rows.
- Around line 451-456: The current partition-list handling in milvus_store.py
lets a wildcard (e.g. "all") in the list disable partition filtering even when
explicit partitions are present, widening scope; update the logic in the block
that inspects partition (the branch that checks isinstance(partition, (list,
tuple)), using _PARTITION_WILDCARDS and _format_value) to detect a mixed list
containing both wildcards and explicit partitions and reject or sanitize it:
either raise a ValueError (or similar) when a mix is provided, or strip
wildcards and only use explicit entries when building parts (i.e. use [p for p
in partition if p not in _PARTITION_WILDCARDS] and only append "partition in
[...]" if that resulting list is non-empty). Ensure the change references
partition, _PARTITION_WILDCARDS, _format_value and parts so callers cannot
accidentally widen scope.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 928d3a6d-f896-4837-9636-0446b4ae9a79
📒 Files selected for processing (22)
REFACTORING_DECISION_LOG.mddocs/content/docs/documentation/milvus_migration.mdxdocs/content/docs/documentation/sql_migration.mdxopenrag/services/persistence/migrations/alembic/READMEopenrag/services/persistence/migrations/alembic/__init__.pyopenrag/services/persistence/migrations/alembic/alembic.iniopenrag/services/persistence/migrations/alembic/env.pyopenrag/services/persistence/migrations/alembic/schema_helpers.pyopenrag/services/persistence/migrations/alembic/script.py.makoopenrag/services/persistence/migrations/alembic/versions/4add4d260575_initial_migration.pyopenrag/services/persistence/migrations/alembic/versions/__init__.pyopenrag/services/persistence/migrations/alembic/versions/a1b2c3d4e5f6_add_document_relationships.pyopenrag/services/persistence/migrations/alembic/versions/c224d4befe71_add_file_count_and_file_quota.pyopenrag/services/persistence/migrations/alembic/versions/cd642e4502d8_create_users_memberships_tables.pyopenrag/services/persistence/migrations/alembic/versions/cd9b84278028_merge_heads.pyopenrag/services/persistence/migrations/alembic/versions/e7f8a9b0c1d2_add_workspaces.pyopenrag/services/persistence/migrations/alembic/versions/f1a2b3c4d5e6_add_workspace_files_file_id_fk.pyopenrag/services/persistence/migrations/alembic/versions/f5b6c918f741_add_oidc_auth.pyopenrag/services/persistence/migrations/milvus/1.add_created_at_temporal_fields.pyopenrag/services/persistence/migrations/milvus/__init__.pyopenrag/services/persistence/migrations/milvus/migrate.pyopenrag/services/storage/milvus_store.py
✅ Files skipped from review due to trivial changes (2)
- REFACTORING_DECISION_LOG.md
- docs/content/docs/documentation/sql_migration.mdx
There was a problem hiding this comment.
🧹 Nitpick comments (2)
tests/integration/test_milvus_store_integration.py (2)
286-286: 💤 Low valueConsider using
hybrid_config.collection_nameinstead of accessing_collection_name.Lines 286, 301, and 311 access the private
_collection_nameattribute. Sincequery_ids_by_filterandquery_chunks_by_filteraccept acollectionparameter, usinghybrid_config.collection_namewould avoid accessing private attributes and make the tests cleaner.♻️ Refactor to use config value
- remaining_p1 = await hybrid_store.query_ids_by_filter(hybrid_store._collection_name, {"partition": "p1"}) + remaining_p1 = await hybrid_store.query_ids_by_filter(hybrid_config.collection_name, {"partition": "p1"})Apply similar changes to lines 301, 311, 342, and 346.
Also applies to: 301-301, 311-311
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/integration/test_milvus_store_integration.py` at line 286, Tests are accessing the private attribute hybrid_store._collection_name; replace those usages with the public config value hybrid_config.collection_name when calling hybrid_store.query_ids_by_filter and hybrid_store.query_chunks_by_filter (e.g., change calls like hybrid_store.query_ids_by_filter(hybrid_store._collection_name, ...) to hybrid_store.query_ids_by_filter(hybrid_config.collection_name, ...)), and apply the same replacement for all occurrences referenced (lines around usages at 286, 301, 311, 342, 346) so tests no longer depend on private attributes.
315-321: Vector field leakage documented as known issue.The NOTE comment correctly documents that
query_chunks_by_filtercurrently leaks the dense vector field when using defaultoutput_fields, which contradicts the method docstring and is asymmetric withsearch()behavior. This is good practice to document current behavior while tracking the fix.Do you want me to open a tracking issue for the vector field stripping in
query_chunks_by_filter?🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/integration/test_milvus_store_integration.py` around lines 315 - 321, Open a tracking issue titled something like "Strip vector field in query_chunks_by_filter to match search() and docstring" and include: a short summary of the bug (query_chunks_by_filter leaks dense vector when called with default output_fields), reproduction steps referencing the test that demonstrates current behavior (tests/integration/test_milvus_store_integration.py which shows _iter_query leaking the vector), the intended behavior (strip vector like search() using _SEARCH_RESULT_DROPPED_KEYS), and link to the relevant symbols/functions to change: query_chunks_by_filter, _iter_query, and the constant _SEARCH_RESULT_DROPPED_KEYS; assign it to the milvus store owner and mark as bug/regression for follow-up work.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@tests/integration/test_milvus_store_integration.py`:
- Line 286: Tests are accessing the private attribute
hybrid_store._collection_name; replace those usages with the public config value
hybrid_config.collection_name when calling hybrid_store.query_ids_by_filter and
hybrid_store.query_chunks_by_filter (e.g., change calls like
hybrid_store.query_ids_by_filter(hybrid_store._collection_name, ...) to
hybrid_store.query_ids_by_filter(hybrid_config.collection_name, ...)), and apply
the same replacement for all occurrences referenced (lines around usages at 286,
301, 311, 342, 346) so tests no longer depend on private attributes.
- Around line 315-321: Open a tracking issue titled something like "Strip vector
field in query_chunks_by_filter to match search() and docstring" and include: a
short summary of the bug (query_chunks_by_filter leaks dense vector when called
with default output_fields), reproduction steps referencing the test that
demonstrates current behavior
(tests/integration/test_milvus_store_integration.py which shows _iter_query
leaking the vector), the intended behavior (strip vector like search() using
_SEARCH_RESULT_DROPPED_KEYS), and link to the relevant symbols/functions to
change: query_chunks_by_filter, _iter_query, and the constant
_SEARCH_RESULT_DROPPED_KEYS; assign it to the milvus store owner and mark as
bug/regression for follow-up work.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 95959b08-8ec7-4ad7-9791-8000eb3b2f70
📒 Files selected for processing (30)
REFACTORING_DECISION_LOG.mddocs/content/docs/documentation/milvus_migration.mdxdocs/content/docs/documentation/sql_migration.mdxopenrag/services/persistence/connection.pyopenrag/services/persistence/migrations/alembic/READMEopenrag/services/persistence/migrations/alembic/__init__.pyopenrag/services/persistence/migrations/alembic/alembic.iniopenrag/services/persistence/migrations/alembic/env.pyopenrag/services/persistence/migrations/alembic/schema_helpers.pyopenrag/services/persistence/migrations/alembic/script.py.makoopenrag/services/persistence/migrations/alembic/versions/4add4d260575_initial_migration.pyopenrag/services/persistence/migrations/alembic/versions/__init__.pyopenrag/services/persistence/migrations/alembic/versions/a1b2c3d4e5f6_add_document_relationships.pyopenrag/services/persistence/migrations/alembic/versions/c224d4befe71_add_file_count_and_file_quota.pyopenrag/services/persistence/migrations/alembic/versions/cd642e4502d8_create_users_memberships_tables.pyopenrag/services/persistence/migrations/alembic/versions/cd9b84278028_merge_heads.pyopenrag/services/persistence/migrations/alembic/versions/e7f8a9b0c1d2_add_workspaces.pyopenrag/services/persistence/migrations/alembic/versions/f1a2b3c4d5e6_add_workspace_files_file_id_fk.pyopenrag/services/persistence/migrations/alembic/versions/f5b6c918f741_add_oidc_auth.pyopenrag/services/persistence/migrations/milvus/1.add_created_at_temporal_fields.pyopenrag/services/persistence/migrations/milvus/__init__.pyopenrag/services/persistence/migrations/milvus/migrate.pyopenrag/services/persistence/test_document_repo.pyopenrag/services/persistence/test_oidc_session_repo.pyopenrag/services/persistence/test_partition_membership_repo.pyopenrag/services/persistence/test_partition_repo.pyopenrag/services/persistence/test_user_repo.pyopenrag/services/persistence/test_workspace_repo.pyopenrag/services/storage/test_milvus_store.pytests/integration/test_milvus_store_integration.py
✅ Files skipped from review due to trivial changes (6)
- openrag/services/persistence/test_partition_membership_repo.py
- openrag/services/persistence/test_document_repo.py
- openrag/services/persistence/test_oidc_session_repo.py
- openrag/services/persistence/test_partition_repo.py
- docs/content/docs/documentation/milvus_migration.mdx
- docs/content/docs/documentation/sql_migration.mdx
🚧 Files skipped from review as they are similar to previous changes (4)
- openrag/services/persistence/migrations/milvus/1.add_created_at_temporal_fields.py
- openrag/services/persistence/migrations/milvus/migrate.py
- openrag/services/persistence/migrations/alembic/env.py
- openrag/services/persistence/connection.py
ConnectionManager (services/persistence/connection.py) owns the asyncpg pool lifecycle: initialize() opens the pool with 5-attempt exponential backoff (1s-16s), shutdown() closes it, run_migrations() invokes Alembic synchronously via asyncio.to_thread. Repositories will consume the pool via cm.pool through a lambda pool_getter wrapper (7A.2 onwards). schema.py declares the 7 existing tables (partitions, files, users, oidc_sessions, partition_memberships, workspaces, workspace_files) as metadata-only sa.Table definitions on a shared MetaData. Column types, indexes, unique constraints and check constraints match the legacy SQLAlchemy ORM models exactly so Alembic autogenerate sees zero drift. ORM relationship() calls are dropped - relationships are reconstructed via raw SQL joins in the new repos. RDBConfig gains database, pool_min_size, pool_max_size, command_timeout fields (with sensible defaults). database stays optional - the 7E DI wiring will derive it from VectorDBConfig.collection_name. The new fields are also wired in core/config/loader.py via POSTGRES_DATABASE, POSTGRES_POOL_MIN_SIZE, POSTGRES_POOL_MAX_SIZE, POSTGRES_COMMAND_TIMEOUT environment variables.
…tions Pulled forward from the 7A.4 subsection to keep env.py's metadata import aligned with the new schema.py target. All 8 existing version scripts (initial through OIDC sessions) carry forward unchanged; env.py imports target_metadata from services.persistence.schema instead of the legacy ORM Base.metadata. Old scripts/migrations/alembic/ directory is removed. Updates docs/content/docs/documentation/sql_migration.mdx in a follow-up.
…ations Mechanical path swap after the 7A.4 migration directory move - all references to openrag/scripts/migrations/alembic/ now point at the new openrag/services/persistence/migrations/ location.
Records four decisions from the 7A.1 implementation: pulling the 7A.4 migration move forward, extending RDBConfig with pool/database fields, keeping the database name optional with a construction-time guard, and the schema-vs-ORM diff used as a one-time acceptance check.
EnjoyBacon7
force-pushed
the
refactor/phase-7-storage-persistence-adapters
branch
2 times, most recently
from
May 13, 2026 10:24
dd0ef95 to
0605969
Compare
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@openrag/core/vector_stores/vector_store.py`:
- Around line 30-42: The abstract hybrid_search method on VectorStore
contradicts the design decision to keep hybrid search Milvus-specific; remove
the abstract async def hybrid_search(...) from the VectorStore ABC and add a
public async hybrid_search(self, embedding, query_text, top_k=10,
collection="default", filters=None) implementation to MilvusVectorStore (or the
Milvus-specific subclass) with the same signature and behavior, ensuring other
stores (Pinecone/Weaviate/Qdrant) are not forced to implement it; update
references/tests that called VectorStore.hybrid_search to target
MilvusVectorStore.hybrid_search instead and optionally add a short note in
REFACTORING_DECISION_LOG.md if you change the design.
In `@REFACTORING_DECISION_LOG.md`:
- Around line 726-729: The VectorStore ABC currently declares an abstract
hybrid_search, contradicting the Phase 7B decision that hybrid_search should be
Milvus-specific; change the implementation so VectorStore does NOT declare
hybrid_search and instead implement hybrid_search only on MilvusVectorStore
(remove `@abstractmethod` hybrid_search from VectorStore in
openrag/core/vector_stores/vector_store.py and add/ensure a public
hybrid_search(embedding, query_text, ...) on the MilvusVectorStore concrete
class), or if you intended to keep it on the ABC update the decision log to
state why VectorStore.hybrid_search remains required by all stores (reference
symbols: VectorStore, hybrid_search, MilvusVectorStore).
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: e4db651e-5634-49a4-bff7-6885e5c8a2ca
📒 Files selected for processing (37)
REFACTORING_DECISION_LOG.mddocs/content/docs/documentation/milvus_migration.mdxdocs/content/docs/documentation/sql_migration.mdxopenrag/core/config/infrastructure.pyopenrag/core/config/loader.pyopenrag/core/vector_stores/vector_store.pyopenrag/services/persistence/__init__.pyopenrag/services/persistence/connection.pyopenrag/services/persistence/migrations/alembic/READMEopenrag/services/persistence/migrations/alembic/__init__.pyopenrag/services/persistence/migrations/alembic/alembic.iniopenrag/services/persistence/migrations/alembic/env.pyopenrag/services/persistence/migrations/alembic/schema_helpers.pyopenrag/services/persistence/migrations/alembic/script.py.makoopenrag/services/persistence/migrations/alembic/versions/4add4d260575_initial_migration.pyopenrag/services/persistence/migrations/alembic/versions/__init__.pyopenrag/services/persistence/migrations/alembic/versions/a1b2c3d4e5f6_add_document_relationships.pyopenrag/services/persistence/migrations/alembic/versions/c224d4befe71_add_file_count_and_file_quota.pyopenrag/services/persistence/migrations/alembic/versions/cd642e4502d8_create_users_memberships_tables.pyopenrag/services/persistence/migrations/alembic/versions/cd9b84278028_merge_heads.pyopenrag/services/persistence/migrations/alembic/versions/e7f8a9b0c1d2_add_workspaces.pyopenrag/services/persistence/migrations/alembic/versions/f1a2b3c4d5e6_add_workspace_files_file_id_fk.pyopenrag/services/persistence/migrations/alembic/versions/f5b6c918f741_add_oidc_auth.pyopenrag/services/persistence/migrations/milvus/1.add_created_at_temporal_fields.pyopenrag/services/persistence/migrations/milvus/__init__.pyopenrag/services/persistence/migrations/milvus/migrate.pyopenrag/services/persistence/schema.pyopenrag/services/persistence/test_document_repo.pyopenrag/services/persistence/test_oidc_session_repo.pyopenrag/services/persistence/test_partition_membership_repo.pyopenrag/services/persistence/test_partition_repo.pyopenrag/services/persistence/test_user_repo.pyopenrag/services/persistence/test_workspace_repo.pyopenrag/services/storage/milvus_store.pyopenrag/services/storage/test_milvus_store.pytests/integration/docker-compose.yamltests/integration/test_milvus_store_integration.py
✅ Files skipped from review due to trivial changes (7)
- openrag/services/persistence/init.py
- openrag/services/persistence/test_partition_repo.py
- openrag/services/persistence/test_user_repo.py
- docs/content/docs/documentation/milvus_migration.mdx
- docs/content/docs/documentation/sql_migration.mdx
- openrag/services/persistence/test_workspace_repo.py
- openrag/services/persistence/test_oidc_session_repo.py
🚧 Files skipped from review as they are similar to previous changes (11)
- openrag/core/config/infrastructure.py
- openrag/services/persistence/test_partition_membership_repo.py
- openrag/services/persistence/migrations/milvus/1.add_created_at_temporal_fields.py
- openrag/services/persistence/migrations/milvus/migrate.py
- openrag/core/config/loader.py
- openrag/services/storage/test_milvus_store.py
- openrag/services/persistence/connection.py
- openrag/services/persistence/schema.py
- openrag/services/persistence/migrations/alembic/env.py
- tests/integration/test_milvus_store_integration.py
- openrag/services/storage/milvus_store.py
Ahmath-Gadji
force-pushed
the
refactor/phase-7-storage-persistence-adapters
branch
2 times, most recently
from
May 13, 2026 11:01
1011d14 to
e0c893f
Compare
EnjoyBacon7
force-pushed
the
refactor/phase-7-storage-persistence-adapters
branch
2 times, most recently
from
May 13, 2026 16:18
f7262de to
67b7aa5
Compare
EnjoyBacon7
added a commit
that referenced
this pull request
May 18, 2026
…ager Switch connection.py from stdlib logging to the project Loguru wrapper (get_logger) and convert the %-style log calls to f-strings. Wrap command.upgrade(cfg, head) in a try/except so a failed migration logs a clear message and raises a RuntimeError instead of a bare Alembic stacktrace. Addresses CodeRabbit review on PR #393.
EnjoyBacon7
added a commit
that referenced
this pull request
May 18, 2026
…ager Switch connection.py from stdlib logging to the project Loguru wrapper (get_logger) and convert the %-style log calls to f-strings. Wrap command.upgrade(cfg, head) in a try/except so a failed migration logs a clear message and raises a RuntimeError instead of a bare Alembic stacktrace. Addresses CodeRabbit review on PR #393.
EnjoyBacon7
force-pushed
the
refactor/phase-7-storage-persistence-adapters
branch
from
May 18, 2026 10:14
9376b9d to
fadd6d5
Compare
EnjoyBacon7
changed the title
MilvusVectorStore is a Milvus 2.6 implementation of the narrow VectorStore ABC — pure vector ops, no embedding, no PG, no Ray. Bound to one Milvus collection from VectorDBConfig with a lazy initialize(embedding_dimension) so the DI container can wire it without pulling the embedder into construction. hybrid_search is on the ABC — its query_text argument carries the raw query that backends with a server-side lexical path (Milvus's BM25 Function, Qdrant sparse, …) need alongside the dense embedding, so the contract covers the realistic backend set instead of artificially restricting itself to embedding-only. delete_by_filter stays Milvus-specific (partition-row deletion used by the 7C shim's delete_partition flow) — no need to widen the cross-store contract for it yet. Five design decisions deviated from / extended the phase 7 plan (hybrid as ABC method, lazy dim wiring, strict collection-vs-partition separation aimed at the SaaS end-state, dropping the unjustified reconnect logic against pymilvus 2.6, and a deliberately narrow store surface) — all recorded in REFACTORING_DECISION_LOG.md so the rationale survives future passes.
EnjoyBacon7
added a commit
that referenced
this pull request
May 18, 2026
…ager Switch connection.py from stdlib logging to the project Loguru wrapper (get_logger) and convert the %-style log calls to f-strings. Wrap command.upgrade(cfg, head) in a try/except so a failed migration logs a clear message and raises a RuntimeError instead of a bare Alembic stacktrace. Addresses CodeRabbit review on PR #393.
EnjoyBacon7
force-pushed
the
refactor/phase-7-storage-persistence-adapters
branch
from
May 18, 2026 15:16
e1a7a06 to
7bad4bf
Compare
…ations/alembic Carves out a backend-specific subdirectory so future migration sets (notably Milvus, landing in the next commit) can sit beside alembic in the same namespace — services/persistence/migrations/alembic for Postgres, services/persistence/migrations/milvus for Milvus. Pure structural move: env.py, alembic.ini, schema_helpers.py, script.py.mako, README, and the eight version scripts all relocate unchanged. env.py's metadata import (services.persistence.schema) keeps resolving via the existing pythonpath setup; alembic.ini's %(here)s/ script_location and prepend_sys_path = . both remain valid relative to their new location. The SQL migration guide is updated in the same commit to point callers at the new alembic.ini path.
…ions/milvus Lands Milvus migration assets in the unified migrations namespace alongside alembic (previous commit), so all schema-evolution scripts live under a single services/persistence/migrations/ root with one sub-directory per backend. The generic runner (migrate.py) and the v0→1 temporal-fields migration carry forward unchanged in behaviour. SCHEMA_VERSION_PROPERTY_KEY now imports from services.storage.milvus_store (the new home for the constant) instead of the legacy components.indexer.vectordb.vectordb; both modules define an identical value today, and the legacy module is on the Phase 9 deletion path. The Milvus migration guide is updated in the same commit — all references to openrag/scripts/migrations/milvus/ now point at the new nested location.
Decision 1 captures the choice of nested backend subdirectories over either the spec-by-symmetry split-root layout (services/persistence/migrations/ + services/storage/migrations/) or a flat single dir. Decision 2 captures the import-source swap for SCHEMA_VERSION_PROPERTY_KEY from the legacy components.indexer.vectordb.vectordb to the new services.storage.milvus_store, eliminating a backwards dependency into the Phase-9 deletion path.
Two test modules:
* openrag/services/storage/test_milvus_store.py — 52 unit tests with both
pymilvus clients monkeypatched. Cover _format_value escapes,
_build_filter_expr partition/list/expr semantics, _resolve_collection
ABC discipline, ID round-trip, chunk-order metadata linking,
_chunk_to_entity layering (typed > metadata, omit-None policy), and
the empty/wildcard guards on upsert/delete/delete_by_filter. Also
asserts hybrid_search refuses on a store built with hybrid_search=False.
Colocated with the SUT — runs by default with no infra.
* tests/integration/test_milvus_store_integration.py — 14 end-to-end
tests against a real Milvus 2.6, gated by the integration marker and
the OPENRAG_TEST_VDB_HOST env (default localhost:19530). Covers
initialize idempotence, ensure_collection dimension-mismatch guard,
upsert/dense-search/hybrid-search round trip, partition-filtered
search, delete-by-filter (partition scope + wildcard guard),
query_ids/query_chunks round trip, drop_collection + re-initialize.
Lives under tests/integration/ per the Phase 13C target test layout
(tests/{unit,integration,load}/) documented in
docs/refactoring/REFACTORING_STRATEGY_v1.md. pytest.ini's
testpaths=openrag means it stays out of the default unit run and only
fires when explicitly targeted or invoked by CI.
Pure-logic tests use monkeypatch on the loaded module rather than
unittest.mock.patch string paths because the project's pythonpath
setup registers the same source file under two module names
(services.storage.milvus_store vs openrag.services.storage.milvus_store),
which makes string-based patch paths fragile.
The integration fixture deliberately reads OPENRAG_TEST_VDB_HOST rather
than VDB_HOST: pymilvus.settings auto-loads .env at import time, which
would inject the docker-network hostname "milvus" into a host-side test
run. The dedicated env keeps runtime config and test config independent.
One test documents (rather than asserts against) an asymmetry: the
_iter_query path in query_chunks_by_filter leaks the dense vector when
called with "*" output_fields, while the search path strips it via
_SEARCH_RESULT_DROPPED_KEYS. Tracked as a follow-up.
A dedicated GitHub Actions job (.github/workflows/integration_tests.yml)
brings up a minimal Milvus stack via tests/integration/docker-compose.yaml
(etcd + minio + milvus 2.6, with 19530 and 9091 exposed to the host)
and runs pytest tests/integration/ -m integration. The compose file
deliberately does not pull in openrag/mock-vllm/rdb — only what the
integration suite touches — so the job stays fast and independent of
the heavier api-tests stack.
Decomposes PartitionFileManager (utils.py, 1095 LOC, 51 methods) into
fifteen asyncpg-backed Pg*Repository classes under services/persistence/.
Five are real (document, user, partition, oidc_session, workspace) and
cover every method moving over from the legacy god-object; ten are stubs
(job, chunk, prompt, conversation, audit_log, idempotency, entity,
topic_tag, model_endpoint, preset) that raise StubRepositoryError until
their tables exist.
Each real repo exposes a dual surface: the typed port ABC methods used
by future orchestrators, plus the legacy method names with the original
positional signatures the 7C shim will call. Repos receive a pool_getter
callable rather than a pool reference so they always see the live pool
when the connection manager is reinitialised in tests.
Supporting changes:
* Add a Workspace domain model and WorkspaceRepository port (omitted in
phase 4) plus a workspace_repo property on CatalogStore.
* Extend OIDCSession with three optional bytes fields for the Fernet
encrypted IdP tokens so create_session can carry them through the
port without a side channel.
* Register json/jsonb codecs in ConnectionManager._init_connection so
repositories read JSON columns as Python dicts and write dicts back
without per-call json.dumps boilerplate.
* Log the six 7A.2 design decisions in REFACTORING_DECISION_LOG.md
(new port, OIDC token bytes, dual surface, user-model field
reconciliation, stub-repo error class, connection-level codec).
PostgresStore composes the Phase 7A.1 connection manager with every Phase 7A.2 repository implementation into a single CatalogStore root. Orchestrators and the upcoming 7C shim consume it through the port ABC; the fifteen repos are wired once in __init__ and share one pool_getter callable so a shutdown/initialize cycle in tests rebuilds the pool without touching the repos. The class lives in services/storage/ alongside MilvusStore rather than inside services/persistence/, matching the layer split the rest of the phase 7 plan assumes: storage is the high-level adapter surface, persistence owns the row-level implementations. initialize() opens the pool, then runs Alembic to head; a run_migrations=False flag is exposed for fast unit tests against an already-migrated database. A pool property gives Phase 8 orchestrators a transactional escape hatch without polluting the CatalogStore ABC with asyncpg specifics. The four 7A.3 design decisions (placement, eager repo construction, single-entry-point lifecycle, ABC-free pool escape hatch) are recorded in REFACTORING_DECISION_LOG.md.
ServiceContainer gains an optional Settings argument that, when provided, builds the Phase 7A.3 PostgresStore via a dedicated factory and exposes every catalog repository as a typed property on the container. The existing no-arg construction path is preserved so the registry-only inference tests keep working unchanged; storage accessors raise a clear RuntimeError when reached without Settings. di/repositories.create_catalog_store centralises the legacy "partitions_for_collection_<collection>" database-name fallback so DI wiring code never has to mention the prefix; an explicit rdb.database still wins. di/vector_stores.create_vector_store is a fail-loud Phase 7B stub with the right signature, so Phase 8 code can already import the symbol — only the body changes when MilvusVectorStore lands. The container exposes async initialize() and shutdown() that open and close the asyncpg pool plus Alembic migrations, all behind the CatalogStore ABC. Per-repo accessors return the same instance the store holds, so orchestrator injection in Phase 8 sees a single source of truth. A new di/test_container.py covers the wiring end to end: legacy no-arg container, Settings-driven container, catalog-store ABC conformance, database-name derivation and override, all fifteen repo properties returning port-typed instances that identity-match the store's, and the vector store factory + property raising the Phase 7B message. Also aligns services/persistence/ and services/storage/ imports to the project's short-form convention (from core.X, not from openrag.core.X), fixing a dual-import isinstance bug surfaced by the new container test: pytest sets pythonpath=./openrag so the same module is reachable under two names and Python sees the two as distinct classes. CLAUDE.md already mandates the short form. The four 7E design decisions (optional Settings, centralised database fallback, fail-loud vector-store stub, import-convention alignment) are recorded in REFACTORING_DECISION_LOG.md.
Phase 7F lands a real-Postgres integration suite under
tests/integration/ for the Phase 7A.3 PostgresStore composite plus the
five real repositories decomposed from PartitionFileManager. The suite
is kept outside pytest.ini's testpaths so the default unit run stays
infra-free; operators invoke it with uv run pytest tests/integration/.
Fixtures (conftest.py):
* Discover a Postgres DSN via POSTGRES_TEST_ADMIN_DSN, falling back
to the local docker-compose rdb container at 172.21.0.4. Auto-skip
the whole suite when neither is reachable.
* Session-scoped: drop+recreate openrag_phase7_test, build a
PostgresStore, run Alembic to head.
* Autouse function-scoped: TRUNCATE the seven user-modifiable tables
with RESTART IDENTITY CASCADE so each test starts at id=1.
* loop_scope=session on every async fixture and test marker so the
pool stays on the one event loop pytest-asyncio creates per
session.
Tests:
* test_postgres_store.py — 20 cases: CatalogStore ABC conformance,
pool lifecycle, all 15 repo properties returning the expected
PgX class, migration idempotency, initialise/shutdown round-trip
on an isolated store.
* test_partition_repo.py — create/list/exists/delete/idempotent
create.
* test_document_repo.py — CRUD, partition-scoped list, status fold
into metadata JSON, metadata-merge semantics, delete-by-partition
count.
* test_user_repo.py — CRUD, email lowercasing, by-email/by-external-
id lookup, legacy token mint/regenerate/hash-lookup, partition
membership assign/idempotent-upsert/remove.
* test_oidc_session_repo.py — round-trip encrypted token bytes,
revoked/expired rows hidden from get_by_token_hash, revoke-by-sid
and revoke-by-user counts, 7-day retention window before
delete_expired removes a row.
* test_workspace_repo.py — workspace CRUD, file membership
(idempotent add, partition-scoped get_file_workspaces, orphan
detection on delete, remove_file_from_all_workspaces).
* test_stores.py — cross-store full-cycle marked xfail(strict=True)
until Phase 7B's MilvusVectorStore lands.
73 integration tests pass; the unit suite is unchanged at 655 passed.
Two real bugs surfaced while writing the suite are fixed in the same
commit so the new tests land green:
* PgOIDCSessionRepository.delete_expired built a tz-aware cutoff and
bound it against the tz-naive session_expires_at column, which
asyncpg rejects at encode time. Strip tzinfo before the query.
* services/persistence/migrations/env.py unconditionally clobbered
sqlalchemy.url with one derived from load_config(), ignoring the
DSN ConnectionManager.run_migrations() had just set. The first run
of the test suite therefore tried to resolve docker's rdb
hostname from the host. Defer to the preset URL when it's not the
alembic.ini placeholder.
The five 7F design decisions (placement under tests/integration/,
per-test truncate vs per-test DB, session-scoped event loop, strict
xfail for the cross-store test, fixing the two surfaced bugs in this
commit rather than splitting them out) are recorded in
REFACTORING_DECISION_LOG.md.
Person A's 7E commit was drafted while Phase 7B was still pending, so
the di/vector_stores.create_vector_store factory was a fail-loud stub
("Phase 7B deliverable") and the di/container.ServiceContainer.vector_store
property raised NotImplementedError at access time. Phase 7B has since
landed on this branch (commit 5617cb3), so the factory swaps to its
final body — return MilvusVectorStore(settings.vectordb) — and the
container caches the instance at construction so repeated property reads
do not open a new pymilvus gRPC channel each time.
services/storage/__init__.py now exports MilvusVectorStore alongside
PostgresStore and the docstring drops the "Phase 7B placeholder"
language.
The di/test_container.TestVectorStoreNotYetWired class renames to
TestVectorStoreWiring and asserts isinstance against MilvusVectorStore.
An autouse module-level fixture patches MilvusClient / AsyncMilvusClient
with MagicMock so the unit suite does not need Milvus reachable — the
real coverage lives in tests/integration/test_milvus_store_integration.py.
tests/integration/test_stores.py keeps its strict xfail because the
cross-store fixture still needs a Milvus-store companion in the
postgres_store conftest — that lands in Phase 7C. The xfail reason and
the matching REFACTORING_DECISION_LOG.md 7F#4 entry are updated so
"Phase 7C handoff" is the named blocker rather than "Phase 7B not
landed".
Decision log entries 7E#3 and 7F#4 reflect the post-rebase reality.
…e ABC The Phase 7B decision log still recorded the original "ABC stays embedding-only; hybrid lives on MilvusVectorStore" choice, but the hybrid_search abstract method has since been added to the VectorStore ABC (see the 7B refactor commit). Rewrite §1 to capture the reversal and the SaaS-backend reasoning that drove it (Qdrant / Weaviate / Pinecone / OpenSearch all have a hybrid path; treating it as a Milvus quirk would force an isinstance leak in the Phase-8 retrieval orchestrator). Trim §5's "store surface" list to reflect that hybrid_search is no longer Milvus-only — the only remaining Milvus-only public method on the store is delete_by_filter. Resolves the CodeRabbit pull-request comments flagging the "abstract method contradicts documented decision" inconsistency in vector_store.py:42 and REFACTORING_DECISION_LOG.md:909.
CI's ruff format --check caught five files that hadn't gone through the formatter: openrag/di/container.py, openrag/services/persistence/document_repo.py, openrag/services/persistence/user_repo.py, tests/integration/conftest.py, tests/integration/test_oidc_session_repo.py. Whitespace and line-wrap adjustments only — no behaviour change. Lint job was the only red on the PR; tests/api-tests/milvus-integration were still in progress.
Replace the legacy PartitionFileManager+MilvusDB actor body in vectordb.py with a thin shim that composes MilvusVectorStore and PostgresStore behind the existing Ray actor surface. Wire dependencies to initialise the actor's stores inside its own asyncio loop so the asyncpg pool is bound to the actor loop, and override the restart hook so a recreated actor goes through the same initialize step.
…he shim Record the trade-off: the per-collection database bootstrap stayed on the MilvusDB shim (not in PostgresStore.initialize) to honour Phase 7C's "vectordb.py only" scope contract. Flagged as a follow-up to move into PostgresStore / ConnectionManager during 7E DI wiring so the helper can be deleted in Phase 9 without regression.
PgPartitionRepository.delete_partition assumed files.partition_name cascaded at the DB level, but the FK has no ON DELETE CASCADE — the legacy ORM relied on SQLAlchemy's Python-side cascade. Issue an explicit DELETE FROM files inside the same transaction before dropping the partition row, matching the legacy bookkeeping. Add an integration test that creates files under an uploader and asserts the cascade plus the per-uploader file_count decrement.
`ray.get()` is synchronous and returns the resolved value, so wrapping it with `await` tried to await a `str` and raised TypeError. Drop the `ray.get()` wrapper and await the ObjectRef directly — the async-actor pattern used elsewhere in this module.
…ager Switch connection.py from stdlib logging to the project Loguru wrapper (get_logger) and convert the %-style log calls to f-strings. Wrap command.upgrade(cfg, head) in a try/except so a failed migration logs a clear message and raises a RuntimeError instead of a bare Alembic stacktrace. Addresses CodeRabbit review on PR #393.
… migration storm get_vectordb() is a FastAPI Depends used by 34 endpoints; it re-invokes the Vectordb actor's initialize on every request. The shim delegates to PostgresStore.initialize(), which re-ran a full Alembic command.upgrade each time because, unlike MilvusVectorStore (guarded by _loaded), PostgresStore had no idempotency guard. Under concurrent api-tests load the repeated migrations starved the Postgres pool and surfaced as intermittent 500s on indexer/workspace endpoints. Guard initialize() with an _initialized flag (reset on shutdown) so re-running on a hot actor is a no-op, matching the documented contract in get_vectordb(). Add an integration regression test asserting three initialize() calls trigger run_migrations exactly once.
Restores the 7A.2 one-repo-per-entity layout: partition memberships get their own port and asyncpg implementation instead of riding on UserRepository. - new core/ports/partition_membership_repo.py (PartitionMembershipRepository ABC) - new services/persistence/partition_membership_repo.py (PgPartitionMembershipRepository) - moved the 6 clean methods + 6 legacy shim methods + _row_to_user_partition off PgUserRepository - UserRepository keeps reading partition_memberships via _fetch_memberships only to hydrate the User aggregate's partitions field (read-only denormalisation, not membership management) - expose membership_repo on CatalogStore, PostgresStore and the DI container - repoint the 7 membership call sites in the vectordb.py shim - split the membership tests into tests/integration/test_partition_membership_repo.py Unit suite unchanged (714 passed); no caller outside the shim touched.
EnjoyBacon7
force-pushed
the
refactor/phase-7-storage-persistence-adapters
branch
from
May 18, 2026 15:35
7bad4bf to
79afed0
Compare
EnjoyBacon7
added a commit
that referenced
this pull request
May 26, 2026
…ager Switch connection.py from stdlib logging to the project Loguru wrapper (get_logger) and convert the %-style log calls to f-strings. Wrap command.upgrade(cfg, head) in a try/except so a failed migration logs a clear message and raises a RuntimeError instead of a bare Alembic stacktrace. Addresses CodeRabbit review on PR #393.
This was referenced May 29, 2026
This was referenced Jun 25, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Phase 7 — Storage & Persistence Adapters
Decomposes the
MilvusDBRay god object (~2,850 LOC: vector search + all PostgreSQL metadata + embedding, in one@ray.remoteclass) into two clean, single-responsibility adapters behind the Phase 4 ports:MilvusVectorStore(services/storage/milvus_store.py) — vector operations only, implements theVectorStoreABCPostgresStore(services/storage/postgres_store.py) — all metadata, implements theCatalogStoreABC, composing 6 asyncpg repositoriesThe old
MilvusDBactor is not removed — it is reduced to a thin shim that delegates to the new stores, so every one of the ~76 existing.remote()call sites keeps working unchanged. Full removal is deferred to Phase 9.What's in this PR
7A — PostgreSQL persistence layer (SQLAlchemy ORM → asyncpg)
connection.py—ConnectionManagerwith pooled asyncpg, retry/backoff, idempotent migration guardschema.py— 7 ORM models rewritten as metadata-onlysa.Tabledefs (Alembic autogenerate only, no ORM)PartitionFileManager(51 methods):document,user,partition,partition_membership,oidc_session,workspacejob,chunk,prompt,conversation,audit_log,idempotency) — signatures only, from the port ABCsPostgresStorecomposite owns theConnectionManagerlifecycle and exposes each repo as a propertyservices/persistence/migrations/(8 existing scripts carried forward unchanged)7B — Milvus vector store
MilvusVectorStoreextracted fromMilvusDB's vector ops, rewritten against theVectorStoreABCRRFRanker)ConnectionNotExistException(distinct from the Phase 6 circuit breaker)7C — God-object shim
MilvusDB.__init__now builds the new stores internally; pure-PG and pure-Milvus methods delegate; the 6 cross-cutting methods (async_add_documents,delete_file,delete_partition,async_search,upsert_file_metadata,add_documents_for_existing_file) call both stores inline until Phase 8 orchestrators landinitialize()made idempotent to stop a per-request migration storm7E — DI wiring
di/vector_stores.py,di/repositories.pyfactories;ServiceContainerwires both stores and exposes individual repos for Phase 8 orchestrator injection7F — Tests
MilvusVectorStoreunit + live integration tests, withtests/integration/conftest and docker-composeScope & risk
This is the highest-risk refactoring phase (god-object coupling, dual database, sync ORM → async driver rewrite). The mitigation is the shim: the system runs at every commit and no caller is rewritten here — that is Phase 8's job.
Out of scope (deferred by design)
vectordb.py→ Phase 9Verification
services/persistence/orservices/storage/(metadata-only)pool_getter; no module-level config72 files changed, ~7.4k insertions. Base:
refactor/hexagonal.