Start end (#97) (#100)

* Start end (#97)

* Update sarc/ldap/revision.py

Co-authored-by: Xavier Bouthillier <[email protected]>

* Update test_functional_revision.py

* Update assert msg

* Update sarc/ldap/revision.py

Co-authored-by: Xavier Bouthillier <[email protected]>

* Add dummy user

---------

Co-authored-by: Xavier Bouthillier <[email protected]>
Co-authored-by: Pierre Delaunay <[email protected]>

sarc/allocations/__init__.py

@@ -1 +1,3 @@
 from .allocations import get_allocation_summaries, get_allocations
+
+__all__ = ["get_allocation_summaries", "get_allocations"]

sarc/jobs/__init__.py

@@ -1,2 +1,11 @@
 from .job import SlurmJob, count_jobs, get_job, get_jobs
 from .series import get_job_time_series, get_job_time_series_metric_names
+
+__all__ = [
+    "SlurmJob",
+    "count_jobs",
+    "get_job",
+    "get_jobs",
+    "get_job_time_series",
+    "get_job_time_series_metric_names",
+]

sarc/ldap/acquire.py

@@ -13,18 +13,19 @@
 import json
 
 import pandas as pd
-from pymongo import UpdateOne
 
 import sarc.account_matching.make_matches
 import sarc.ldap.mymila
 from sarc.config import config
 from sarc.ldap.read_mila_ldap import fetch_ldap
+from sarc.ldap.revision import commit_matches_to_database
 
 
 def run(prompt=False):
     """If prompt is True, script will prompt for manual matching."""
 
     cfg = config()
+    user_collection = cfg.mongo.database_instance[cfg.ldap.mongo_collection_name]
 
     LD_users = fetch_ldap(ldap=cfg.ldap)
 
@@ -97,9 +98,12 @@ def run(prompt=False):
     #       "drac_members": {...} or None
     #     }
 
+    for _, user in DD_persons_matched.items():
+        fill_computed_fields(user)
+
     # These associations can now be propagated to the database.
-    save(
-        cfg.mongo.database_instance[cfg.ldap.mongo_collection_name],
+    commit_matches_to_database(
+        user_collection,
         DD_persons_matched,
     )
 
@@ -148,45 +152,5 @@ def fill_computed_fields(data: dict):
     return data
 
 
-def save(collection, users, verbose=False):
-    updates = []
-
-    for username, user in users.items():
-
-        user = fill_computed_fields(user)
-
-        updates.append(
-            UpdateOne(
-                {"mila_ldap.mila_email_username": username},
-                {
-                    # We set all the fields corresponding to the fields from `updated_user`,
-                    # so that's a convenient way to do it. Note that this does not affect
-                    # the fields in the database that are already present for that user.
-                    "$set": {
-                        "mila_ldap": user["mila_ldap"],
-                        "name": user["name"],
-                        "mila": user["mila"],
-                        "drac": user["drac"],
-                        "drac_roles": user["drac_roles"],
-                        "drac_members": user["drac_members"],
-                    },
-                },
-                upsert=True,
-            )
-        )
-
-    result = 0
-    if updates:
-        result = collection.bulk_write(updates)  #  <- the actual commit
-        if verbose:
-            print(result.bulk_api_result)
-    else:
-        if verbose:
-            print("Nothing to do.")
-
-    # might as well return this result in case we'd like to write tests for it
-    return result
-
-
 if __name__ == "__main__":
     run()

sarc/ldap/api.py

@@ -6,12 +6,15 @@
 
 from __future__ import annotations
 
+from datetime import date
 from typing import Optional
 
 from pydantic_mongo import AbstractRepository, ObjectIdField
 
 from sarc.config import BaseModel, config
 
+from .revision import query_latest_records
+
 
 class Credentials(BaseModel):
     username: str
@@ -31,20 +34,18 @@ class User(BaseModel):
     drac_members: Optional[dict]
     drac_roles: Optional[dict]
 
+    record_start: Optional[date] = None
+    record_end: Optional[date] = None
+
 
 class UserRepository(AbstractRepository[User]):
     class Meta:
         collection_name = "users"
 
-    def save_user(self, model: User):
-        document = self.to_document(model)
-        return self.get_collection().update_one(
-            {
-                "_id": model.id,
-            },
-            {"$set": document},
-            upsert=True,
-        )
+    # The API created by pydantic is too simplistic
+    # inserting into the users collection need to
+    # take into account revisions
+    # use: revision.update_user
 
 
 def users_collection():
@@ -53,36 +54,62 @@ def users_collection():
     return UserRepository(database=db)
 
 
-def get_users(query=None, query_options: dict | None = None):
+def get_users(query=None, query_options: dict | None = None, latest=True) -> list[User]:
     if query_options is None:
         query_options = {}
 
     if query is None:
-        query = {}
-    return list(users_collection().find_by(query, query_options))
+        return {}
+
+    if latest:
+        query = {
+            "$and": [
+                query_latest_records(),
+                query,
+            ]
+        }
+
+    results = users_collection().find_by(query, query_options)
+
+    return list(results)
 
 
 def get_user(
     mila_email_username=None, mila_cluster_username=None, drac_account_username=None
-):
+) -> Optional[User]:
     if mila_email_username is not None:
-        query = {"mila_ldap.mila_email_username": mila_email_username}
+        query = {
+            "$and": [
+                query_latest_records(),
+                {"mila_ldap.mila_email_username": mila_email_username},
+            ]
+        }
     elif mila_cluster_username is not None:
-        query = {"mila_ldap.mila_cluster_username": mila_cluster_username}
+        query = {
+            "$and": [
+                query_latest_records(),
+                {"mila_ldap.mila_cluster_username": mila_cluster_username},
+            ]
+        }
     elif drac_account_username is not None:
         query = {
-            "$or": [
-                {"drac_roles.username": drac_account_username},
-                {"drac_members.username": drac_account_username},
+            "$and": [
+                query_latest_records(),
+                {
+                    "$or": [
+                        {"drac_roles.username": drac_account_username},
+                        {"drac_members.username": drac_account_username},
+                    ]
+                },
             ]
         }
     else:
         raise ValueError("At least one of the arguments must be provided.")
 
-    L = get_users(query)
+    users = get_users(query)
 
-    assert len(L) <= 1
-    if len(L) == 1:
-        return L[0]
+    assert len(users) <= 1
+    if len(users) == 1:
+        return users[0]
     else:
         return None

sarc/ldap/read_mila_ldap.py

@@ -249,7 +249,7 @@ def client_side_user_updates(LD_users_DB, LD_users_LDAP):
     for meu in set(list(DD_users_DB.keys()) + list(DD_users_LDAP.keys())):
         # `meu` is short for the mila_email_username value
 
-        if meu in DD_users_DB and not meu in DD_users_LDAP:
+        if meu in DD_users_DB and meu not in DD_users_LDAP:
             # User is in DB but not in the LDAP.
             # Let's mark it as archived.
             entry = DD_users_DB[meu]