refactor: organize model registry and infrastructure packages

mlflow_mongodb/__init__.py

@@ -1,6 +1,6 @@
 """MongoDB model registry store plugin for MLflow."""
 
-from mlflow_mongodb.model_registry_store import MongoDBModelRegistryStore
+from mlflow_mongodb.model_registry.store import MongoDBModelRegistryStore
 
 __all__ = ["MongoDBModelRegistryStore"]
 __version__ = "0.1.0.dev0"

mlflow_mongodb/infrastructure/__init__.py

@@ -0,0 +1 @@
+"""Shared infrastructure for the MongoDB MLflow stores."""

mlflow_mongodb/model_registry/__init__.py

@@ -0,0 +1 @@
+"""MongoDB-backed MLflow model-registry store."""

mlflow_mongodb/model_registry/errors.py

@@ -0,0 +1,17 @@
+"""Domain exceptions raised by the model-registry store and repositories."""
+
+
+class ModelVersionAlreadyExistsError(Exception):
+    """Raised when a model version number is already stored for a registered model."""
+
+
+class ModelVersionNotFoundError(Exception):
+    """Raised when a model version is not stored for a registered model."""
+
+
+class RegisteredModelAlreadyExistsError(Exception):
+    """Raised when a registered model name is already stored."""
+
+
+class RegisteredModelNotFoundError(Exception):
+    """Raised when a registered model name is not stored."""

mlflow_mongodb/model_registry/repositories/__init__.py

@@ -0,0 +1,25 @@
+"""MongoDB repositories owned by the model-registry store."""
+
+from mlflow_mongodb.model_registry.repositories.model_versions import (
+    ModelVersionFilter,
+    ModelVersionOrder,
+    ModelVersionPage,
+    ModelVersionRepository,
+)
+from mlflow_mongodb.model_registry.repositories.registered_models import (
+    RegisteredModelFilter,
+    RegisteredModelOrder,
+    RegisteredModelPage,
+    RegisteredModelRepository,
+)
+
+__all__ = [
+    "ModelVersionFilter",
+    "ModelVersionOrder",
+    "ModelVersionPage",
+    "ModelVersionRepository",
+    "RegisteredModelFilter",
+    "RegisteredModelOrder",
+    "RegisteredModelPage",
+    "RegisteredModelRepository",
+]

mlflow_mongodb/repositories/model_versions.py → mlflow_mongodb/model_registry/repositories/model_versions.py

@@ -16,16 +16,20 @@
 from pymongo.database import Database
 from pymongo.errors import DuplicateKeyError
 
