py: fix order by ID (#355)

* api: fix ID field in orderBy

Fixes: #353

Signed-off-by: Isabella do Amaral <[email protected]>

* py: please mypy

Signed-off-by: Isabella do Amaral <[email protected]>

* add stubs for dateutil

Signed-off-by: Isabella do Amaral <[email protected]>

---------

Signed-off-by: Isabella do Amaral <[email protected]>

api/openapi/model-registry.yaml

@@ -1348,7 +1348,7 @@ components:
       enum:
         - CREATE_TIME
         - LAST_UPDATE_TIME
-        - Id
+        - ID
       type: string
     Artifact:
       oneOf:
@@ -1661,7 +1661,7 @@ components:
       explode: true
       examples:
         orderBy:
-          value: Id
+          value: ID
       name: orderBy
       description: Specifies the order by criteria for listing entities.
       schema:

clients/python/noxfile.py

@@ -42,7 +42,10 @@ def lint(session: Session) -> None:
 def mypy(session: Session) -> None:
     """Type check using mypy."""
     session.install(".")
-    session.install("mypy")
+    session.install(
+        "mypy",
+        "types-python-dateutil",
+    )
 
     session.run("mypy", "src/model_registry")
 

clients/python/pyproject.toml

@@ -49,6 +49,7 @@ mypy = "^1.7.0"
 pytest-asyncio = ">=0.23.7,<0.25.0"
 requests = "^2.32.2"
 black = "^24.4.2"
+types-python-dateutil = "^2.9.0.20240906"
 
 [tool.coverage.run]
 branch = true

clients/python/src/model_registry/_client.py

@@ -32,7 +32,7 @@ def __init__(
         *,
         author: str,
         is_secure: bool = True,
-        user_token: bytes | None = None,
+        user_token: str | None = None,
         custom_ca: str | None = None,
     ):
         """Constructor.
@@ -44,8 +44,8 @@ def __init__(
         Keyword Args:
             author: Name of the author.
             is_secure: Whether to use a secure connection. Defaults to True.
-            user_token: The PEM-encoded user token as a byte string. Defaults to content of path on envvar KF_PIPELINES_SA_TOKEN_PATH.
-            custom_ca: Path to the PEM-encoded root certificates as a byte string. Defaults to path on envvar CERT.
+            user_token: The PEM-encoded user token as a string. Defaults to content of path on envvar KF_PIPELINES_SA_TOKEN_PATH.
+            custom_ca: Path to the PEM-encoded root certificates as a string. Defaults to path on envvar CERT.
         """
         import nest_asyncio
 

clients/python/src/model_registry/core.py

@@ -42,7 +42,7 @@ def secure_connection(
         server_address: str,
         port: int = 443,
         *,
-        user_token: bytes,
+        user_token: str,
         custom_ca: str | None = None,
     ) -> ModelRegistryAPIClient:
         """Constructor.
@@ -52,7 +52,7 @@ def secure_connection(
             port: Server port. Defaults to 443.
 
         Keyword Args:
-            user_token: The PEM-encoded user token as a byte string.
+            user_token: The PEM-encoded user token as a string.
             custom_ca: The path to a PEM-
         """
         return cls(
@@ -68,14 +68,14 @@ def insecure_connection(
         cls,
         server_address: str,
         port: int,
-        user_token: bytes | None = None,
+        user_token: str | None = None,
     ) -> ModelRegistryAPIClient:
         """Constructor.
 
         Args:
             server_address: Server address.
             port: Server port.
-            user_token: The PEM-encoded user token as a byte string.
+            user_token: The PEM-encoded user token as a string.
         """
         return cls(
             Configuration(host=f"{server_address}:{port}", access_token=user_token)

clients/python/src/model_registry/types/pager.py

@@ -42,8 +42,8 @@ def restart(self) -> Pager[T]:
         This keeps the current options and page function, but resets the internal state.
         """
         # as MLMD loops over pages, we need to keep track of the first page or we'll loop forever
-        self._start = None
-        self._current_page = None
+        self._start: str | None = None
+        self._current_page: list[T] | None = None
         # tracks the next item on the current page
         self._i = 0
         self.options.next_page_token = None

clients/python/src/mr_openapi/models/order_by_field.py

@@ -24,7 +24,7 @@ class OrderByField(str, Enum):
     """
     CREATE_TIME = "CREATE_TIME"
     LAST_UPDATE_TIME = "LAST_UPDATE_TIME"
-    ID = "Id"
+    ID = "ID"
 
     @classmethod
     def from_json(cls, json_str: str) -> Self:

clients/python/tests/test_client.py

@@ -203,6 +203,72 @@ def test_get_registered_models(client: ModelRegistry):
     assert i == models
 
 
+@pytest.mark.e2e
+def test_get_registered_models_order_by(client: ModelRegistry):
+    models = 5
+
+    rms = []
+    for name in [f"test_model{i}" for i in range(models)]:
+        rms.append(
+            client.register_model(
+                name,
+                "s3",
+                model_format_name="test_format",
+                model_format_version="test_version",
+                version="1.0.0",
+            )
+        )
+
+    # id ordering should match creation order
+    i = 0
+    for rm, by_id in zip(
+        rms,
+        client.get_registered_models().order_by_id(),
+    ):
+        assert rm.id == by_id.id
+        i += 1
+
+    assert i == models
+
+    # and obviously, creation ordering should match creation ordering
+    i = 0
+    for rm, by_creation in zip(
+        rms,
+        client.get_registered_models().order_by_creation_time(),
+    ):
+        assert rm.id == by_creation.id
+        i += 1
+
+    assert i == models
+
+    # update order should match creation ordering by default
+    i = 0
+    for rm, by_update in zip(
+        rms,
+        client.get_registered_models().order_by_update_time(),
+    ):
+        assert rm.id == by_update.id
+        i += 1
+
+    assert i == models
+
+    # now update the models in reverse order
+    for rm in reversed(rms):
+        rm.description = "updated"
+        client.update(rm)
+
+    # and they should match in reverse
+    i = 0
+    for rm, by_update in zip(
+        reversed(rms),
+        client.get_registered_models().order_by_update_time(),
+    ):
+        assert rm.id == by_update.id
+        i += 1
+
+    assert i == models
+
+
 @pytest.mark.e2e
 def test_get_registered_models_and_reset(client: ModelRegistry):
     model_count = 6
@@ -260,6 +326,64 @@ def test_get_model_versions(client: ModelRegistry):
     assert i == models
 
 
+@pytest.mark.e2e
+def test_get_model_versions_order_by(client: ModelRegistry):
+    name = "test_model"
+    models = 5
+    mvs = []
+    for v in [f"1.0.{i}" for i in range(models)]:
+        client.register_model(
+            name,
+            "s3",
+            model_format_name="test_format",
+            model_format_version="test_version",
+            version=v,
+        )
+        mvs.append(client.get_model_version(name, v))
+
+    i = 0
+    for mv, by_id in zip(
+        mvs,
+        client.get_model_versions(name).order_by_id(),
+    ):
+        assert mv.id == by_id.id
+        i += 1
+
+    assert i == models
+
+    i = 0
+    for mv, by_creation in zip(
+        mvs,
+        client.get_model_versions(name).order_by_creation_time(),
+    ):
+        assert mv.id == by_creation.id
+        i += 1
+
+    assert i == models
+
+    i = 0
+    for mv, by_update in zip(
+        mvs,
+        client.get_model_versions(name).order_by_update_time(),
+    ):
+        assert mv.id == by_update.id
+        i += 1
+
+    assert i == models
+
+    for mv in reversed(mvs):
+        mv.description = "updated"
+        client.update(mv)
+
+    i = 0
+    for mv, by_update in zip(
+        reversed(mvs),
+        client.get_model_versions(name).order_by_update_time(),
+    ):
+        assert mv.id == by_update.id
+        i += 1
+
+
 @pytest.mark.e2e
 def test_get_model_versions_and_reset(client: ModelRegistry):
     name = "test_model"