Add model_rebuild for job models

NERSC · Apr 10, 2024 · 376fe1f · 376fe1f
1 parent 284acda
commit 376fe1f
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 0 deletions.
diff --git a/src/sfapi_client/_async/client.py b/src/sfapi_client/_async/client.py
@@ -512,3 +512,9 @@ def resources(self) -> AsyncResources:
 AsyncRemotePath.model_rebuild()
 AsyncRole.model_rebuild()
 AsyncGroupMember.model_rebuild()
+
+# Ensure that the job models are built, we need to import here to
+# avoid circular imports
+from .jobs import AsyncJobSacct, AsyncJobSqueue
+AsyncJobSqueue.model_rebuild()
+AsyncJobSacct.model_rebuild()
diff --git a/src/sfapi_client/_sync/client.py b/src/sfapi_client/_sync/client.py
@@ -512,3 +512,9 @@ def resources(self) -> Resources:
 RemotePath.model_rebuild()
 Role.model_rebuild()
 GroupMember.model_rebuild()
+
+# Ensure that the job models are built, we need to import here to
+# avoid circular imports
+from .jobs import JobSacct, JobSqueue
+JobSqueue.model_rebuild()
+JobSacct.model_rebuild()