-from mlflow_mongodb.repositories._helpers import (
+from mlflow_mongodb.infrastructure._array_updates import (
     build_remove_array_element_update,
     build_replace_array_element_pipeline,
 )
-from mlflow_mongodb.repositories.types import (
+from mlflow_mongodb.infrastructure.settings import MongoDBSettings
+from mlflow_mongodb.model_registry.errors import (
+    ModelVersionAlreadyExistsError,
+    ModelVersionNotFoundError,
+)
+from mlflow_mongodb.model_registry.types import (
     ModelVersionRecord,
     ModelVersionSearchResult,
     RegisteredModelRecord,
 )
-from mlflow_mongodb.settings import MongoDBSettings
 
 logger = logging.getLogger(__name__)
 
@@ -61,14 +65,6 @@ class ModelVersionPage:
     has_more: bool
 
 
-class ModelVersionAlreadyExistsError(Exception):
-    """Raised when a model version number is already stored for a registered model."""
-
-
-class ModelVersionNotFoundError(Exception):
-    """Raised when a model version is not stored for a registered model."""
-
-
 class ModelVersionRepository:
     """Read and write model-version documents in MongoDB."""
 

mlflow_mongodb/repositories/registered_models.py → mlflow_mongodb/model_registry/repositories/registered_models.py

@@ -13,16 +13,20 @@
 from pymongo.database import Database
 from pymongo.errors import DuplicateKeyError
 
-from mlflow_mongodb.repositories._helpers import (
+from mlflow_mongodb.infrastructure._array_updates import (
     build_remove_array_element_update,
     build_replace_array_element_pipeline,
 )
-from mlflow_mongodb.repositories.types import (
+from mlflow_mongodb.infrastructure.settings import MongoDBSettings
+from mlflow_mongodb.model_registry.errors import (
+    RegisteredModelAlreadyExistsError,
+    RegisteredModelNotFoundError,
+)
+from mlflow_mongodb.model_registry.types import (
     ModelVersionRecord,
     RegisteredModelDetails,
     RegisteredModelRecord,
 )
-from mlflow_mongodb.settings import MongoDBSettings
 
 logger = logging.getLogger(__name__)
 
@@ -53,14 +57,6 @@ class RegisteredModelPage:
     has_more: bool
 
 
-class RegisteredModelAlreadyExistsError(Exception):
-    """Raised when a registered model name is already stored."""
-
-
-class RegisteredModelNotFoundError(Exception):
-    """Raised when a registered model name is not stored."""
-
-
 class RegisteredModelRepository:
     """Store registered-model documents in MongoDB."""
 

mlflow_mongodb/model_registry_store.py → mlflow_mongodb/model_registry/store.py

@@ -54,22 +54,26 @@
 from pymongo import MongoClient
 from pymongo.errors import ConfigurationError
 
-from mlflow_mongodb.repositories import (
+from mlflow_mongodb.infrastructure.settings import MongoDBSettings
+from mlflow_mongodb.model_registry.errors import (
     ModelVersionAlreadyExistsError,
-    ModelVersionFilter,
     ModelVersionNotFoundError,
+    RegisteredModelAlreadyExistsError,
+    RegisteredModelNotFoundError,
+)
+from mlflow_mongodb.model_registry.repositories import (
+    ModelVersionFilter,
     ModelVersionOrder,
-    ModelVersionRecord,
     ModelVersionRepository,
-    RegisteredModelAlreadyExistsError,
-    RegisteredModelDetails,
     RegisteredModelFilter,
-    RegisteredModelNotFoundError,
     RegisteredModelOrder,
-    RegisteredModelRecord,
     RegisteredModelRepository,
 )
-from mlflow_mongodb.settings import MongoDBSettings
+from mlflow_mongodb.model_registry.types import (
+    ModelVersionRecord,
+    RegisteredModelDetails,
+    RegisteredModelRecord,
+)
 
 logger = logging.getLogger(__name__)
 

mlflow_mongodb/repositories/types.py → mlflow_mongodb/model_registry/types.py

@@ -1,4 +1,4 @@
-"""Shared persistence DTOs for the MongoDB repositories."""
+"""Persistence DTOs owned by the model-registry store."""
 
 from collections.abc import Mapping
 from dataclasses import dataclass

pyproject.toml

@@ -52,7 +52,7 @@ ignore = [
 known-first-party = ["mlflow_mongodb"]
 
 [tool.ruff.lint.per-file-ignores]
-"mlflow_mongodb/model_registry_store.py" = ["import-private-name"]
+"mlflow_mongodb/model_registry/store.py" = ["import-private-name"]
 
 [tool.ruff.lint.pydocstyle]
 convention = "google"

setup.py

@@ -32,8 +32,8 @@
     },
     entry_points={
         "mlflow.model_registry_store": [
-            "mongodb=mlflow_mongodb.model_registry_store:MongoDBModelRegistryStore",
-            "mongodb+srv=mlflow_mongodb.model_registry_store:MongoDBModelRegistryStore",
+            "mongodb=mlflow_mongodb.model_registry.store:MongoDBModelRegistryStore",
+            "mongodb+srv=mlflow_mongodb.model_registry.store:MongoDBModelRegistryStore",
         ],
     },
     classifiers=[

tests/functional/test_indexes.py

@@ -1,7 +1,7 @@
 """Functional assertions for production MongoDB index definitions."""
 
 from mlflow_mongodb import MongoDBModelRegistryStore
-from mlflow_mongodb.repositories import (
+from mlflow_mongodb.model_registry.repositories import (
     ModelVersionRepository,
     RegisteredModelRepository,
 )

tests/functional/test_model_version_contract.py

@@ -20,7 +20,7 @@ def test_model_version_create_update_and_stage_validation_contract(
     name = "model-version-metadata-contract"
     clock = {"now": 1_700_001_000_000}
     monkeypatch.setattr(
-        "mlflow_mongodb.model_registry_store.get_current_time_millis",
+        "mlflow_mongodb.model_registry.store.get_current_time_millis",
         lambda: clock["now"],
     )
     store.create_registered_model(name)
@@ -129,7 +129,7 @@ def test_transition_with_archiving_rejects_inactive_stages_and_archives_peers(
 
     transition_timestamp = 1_700_002_000_000
     monkeypatch.setattr(
-        "mlflow_mongodb.model_registry_store.get_current_time_millis",
+        "mlflow_mongodb.model_registry.store.get_current_time_millis",
         lambda: transition_timestamp,
     )
     store.transition_model_version_stage(name, version_three.version, "STAGING", True)

tests/functional/test_registered_model_contract.py

@@ -23,7 +23,7 @@ def test_registered_model_create_get_and_update_contract(
 ):
     clock = {"now": 1_700_000_000_000}
     monkeypatch.setattr(
-        "mlflow_mongodb.model_registry_store.get_current_time_millis",
+        "mlflow_mongodb.model_registry.store.get_current_time_millis",
         lambda: clock["now"],
     )
 
@@ -93,7 +93,7 @@ def test_rename_registered_model_preserves_versions_and_rejects_conflicts(
 
     rename_timestamp = 1_700_000_100_000
     monkeypatch.setattr(
-        "mlflow_mongodb.model_registry_store.get_current_time_millis",
+        "mlflow_mongodb.model_registry.store.get_current_time_millis",
         lambda: rename_timestamp,
     )
 

tests/functional/test_search_contract.py

@@ -113,7 +113,7 @@ def test_search_model_versions_orders_and_paginates_stably(
 ):
     clock = {"now": 1_700_003_000_000}
     monkeypatch.setattr(
-        "mlflow_mongodb.model_registry_store.get_current_time_millis",
+        "mlflow_mongodb.model_registry.store.get_current_time_millis",
         lambda: clock["now"],
     )
     names = [
@@ -294,7 +294,7 @@ def test_search_registered_models_orders_and_paginates_stably(
 ):
     clock = {"now": 1_700_004_000_000}
     monkeypatch.setattr(
-        "mlflow_mongodb.model_registry_store.get_current_time_millis",
+        "mlflow_mongodb.model_registry.store.get_current_time_millis",
         lambda: clock["now"],
     )
 

tests/unit/conftest.py

@@ -10,14 +10,16 @@
 from mlflow.tracking.client import MlflowClient
 
 from mlflow_mongodb import MongoDBModelRegistryStore
-from mlflow_mongodb.repositories import (
-    ModelVersionRecord,
+from mlflow_mongodb.model_registry.repositories import (
     ModelVersionRepository,
+    RegisteredModelRepository,
+)
+from mlflow_mongodb.model_registry.types import (
+    ModelVersionRecord,
     ModelVersionTagRecord,
     RegisteredModelAliasRecord,
     RegisteredModelDetails,
     RegisteredModelRecord,
-    RegisteredModelRepository,
     RegisteredModelTagRecord,
 )
 
@@ -48,7 +50,7 @@ def tracking_client():
 def store(monkeypatch, registered_model_repository, model_version_repository, tracking_client):
     """Return a store wired to mocks so unit tests never connect to MongoDB."""
     monkeypatch.setattr(
-        "mlflow_mongodb.model_registry_store.get_current_time_millis",
+        "mlflow_mongodb.model_registry.store.get_current_time_millis",
         lambda: FIXED_TIMESTAMP,
     )
 

tests/unit/test_prompt_store.py

@@ -4,7 +4,7 @@
 from mlflow.exceptions import MlflowException
 from mlflow.prompt.constants import IS_PROMPT_TAG_KEY, PROMPT_TEXT_TAG_KEY
 
-from mlflow_mongodb.repositories import (
+from mlflow_mongodb.model_registry.errors import (
     RegisteredModelAlreadyExistsError,
     RegisteredModelNotFoundError,
 )

tests/unit/test_store.py

@@ -15,7 +15,7 @@
 )
 
 from mlflow_mongodb import MongoDBModelRegistryStore
-from mlflow_mongodb.repositories import (
+from mlflow_mongodb.model_registry.errors import (
     ModelVersionAlreadyExistsError,
     ModelVersionNotFoundError,
     RegisteredModelAlreadyExistsError,
@@ -33,7 +33,7 @@ def test_tracking_client_is_lazy_and_cached(monkeypatch):
     tracking_client = Mock()
     client_factory = Mock(return_value=tracking_client)
     monkeypatch.setattr(
-        "mlflow_mongodb.model_registry_store.MlflowClient",
+        "mlflow_mongodb.model_registry.store.MlflowClient",
         client_factory,
     )
 

tests/unit/test_store_configuration.py

@@ -6,7 +6,7 @@
 from mlflow.exceptions import MlflowException
 
 from mlflow_mongodb import MongoDBModelRegistryStore
-from mlflow_mongodb.settings import MongoDBSettings
+from mlflow_mongodb.infrastructure.settings import MongoDBSettings
 
 
 def test_mongo_client_requires_registry_uri():

tests/unit/test_store_search.py

@@ -7,7 +7,7 @@
 from mlflow.utils.search_utils import SearchModelUtils, SearchModelVersionUtils, SearchUtils
 
 from mlflow_mongodb import MongoDBModelRegistryStore
-from mlflow_mongodb.repositories import (
+from mlflow_mongodb.model_registry.repositories import (
     ModelVersionFilter,
     ModelVersionOrder,
     RegisteredModelFilter,