simstack 2.0.1__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (354) hide show
  1. simstack/__init__.py +0 -0
  2. simstack/_version.py +24 -0
  3. simstack/core/__init__.py +0 -0
  4. simstack/core/artifacts.py +406 -0
  5. simstack/core/asnyc_helper.py +14 -0
  6. simstack/core/config_file.py +60 -0
  7. simstack/core/context.py +283 -0
  8. simstack/core/definitions.py +22 -0
  9. simstack/core/engine.py +115 -0
  10. simstack/core/find_simstack_modules.py +41 -0
  11. simstack/core/hash.py +179 -0
  12. simstack/core/node.py +1137 -0
  13. simstack/core/node_claim.py +31 -0
  14. simstack/core/node_runner.py +362 -0
  15. simstack/core/process_results.py +102 -0
  16. simstack/core/recompute_artifacts.py +117 -0
  17. simstack/core/reset_db.py +12 -0
  18. simstack/core/resource_assignment.py +222 -0
  19. simstack/core/resources.py +90 -0
  20. simstack/core/route_table.py +24 -0
  21. simstack/core/run_docker.py +109 -0
  22. simstack/core/run_node.py +71 -0
  23. simstack/core/runner.py +169 -0
  24. simstack/core/services/__init__.py +11 -0
  25. simstack/core/services/base_service.py +191 -0
  26. simstack/core/services/file_transfer_service.py +117 -0
  27. simstack/core/services/git_restart_service.py +12 -0
  28. simstack/core/services/git_uv_update_service.py +105 -0
  29. simstack/core/services/node_execution_service.py +182 -0
  30. simstack/core/services/resource_branch_monitor_service.py +71 -0
  31. simstack/core/services/runner_cleanup_service.py +149 -0
  32. simstack/core/services/runner_manager.py +163 -0
  33. simstack/core/services/runner_status_service.py +10 -0
  34. simstack/core/services/slurm_status_service.py +67 -0
  35. simstack/core/services/stop_check_service.py +27 -0
  36. simstack/core/services/timeout_restart_service.py +26 -0
  37. simstack/core/simstack_result.py +39 -0
  38. simstack/core/submit_node.py +217 -0
  39. simstack/core/task_id.py +22 -0
  40. simstack/methods/__init__.py +0 -0
  41. simstack/methods/archive_files.py +292 -0
  42. simstack/methods/fire_and_forget_runner.py +113 -0
  43. simstack/methods/generate_test.py +125 -0
  44. simstack/methods/get_file.py +26 -0
  45. simstack/methods/mass_runner.py +154 -0
  46. simstack/methods/multiple_executor.py +307 -0
  47. simstack/methods/switch_git.py +16 -0
  48. simstack/methods/upload_helpers.py +81 -0
  49. simstack/models/__init__.py +74 -0
  50. simstack/models/array_list.py +21 -0
  51. simstack/models/array_storage.py +83 -0
  52. simstack/models/artifact_models.py +127 -0
  53. simstack/models/base_lists.py +51 -0
  54. simstack/models/base_types.py +209 -0
  55. simstack/models/charts_artifact.py +690 -0
  56. simstack/models/dataset.py +451 -0
  57. simstack/models/dataset_metadata.py +390 -0
  58. simstack/models/datasettuple.py +587 -0
  59. simstack/models/file_instance.py +150 -0
  60. simstack/models/file_list.py +45 -0
  61. simstack/models/files.py +424 -0
  62. simstack/models/fire_and_forget_result.py +19 -0
  63. simstack/models/images2d.py +113 -0
  64. simstack/models/log_entry_model.py +37 -0
  65. simstack/models/models.py +52 -0
  66. simstack/models/named_data_reference.py +36 -0
  67. simstack/models/node_registry.py +96 -0
  68. simstack/models/pandas_model.py +239 -0
  69. simstack/models/parameters.py +402 -0
  70. simstack/models/pickle_models.py +300 -0
  71. simstack/models/project.py +13 -0
  72. simstack/models/resource_assignment.py +170 -0
  73. simstack/models/resource_definition.py +125 -0
  74. simstack/models/runner_model.py +39 -0
  75. simstack/models/simple_table.py +79 -0
  76. simstack/models/simstack_model.py +97 -0
  77. simstack/models/slurm_info.py +17 -0
  78. simstack/models/table_artifact.py +87 -0
  79. simstack/models/tag.py +12 -0
  80. simstack/models/unused_test_file.py +5 -0
  81. simstack/simstack_marker.com +1 -0
  82. simstack/static/docs/.buildinfo +4 -0
  83. simstack/static/docs/_images/GUI Concept.webp +0 -0
  84. simstack/static/docs/_modules/index.html +434 -0
  85. simstack/static/docs/_modules/simstack/core/artifacts.html +743 -0
  86. simstack/static/docs/_modules/simstack/core/asnyc_helper.html +330 -0
  87. simstack/static/docs/_modules/simstack/core/config_file.html +379 -0
  88. simstack/static/docs/_modules/simstack/core/definitions.html +341 -0
  89. simstack/static/docs/_modules/simstack/core/engine.html +437 -0
  90. simstack/static/docs/_modules/simstack/core/find_simstack_modules.html +360 -0
  91. simstack/static/docs/_modules/simstack/core/hash.html +528 -0
  92. simstack/static/docs/_modules/simstack/core/node.html +1392 -0
  93. simstack/static/docs/_modules/simstack/core/node_claim.html +347 -0
  94. simstack/static/docs/_modules/simstack/core/node_runner.html +711 -0
  95. simstack/static/docs/_modules/simstack/core/process_results.html +418 -0
  96. simstack/static/docs/_modules/simstack/core/recompute_artifacts.html +433 -0
  97. simstack/static/docs/_modules/simstack/core/reset_db.html +328 -0
  98. simstack/static/docs/_modules/simstack/core/resource_assignment.html +553 -0
  99. simstack/static/docs/_modules/simstack/core/resources.html +424 -0
  100. simstack/static/docs/_modules/simstack/core/route_table.html +349 -0
  101. simstack/static/docs/_modules/simstack/core/run_docker.html +425 -0
  102. simstack/static/docs/_modules/simstack/core/run_node.html +390 -0
  103. simstack/static/docs/_modules/simstack/core/runner.html +491 -0
  104. simstack/static/docs/_modules/simstack/core/services/base_service.html +522 -0
  105. simstack/static/docs/_modules/simstack/core/services/file_transfer_service.html +436 -0
  106. simstack/static/docs/_modules/simstack/core/services/git_restart_service.html +331 -0
  107. simstack/static/docs/_modules/simstack/core/services/git_uv_update_service.html +427 -0
  108. simstack/static/docs/_modules/simstack/core/services/node_execution_service.html +507 -0
  109. simstack/static/docs/_modules/simstack/core/services/resource_branch_monitor_service.html +390 -0
  110. simstack/static/docs/_modules/simstack/core/services/runner_cleanup_service.html +468 -0
  111. simstack/static/docs/_modules/simstack/core/services/runner_manager.html +485 -0
  112. simstack/static/docs/_modules/simstack/core/services/runner_status_service.html +329 -0
  113. simstack/static/docs/_modules/simstack/core/services/slurm_status_service.html +386 -0
  114. simstack/static/docs/_modules/simstack/core/services/stop_check_service.html +346 -0
  115. simstack/static/docs/_modules/simstack/core/services/timeout_restart_service.html +345 -0
  116. simstack/static/docs/_modules/simstack/core/simstack_result.html +355 -0
  117. simstack/static/docs/_modules/simstack/core/submit_node.html +536 -0
  118. simstack/static/docs/_modules/simstack/core/task_id.html +344 -0
  119. simstack/static/docs/_modules/simstack/methods/archive_files.html +598 -0
  120. simstack/static/docs/_modules/simstack/methods/fire_and_forget_runner.html +432 -0
  121. simstack/static/docs/_modules/simstack/methods/generate_test.html +443 -0
  122. simstack/static/docs/_modules/simstack/methods/get_file.html +342 -0
  123. simstack/static/docs/_modules/simstack/methods/mass_runner.html +476 -0
  124. simstack/static/docs/_modules/simstack/methods/multiple_executor.html +674 -0
  125. simstack/static/docs/_modules/simstack/methods/switch_git.html +332 -0
  126. simstack/static/docs/_modules/simstack/methods/upload_helpers.html +403 -0
  127. simstack/static/docs/_modules/simstack/models/array_list.html +337 -0
  128. simstack/static/docs/_modules/simstack/models/array_storage.html +417 -0
  129. simstack/static/docs/_modules/simstack/models/base_lists.html +374 -0
  130. simstack/static/docs/_modules/simstack/models/base_types.html +588 -0
  131. simstack/static/docs/_modules/simstack/models/charts_artifact.html +1035 -0
  132. simstack/static/docs/_modules/simstack/models/dataset.html +866 -0
  133. simstack/static/docs/_modules/simstack/models/dataset_metadata.html +711 -0
  134. simstack/static/docs/_modules/simstack/models/datasettuple.html +1005 -0
  135. simstack/static/docs/_modules/simstack/models/file_instance.html +472 -0
  136. simstack/static/docs/_modules/simstack/models/file_list.html +367 -0
  137. simstack/static/docs/_modules/simstack/models/files.html +773 -0
  138. simstack/static/docs/_modules/simstack/models/fire_and_forget_result.html +335 -0
  139. simstack/static/docs/_modules/simstack/models/images2d.html +447 -0
  140. simstack/static/docs/_modules/simstack/models/log_entry_model.html +356 -0
  141. simstack/static/docs/_modules/simstack/models/models.html +374 -0
  142. simstack/static/docs/_modules/simstack/models/named_data_reference.html +355 -0
  143. simstack/static/docs/_modules/simstack/models/node_registry.html +415 -0
  144. simstack/static/docs/_modules/simstack/models/pandas_model.html +570 -0
  145. simstack/static/docs/_modules/simstack/models/parameters.html +739 -0
  146. simstack/static/docs/_modules/simstack/models/pickle_models.html +631 -0
  147. simstack/static/docs/_modules/simstack/models/project.html +329 -0
  148. simstack/static/docs/_modules/simstack/models/resource_assignment.html +501 -0
  149. simstack/static/docs/_modules/simstack/models/resource_definition.html +477 -0
  150. simstack/static/docs/_modules/simstack/models/runner_model.html +361 -0
  151. simstack/static/docs/_modules/simstack/models/simple_table.html +413 -0
  152. simstack/static/docs/_modules/simstack/models/simstack_model.html +416 -0
  153. simstack/static/docs/_modules/simstack/models/slurm_info.html +333 -0
  154. simstack/static/docs/_modules/simstack/models/table_artifact.html +406 -0
  155. simstack/static/docs/_modules/simstack/models/tag.html +328 -0
  156. simstack/static/docs/_modules/simstack/tables/model_table.html +519 -0
  157. simstack/static/docs/_modules/simstack/tables/node_children.html +538 -0
  158. simstack/static/docs/_modules/simstack/tables/node_table.html +854 -0
  159. simstack/static/docs/_modules/simstack/tables/table_builder.html +642 -0
  160. simstack/static/docs/_modules/simstack/util/async_zip_utils.html +337 -0
  161. simstack/static/docs/_modules/simstack/util/b64mixin.html +339 -0
  162. simstack/static/docs/_modules/simstack/util/cleaned_json_schema.html +336 -0
  163. simstack/static/docs/_modules/simstack/util/config_reader.html +498 -0
  164. simstack/static/docs/_modules/simstack/util/custom_model_dump.html +391 -0
  165. simstack/static/docs/_modules/simstack/util/database_information.html +488 -0
  166. simstack/static/docs/_modules/simstack/util/db.html +1392 -0
  167. simstack/static/docs/_modules/simstack/util/db_logger.html +448 -0
  168. simstack/static/docs/_modules/simstack/util/default_from_dict.html +376 -0
  169. simstack/static/docs/_modules/simstack/util/directory_iterator.html +553 -0
  170. simstack/static/docs/_modules/simstack/util/docstring_parser.html +526 -0
  171. simstack/static/docs/_modules/simstack/util/file_hashing.html +394 -0
  172. simstack/static/docs/_modules/simstack/util/file_transfer_client.html +791 -0
  173. simstack/static/docs/_modules/simstack/util/files.html +733 -0
  174. simstack/static/docs/_modules/simstack/util/find_python_files.html +349 -0
  175. simstack/static/docs/_modules/simstack/util/generate_ui_schema.html +396 -0
  176. simstack/static/docs/_modules/simstack/util/generic_list_mixin.html +487 -0
  177. simstack/static/docs/_modules/simstack/util/get_module_path.html +336 -0
  178. simstack/static/docs/_modules/simstack/util/git_pull.html +381 -0
  179. simstack/static/docs/_modules/simstack/util/git_repository_status.html +411 -0
  180. simstack/static/docs/_modules/simstack/util/import_module.html +377 -0
  181. simstack/static/docs/_modules/simstack/util/importer.html +637 -0
  182. simstack/static/docs/_modules/simstack/util/init_data_source.html +357 -0
  183. simstack/static/docs/_modules/simstack/util/make_graph.html +423 -0
  184. simstack/static/docs/_modules/simstack/util/make_table.html +1171 -0
  185. simstack/static/docs/_modules/simstack/util/mappings.html +378 -0
  186. simstack/static/docs/_modules/simstack/util/minimal_route_finder.html +388 -0
  187. simstack/static/docs/_modules/simstack/util/mock_db.html +472 -0
  188. simstack/static/docs/_modules/simstack/util/module_path_checker.html +354 -0
  189. simstack/static/docs/_modules/simstack/util/object_list_mixin.html +649 -0
  190. simstack/static/docs/_modules/simstack/util/path_manager.html +566 -0
  191. simstack/static/docs/_modules/simstack/util/project_root_finder.html +357 -0
  192. simstack/static/docs/_modules/simstack/util/queue_watcher.html +593 -0
  193. simstack/static/docs/_modules/simstack/util/resource_config.html +532 -0
  194. simstack/static/docs/_modules/simstack/util/route_finder.html +390 -0
  195. simstack/static/docs/_modules/simstack/util/runner_util.html +487 -0
  196. simstack/static/docs/_modules/simstack/util/runner_utils.html +468 -0
  197. simstack/static/docs/_modules/simstack/util/safe_code_executor.html +510 -0
  198. simstack/static/docs/_modules/simstack/util/sample_config.html +414 -0
  199. simstack/static/docs/_modules/simstack/util/sensitive_content_filter.html +364 -0
  200. simstack/static/docs/_modules/simstack/util/setup_logging.html +365 -0
  201. simstack/static/docs/_modules/simstack/util/submit_to_watchdog.html +425 -0
  202. simstack/static/docs/_modules/simstack/util/toml_reader.html +482 -0
  203. simstack/static/docs/_modules/simstack/util/transform_file_name.html +357 -0
  204. simstack/static/docs/_modules/simstack/util/ui_tools.html +541 -0
  205. simstack/static/docs/_sources/architecture.rst.txt +25 -0
  206. simstack/static/docs/_sources/array_lists.rst.txt +60 -0
  207. simstack/static/docs/_sources/ci_pipeline.rst.txt +34 -0
  208. simstack/static/docs/_sources/data.rst.txt +175 -0
  209. simstack/static/docs/_sources/database_migrations.rst.txt +310 -0
  210. simstack/static/docs/_sources/datasets.rst.txt +93 -0
  211. simstack/static/docs/_sources/developer_guide.rst.txt +25 -0
  212. simstack/static/docs/_sources/docker_nodes.rst.txt +149 -0
  213. simstack/static/docs/_sources/files.rst.txt +159 -0
  214. simstack/static/docs/_sources/fire_and_forget_runner.rst.txt +67 -0
  215. simstack/static/docs/_sources/index.rst.txt +25 -0
  216. simstack/static/docs/_sources/installation.rst.txt +213 -0
  217. simstack/static/docs/_sources/lists.rst.txt +111 -0
  218. simstack/static/docs/_sources/mass_runner.rst.txt +91 -0
  219. simstack/static/docs/_sources/misc.rst.txt +106 -0
  220. simstack/static/docs/_sources/overview.rst.txt +91 -0
  221. simstack/static/docs/_sources/recovery.rst.txt +38 -0
  222. simstack/static/docs/_sources/resource_config.rst.txt +141 -0
  223. simstack/static/docs/_sources/simstack/modules/modules.rst.txt +7 -0
  224. simstack/static/docs/_sources/simstack/modules/simstack.core.rst.txt +205 -0
  225. simstack/static/docs/_sources/simstack/modules/simstack.core.services.rst.txt +109 -0
  226. simstack/static/docs/_sources/simstack/modules/simstack.methods.rst.txt +77 -0
  227. simstack/static/docs/_sources/simstack/modules/simstack.models.rst.txt +261 -0
  228. simstack/static/docs/_sources/simstack/modules/simstack.rst.txt +22 -0
  229. simstack/static/docs/_sources/simstack/modules/simstack.tables.rst.txt +45 -0
  230. simstack/static/docs/_sources/simstack/modules/simstack.util.rst.txt +381 -0
  231. simstack/static/docs/_sources/submit.rst.txt +141 -0
  232. simstack/static/docs/_sources/tables.rst.txt +50 -0
  233. simstack/static/docs/_sources/testing_complex_nodes.rst.txt +92 -0
  234. simstack/static/docs/_sources/ui_documentation.rst.txt +272 -0
  235. simstack/static/docs/_sources/workflows.rst.txt +446 -0
  236. simstack/static/docs/_static/base-stemmer.js +476 -0
  237. simstack/static/docs/_static/basic.css +906 -0
  238. simstack/static/docs/_static/cobalt2-theme.css +1 -0
  239. simstack/static/docs/_static/debug.css +69 -0
  240. simstack/static/docs/_static/docs-build-badge.css +51 -0
  241. simstack/static/docs/_static/docs-build-badge.js +51 -0
  242. simstack/static/docs/_static/doctools.js +150 -0
  243. simstack/static/docs/_static/documentation_options.js +13 -0
  244. simstack/static/docs/_static/english-stemmer.js +1066 -0
  245. simstack/static/docs/_static/file.png +0 -0
  246. simstack/static/docs/_static/language_data.js +13 -0
  247. simstack/static/docs/_static/minus.png +0 -0
  248. simstack/static/docs/_static/plus.png +0 -0
  249. simstack/static/docs/_static/pygments.css +250 -0
  250. simstack/static/docs/_static/scripts/furo-extensions.js +0 -0
  251. simstack/static/docs/_static/scripts/furo.js +3 -0
  252. simstack/static/docs/_static/scripts/furo.js.LICENSE.txt +7 -0
  253. simstack/static/docs/_static/scripts/furo.js.map +1 -0
  254. simstack/static/docs/_static/searchtools.js +693 -0
  255. simstack/static/docs/_static/skeleton.css +296 -0
  256. simstack/static/docs/_static/sphinx_highlight.js +159 -0
  257. simstack/static/docs/_static/styles/furo-extensions.css +2 -0
  258. simstack/static/docs/_static/styles/furo-extensions.css.map +1 -0
  259. simstack/static/docs/_static/styles/furo.css +2 -0
  260. simstack/static/docs/_static/styles/furo.css.map +1 -0
  261. simstack/static/docs/architecture.html +340 -0
  262. simstack/static/docs/array_lists.html +678 -0
  263. simstack/static/docs/ci_pipeline.html +385 -0
  264. simstack/static/docs/data.html +545 -0
  265. simstack/static/docs/database_migrations.html +652 -0
  266. simstack/static/docs/datasets.html +1268 -0
  267. simstack/static/docs/developer_guide.html +359 -0
  268. simstack/static/docs/docker_nodes.html +503 -0
  269. simstack/static/docs/files.html +718 -0
  270. simstack/static/docs/fire_and_forget_runner.html +641 -0
  271. simstack/static/docs/genindex.html +6390 -0
  272. simstack/static/docs/index.html +516 -0
  273. simstack/static/docs/installation.html +553 -0
  274. simstack/static/docs/lists.html +478 -0
  275. simstack/static/docs/mass_runner.html +481 -0
  276. simstack/static/docs/misc.html +454 -0
  277. simstack/static/docs/objects.inv +0 -0
  278. simstack/static/docs/overview.html +469 -0
  279. simstack/static/docs/py-modindex.html +1111 -0
  280. simstack/static/docs/recovery.html +371 -0
  281. simstack/static/docs/resource_config.html +590 -0
  282. simstack/static/docs/search.html +327 -0
  283. simstack/static/docs/searchindex.js +1 -0
  284. simstack/static/docs/simstack/modules/modules.html +489 -0
  285. simstack/static/docs/simstack/modules/simstack.core.html +1883 -0
  286. simstack/static/docs/simstack/modules/simstack.core.services.html +698 -0
  287. simstack/static/docs/simstack/modules/simstack.html +3060 -0
  288. simstack/static/docs/simstack/modules/simstack.methods.html +1150 -0
  289. simstack/static/docs/simstack/modules/simstack.models.html +13113 -0
  290. simstack/static/docs/simstack/modules/simstack.tables.html +635 -0
  291. simstack/static/docs/simstack/modules/simstack.util.html +3320 -0
  292. simstack/static/docs/submit.html +487 -0
  293. simstack/static/docs/tables.html +539 -0
  294. simstack/static/docs/testing_complex_nodes.html +458 -0
  295. simstack/static/docs/ui_documentation.html +788 -0
  296. simstack/static/docs/workflows.html +766 -0
  297. simstack/tables/__init__.py +0 -0
  298. simstack/tables/model_table.py +191 -0
  299. simstack/tables/node_children.py +219 -0
  300. simstack/tables/node_table.py +517 -0
  301. simstack/tables/table_builder.py +314 -0
  302. simstack/util/__init__.py +0 -0
  303. simstack/util/async_zip_utils.py +18 -0
  304. simstack/util/b64mixin.py +23 -0
  305. simstack/util/cleaned_json_schema.py +20 -0
  306. simstack/util/config_reader.py +179 -0
  307. simstack/util/custom_model_dump.py +75 -0
  308. simstack/util/database_information.py +160 -0
  309. simstack/util/db.py +1013 -0
  310. simstack/util/db_logger.py +117 -0
  311. simstack/util/default_from_dict.py +57 -0
  312. simstack/util/directory_iterator.py +216 -0
  313. simstack/util/docstring_parser.py +186 -0
  314. simstack/util/file_hashing.py +72 -0
  315. simstack/util/file_transfer_client.py +427 -0
  316. simstack/util/files.py +396 -0
  317. simstack/util/find_python_files.py +33 -0
  318. simstack/util/generate_ui_schema.py +80 -0
  319. simstack/util/generic_list_mixin.py +120 -0
  320. simstack/util/geometry_utils.py +118 -0
  321. simstack/util/get_module_path.py +20 -0
  322. simstack/util/git_pull.py +65 -0
  323. simstack/util/git_repository_status.py +95 -0
  324. simstack/util/import_module.py +63 -0
  325. simstack/util/importer.py +312 -0
  326. simstack/util/init_data_source.py +41 -0
  327. simstack/util/make_graph.py +107 -0
  328. simstack/util/make_table.py +837 -0
  329. simstack/util/mappings.py +41 -0
  330. simstack/util/minimal_route_finder.py +72 -0
  331. simstack/util/mock_db.py +144 -0
  332. simstack/util/module_path_checker.py +38 -0
  333. simstack/util/object_list_mixin.py +282 -0
  334. simstack/util/path_manager.py +223 -0
  335. simstack/util/project_root_finder.py +41 -0
  336. simstack/util/queue_watcher.py +262 -0
  337. simstack/util/resource_config.py +198 -0
  338. simstack/util/route_finder.py +74 -0
  339. simstack/util/runner_util.py +166 -0
  340. simstack/util/runner_utils.py +143 -0
  341. simstack/util/safe_code_executor.py +194 -0
  342. simstack/util/sample_config.py +98 -0
  343. simstack/util/sensitive_content_filter.py +48 -0
  344. simstack/util/setup_logging.py +49 -0
  345. simstack/util/submit_to_watchdog.py +106 -0
  346. simstack/util/toml_reader.py +142 -0
  347. simstack/util/transform_file_name.py +41 -0
  348. simstack/util/ui_tools.py +213 -0
  349. simstack-2.0.1.dist-info/METADATA +454 -0
  350. simstack-2.0.1.dist-info/RECORD +354 -0
  351. simstack-2.0.1.dist-info/WHEEL +5 -0
  352. simstack-2.0.1.dist-info/entry_points.txt +12 -0
  353. simstack-2.0.1.dist-info/licenses/LICENSE +21 -0
  354. simstack-2.0.1.dist-info/top_level.txt +1 -0
simstack/static/docs/_modules/simstack/util/db.html ADDED
@@ -0,0 +1,1392 @@
1
+ <!doctype html>
2
+ <html class="no-js" lang="en" data-content_root="../../../">
3
+ <head><meta charset="utf-8">
4
+ <meta name="viewport" content="width=device-width,initial-scale=1">
5
+ <meta name="color-scheme" content="light dark"><link rel="index" title="Index" href="../../../genindex.html"><link rel="search" title="Search" href="../../../search.html">
6
+
7
+ <!-- Generated with Sphinx 9.1.0 and Furo 2025.12.19 -->
8
+ <title>simstack.util.db - SimStack II documentation</title>
9
+ <link rel="stylesheet" type="text/css" href="../../../_static/pygments.css?v=d111a655" />
10
+ <link rel="stylesheet" type="text/css" href="../../../_static/styles/furo.css?v=7bdb33bb" />
11
+ <link rel="stylesheet" type="text/css" href="../../../_static/styles/furo-extensions.css?v=8dab3a3b" />
12
+ <link rel="stylesheet" type="text/css" href="../../../_static/cobalt2-theme.css?v=2d651046" />
13
+ <link rel="stylesheet" type="text/css" href="../../../_static/docs-build-badge.css?v=803faff8" />
14
+
15
+
16
+
17
+
18
+ <style>
19
+ body {
20
+ --color-code-background: #f2f2f2;
21
+ --color-code-foreground: #1e1e1e;
22
+
23
+ }
24
+ @media not print {
25
+ body[data-theme="dark"] {
26
+ --color-code-background: #202020;
27
+ --color-code-foreground: #d0d0d0;
28
+ --color-background-primary: #132738;
29
+ --color-background-secondary: #1e3a52;
30
+ --color-foreground-primary: #ffffff;
31
+ --color-foreground-secondary: #bac8d3;
32
+ --color-brand-primary: #22d5ff;
33
+ --color-brand-content: #0088ff;
34
+
35
+ }
36
+ @media (prefers-color-scheme: dark) {
37
+ body:not([data-theme="light"]) {
38
+ --color-code-background: #202020;
39
+ --color-code-foreground: #d0d0d0;
40
+ --color-background-primary: #132738;
41
+ --color-background-secondary: #1e3a52;
42
+ --color-foreground-primary: #ffffff;
43
+ --color-foreground-secondary: #bac8d3;
44
+ --color-brand-primary: #22d5ff;
45
+ --color-brand-content: #0088ff;
46
+
47
+ }
48
+ }
49
+ }
50
+ </style></head>
51
+ <body>
52
+
53
+ <script>
54
+ document.body.dataset.theme = localStorage.getItem("theme") || "auto";
55
+ </script>
56
+
57
+
58
+ <svg xmlns="http://www.w3.org/2000/svg" style="display: none;">
59
+ <symbol id="svg-toc" viewBox="0 0 24 24">
60
+ <title>Contents</title>
61
+ <svg stroke="currentColor" fill="currentColor" stroke-width="0" viewBox="0 0 1024 1024">
62
+ <path d="M408 442h480c4.4 0 8-3.6 8-8v-56c0-4.4-3.6-8-8-8H408c-4.4 0-8 3.6-8 8v56c0 4.4 3.6 8 8 8zm-8 204c0 4.4 3.6 8 8 8h480c4.4 0 8-3.6 8-8v-56c0-4.4-3.6-8-8-8H408c-4.4 0-8 3.6-8 8v56zm504-486H120c-4.4 0-8 3.6-8 8v56c0 4.4 3.6 8 8 8h784c4.4 0 8-3.6 8-8v-56c0-4.4-3.6-8-8-8zm0 632H120c-4.4 0-8 3.6-8 8v56c0 4.4 3.6 8 8 8h784c4.4 0 8-3.6 8-8v-56c0-4.4-3.6-8-8-8zM115.4 518.9L271.7 642c5.8 4.6 14.4.5 14.4-6.9V388.9c0-7.4-8.5-11.5-14.4-6.9L115.4 505.1a8.74 8.74 0 0 0 0 13.8z"/>
63
+ </svg>
64
+ </symbol>
65
+ <symbol id="svg-menu" viewBox="0 0 24 24">
66
+ <title>Menu</title>
67
+ <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="currentColor"
68
+ stroke-width="2" stroke-linecap="round" stroke-linejoin="round" class="feather-menu">
69
+ <line x1="3" y1="12" x2="21" y2="12"></line>
70
+ <line x1="3" y1="6" x2="21" y2="6"></line>
71
+ <line x1="3" y1="18" x2="21" y2="18"></line>
72
+ </svg>
73
+ </symbol>
74
+ <symbol id="svg-arrow-right" viewBox="0 0 24 24">
75
+ <title>Expand</title>
76
+ <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="currentColor"
77
+ stroke-width="2" stroke-linecap="round" stroke-linejoin="round" class="feather-chevron-right">
78
+ <polyline points="9 18 15 12 9 6"></polyline>
79
+ </svg>
80
+ </symbol>
81
+ <symbol id="svg-sun" viewBox="0 0 24 24">
82
+ <title>Light mode</title>
83
+ <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="currentColor"
84
+ stroke-width="1" stroke-linecap="round" stroke-linejoin="round" class="feather-sun">
85
+ <circle cx="12" cy="12" r="5"></circle>
86
+ <line x1="12" y1="1" x2="12" y2="3"></line>
87
+ <line x1="12" y1="21" x2="12" y2="23"></line>
88
+ <line x1="4.22" y1="4.22" x2="5.64" y2="5.64"></line>
89
+ <line x1="18.36" y1="18.36" x2="19.78" y2="19.78"></line>
90
+ <line x1="1" y1="12" x2="3" y2="12"></line>
91
+ <line x1="21" y1="12" x2="23" y2="12"></line>
92
+ <line x1="4.22" y1="19.78" x2="5.64" y2="18.36"></line>
93
+ <line x1="18.36" y1="5.64" x2="19.78" y2="4.22"></line>
94
+ </svg>
95
+ </symbol>
96
+ <symbol id="svg-moon" viewBox="0 0 24 24">
97
+ <title>Dark mode</title>
98
+ <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="currentColor"
99
+ stroke-width="1" stroke-linecap="round" stroke-linejoin="round" class="icon-tabler-moon">
100
+ <path stroke="none" d="M0 0h24v24H0z" fill="none" />
101
+ <path d="M12 3c.132 0 .263 0 .393 0a7.5 7.5 0 0 0 7.92 12.446a9 9 0 1 1 -8.313 -12.454z" />
102
+ </svg>
103
+ </symbol>
104
+ <symbol id="svg-sun-with-moon" viewBox="0 0 24 24">
105
+ <title>Auto light/dark, in light mode</title>
106
+ <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="currentColor"
107
+ stroke-width="1" stroke-linecap="round" stroke-linejoin="round"
108
+ class="icon-custom-derived-from-feather-sun-and-tabler-moon">
109
+ <path style="opacity: 50%" d="M 5.411 14.504 C 5.471 14.504 5.532 14.504 5.591 14.504 C 3.639 16.319 4.383 19.569 6.931 20.352 C 7.693 20.586 8.512 20.551 9.25 20.252 C 8.023 23.207 4.056 23.725 2.11 21.184 C 0.166 18.642 1.702 14.949 4.874 14.536 C 5.051 14.512 5.231 14.5 5.411 14.5 L 5.411 14.504 Z"/>
110
+ <line x1="14.5" y1="3.25" x2="14.5" y2="1.25"/>
111
+ <line x1="14.5" y1="15.85" x2="14.5" y2="17.85"/>
112
+ <line x1="10.044" y1="5.094" x2="8.63" y2="3.68"/>
113
+ <line x1="19" y1="14.05" x2="20.414" y2="15.464"/>
114
+ <line x1="8.2" y1="9.55" x2="6.2" y2="9.55"/>
115
+ <line x1="20.8" y1="9.55" x2="22.8" y2="9.55"/>
116
+ <line x1="10.044" y1="14.006" x2="8.63" y2="15.42"/>
117
+ <line x1="19" y1="5.05" x2="20.414" y2="3.636"/>
118
+ <circle cx="14.5" cy="9.55" r="3.6"/>
119
+ </svg>
120
+ </symbol>
121
+ <symbol id="svg-moon-with-sun" viewBox="0 0 24 24">
122
+ <title>Auto light/dark, in dark mode</title>
123
+ <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="currentColor"
124
+ stroke-width="1" stroke-linecap="round" stroke-linejoin="round"
125
+ class="icon-custom-derived-from-feather-sun-and-tabler-moon">
126
+ <path d="M 8.282 7.007 C 8.385 7.007 8.494 7.007 8.595 7.007 C 5.18 10.184 6.481 15.869 10.942 17.24 C 12.275 17.648 13.706 17.589 15 17.066 C 12.851 22.236 5.91 23.143 2.505 18.696 C -0.897 14.249 1.791 7.786 7.342 7.063 C 7.652 7.021 7.965 7 8.282 7 L 8.282 7.007 Z"/>
127
+ <line style="opacity: 50%" x1="18" y1="3.705" x2="18" y2="2.5"/>
128
+ <line style="opacity: 50%" x1="18" y1="11.295" x2="18" y2="12.5"/>
129
+ <line style="opacity: 50%" x1="15.316" y1="4.816" x2="14.464" y2="3.964"/>
130
+ <line style="opacity: 50%" x1="20.711" y1="10.212" x2="21.563" y2="11.063"/>
131
+ <line style="opacity: 50%" x1="14.205" y1="7.5" x2="13.001" y2="7.5"/>
132
+ <line style="opacity: 50%" x1="21.795" y1="7.5" x2="23" y2="7.5"/>
133
+ <line style="opacity: 50%" x1="15.316" y1="10.184" x2="14.464" y2="11.036"/>
134
+ <line style="opacity: 50%" x1="20.711" y1="4.789" x2="21.563" y2="3.937"/>
135
+ <circle style="opacity: 50%" cx="18" cy="7.5" r="2.169"/>
136
+ </svg>
137
+ </symbol>
138
+ <symbol id="svg-pencil" viewBox="0 0 24 24">
139
+ <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="currentColor"
140
+ stroke-width="1" stroke-linecap="round" stroke-linejoin="round" class="icon-tabler-pencil-code">
141
+ <path d="M4 20h4l10.5 -10.5a2.828 2.828 0 1 0 -4 -4l-10.5 10.5v4" />
142
+ <path d="M13.5 6.5l4 4" />
143
+ <path d="M20 21l2 -2l-2 -2" />
144
+ <path d="M17 17l-2 2l2 2" />
145
+ </svg>
146
+ </symbol>
147
+ <symbol id="svg-eye" viewBox="0 0 24 24">
148
+ <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="currentColor"
149
+ stroke-width="1" stroke-linecap="round" stroke-linejoin="round" class="icon-tabler-eye-code">
150
+ <path stroke="none" d="M0 0h24v24H0z" fill="none" />
151
+ <path d="M10 12a2 2 0 1 0 4 0a2 2 0 0 0 -4 0" />
152
+ <path
153
+ d="M11.11 17.958c-3.209 -.307 -5.91 -2.293 -8.11 -5.958c2.4 -4 5.4 -6 9 -6c3.6 0 6.6 2 9 6c-.21 .352 -.427 .688 -.647 1.008" />
154
+ <path d="M20 21l2 -2l-2 -2" />
155
+ <path d="M17 17l-2 2l2 2" />
156
+ </svg>
157
+ </symbol>
158
+ </svg>
159
+
160
+ <input type="checkbox" class="sidebar-toggle" name="__navigation" id="__navigation" aria-label="Toggle site navigation sidebar">
161
+ <input type="checkbox" class="sidebar-toggle" name="__toc" id="__toc" aria-label="Toggle table of contents sidebar">
162
+ <label class="overlay sidebar-overlay" for="__navigation"></label>
163
+ <label class="overlay toc-overlay" for="__toc"></label>
164
+
165
+ <a class="skip-to-content muted-link" href="#furo-main-content">Skip to content</a>
166
+
167
+
168
+
169
+ <div class="page">
170
+ <header class="mobile-header">
171
+ <div class="header-left">
172
+ <label class="nav-overlay-icon" for="__navigation">
173
+ <span class="icon"><svg><use href="#svg-menu"></use></svg></span>
174
+ </label>
175
+ </div>
176
+ <div class="header-center">
177
+ <a href="../../../index.html"><div class="brand">SimStack II documentation</div></a>
178
+ </div>
179
+ <div class="header-right">
180
+ <div class="theme-toggle-container theme-toggle-header">
181
+ <button class="theme-toggle" aria-label="Toggle Light / Dark / Auto color theme">
182
+ <svg class="theme-icon-when-auto-light"><use href="#svg-sun-with-moon"></use></svg>
183
+ <svg class="theme-icon-when-auto-dark"><use href="#svg-moon-with-sun"></use></svg>
184
+ <svg class="theme-icon-when-dark"><use href="#svg-moon"></use></svg>
185
+ <svg class="theme-icon-when-light"><use href="#svg-sun"></use></svg>
186
+ </button>
187
+ </div>
188
+ <label class="toc-overlay-icon toc-header-icon no-toc" for="__toc">
189
+ <span class="icon"><svg><use href="#svg-toc"></use></svg></span>
190
+ </label>
191
+ </div>
192
+ </header>
193
+ <aside class="sidebar-drawer">
194
+ <div class="sidebar-container">
195
+
196
+ <div class="sidebar-sticky"><a class="sidebar-brand" href="../../../index.html">
197
+
198
+
199
+ </a><form class="sidebar-search-container" method="get" action="../../../search.html" role="search">
200
+ <input class="sidebar-search" placeholder="Search" name="q" aria-label="Search">
201
+ <input type="hidden" name="check_keywords" value="yes">
202
+ <input type="hidden" name="area" value="default">
203
+ </form>
204
+ <div id="searchbox"></div><div class="sidebar-scroll"><div class="sidebar-tree">
205
+ <p class="caption" role="heading"><span class="caption-text">Contents:</span></p>
206
+ <ul>
207
+ <li class="toctree-l1"><a class="reference internal" href="../../../overview.html">Overview</a></li>
208
+ <li class="toctree-l1"><a class="reference internal" href="../../../workflows.html">Workflows</a></li>
209
+ <li class="toctree-l1"><a class="reference internal" href="../../../ui_documentation.html">Graphical User Interface</a></li>
210
+ <li class="toctree-l1 has-children"><a class="reference internal" href="../../../developer_guide.html">Developer Guide</a><input aria-label="Toggle navigation of Developer Guide" class="toctree-checkbox" id="toctree-checkbox-1" name="toctree-checkbox-1" role="switch" type="checkbox"/><label for="toctree-checkbox-1"><span class="icon"><svg><use href="#svg-arrow-right"></use></svg></span></label><ul>
211
+ <li class="toctree-l2"><a class="reference internal" href="../../../data.html">Persisting and Visualizing Data</a></li>
212
+ <li class="toctree-l2"><a class="reference internal" href="../../../datasets.html">Datasets</a></li>
213
+ <li class="toctree-l2"><a class="reference internal" href="../../../lists.html">List Mixins</a></li>
214
+ <li class="toctree-l2"><a class="reference internal" href="../../../array_lists.html">Array Lists</a></li>
215
+ <li class="toctree-l2"><a class="reference internal" href="../../../tables.html">Tables</a></li>
216
+ <li class="toctree-l2"><a class="reference internal" href="../../../files.html">File Management</a></li>
217
+ <li class="toctree-l2"><a class="reference internal" href="../../../resource_config.html">Program Execution</a></li>
218
+ <li class="toctree-l2"><a class="reference internal" href="../../../mass_runner.html">MassRunner</a></li>
219
+ <li class="toctree-l2"><a class="reference internal" href="../../../fire_and_forget_runner.html">FireAndForgetRunner</a></li>
220
+ <li class="toctree-l2"><a class="reference internal" href="../../../docker_nodes.html">Creating Docker Containers for SimStack Nodes</a></li>
221
+ <li class="toctree-l2"><a class="reference internal" href="../../../testing_complex_nodes.html">Testing Complex Nodes</a></li>
222
+ <li class="toctree-l2"><a class="reference internal" href="../../../database_migrations.html">Database Migrations</a></li>
223
+ <li class="toctree-l2"><a class="reference internal" href="../../../ci_pipeline.html">CI Pipeline and Documentation Artifacts</a></li>
224
+ <li class="toctree-l2"><a class="reference internal" href="../../../misc.html">Miscellaneous</a></li>
225
+ </ul>
226
+ </li>
227
+ <li class="toctree-l1"><a class="reference internal" href="../../../installation.html">Simstack II Setup</a></li>
228
+ <li class="toctree-l1 has-children"><a class="reference internal" href="../../../simstack/modules/modules.html">simstack</a><input aria-label="Toggle navigation of simstack" class="toctree-checkbox" id="toctree-checkbox-2" name="toctree-checkbox-2" role="switch" type="checkbox"/><label for="toctree-checkbox-2"><span class="icon"><svg><use href="#svg-arrow-right"></use></svg></span></label><ul>
229
+ <li class="toctree-l2 has-children"><a class="reference internal" href="../../../simstack/modules/simstack.html">simstack package</a><input aria-label="Toggle navigation of simstack package" class="toctree-checkbox" id="toctree-checkbox-3" name="toctree-checkbox-3" role="switch" type="checkbox"/><label for="toctree-checkbox-3"><span class="icon"><svg><use href="#svg-arrow-right"></use></svg></span></label><ul>
230
+ <li class="toctree-l3 has-children"><a class="reference internal" href="../../../simstack/modules/simstack.core.html">simstack.core package</a><input aria-label="Toggle navigation of simstack.core package" class="toctree-checkbox" id="toctree-checkbox-4" name="toctree-checkbox-4" role="switch" type="checkbox"/><label for="toctree-checkbox-4"><span class="icon"><svg><use href="#svg-arrow-right"></use></svg></span></label><ul>
231
+ <li class="toctree-l4"><a class="reference internal" href="../../../simstack/modules/simstack.core.services.html">simstack.core.services package</a></li>
232
+ </ul>
233
+ </li>
234
+ <li class="toctree-l3"><a class="reference internal" href="../../../simstack/modules/simstack.methods.html">simstack.methods package</a></li>
235
+ <li class="toctree-l3"><a class="reference internal" href="../../../simstack/modules/simstack.models.html">simstack.models package</a></li>
236
+ <li class="toctree-l3"><a class="reference internal" href="../../../simstack/modules/simstack.tables.html">simstack.tables package</a></li>
237
+ <li class="toctree-l3"><a class="reference internal" href="../../../simstack/modules/simstack.util.html">simstack.util package</a></li>
238
+ </ul>
239
+ </li>
240
+ </ul>
241
+ </li>
242
+ </ul>
243
+
244
+ </div>
245
+ </div>
246
+
247
+ </div>
248
+
249
+ </div>
250
+ </aside>
251
+ <div class="main">
252
+ <div class="content">
253
+ <div class="article-container">
254
+ <a href="#" class="back-to-top muted-link">
255
+ <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24">
256
+ <path d="M13 20h-2V8l-5.5 5.5-1.42-1.42L12 4.16l7.92 7.92-1.42 1.42L13 8v12z"></path>
257
+ </svg>
258
+ <span>Back to top</span>
259
+ </a>
260
+ <div class="content-icon-container">
261
+ <div class="theme-toggle-container theme-toggle-content">
262
+ <button class="theme-toggle" aria-label="Toggle Light / Dark / Auto color theme">
263
+ <svg class="theme-icon-when-auto-light"><use href="#svg-sun-with-moon"></use></svg>
264
+ <svg class="theme-icon-when-auto-dark"><use href="#svg-moon-with-sun"></use></svg>
265
+ <svg class="theme-icon-when-dark"><use href="#svg-moon"></use></svg>
266
+ <svg class="theme-icon-when-light"><use href="#svg-sun"></use></svg>
267
+ </button>
268
+ </div>
269
+ <label class="toc-overlay-icon toc-content-icon no-toc" for="__toc">
270
+ <span class="icon"><svg><use href="#svg-toc"></use></svg></span>
271
+ </label>
272
+ </div>
273
+ <article role="main" id="furo-main-content">
274
+ <h1>Source code for simstack.util.db</h1><div class="highlight"><pre>
275
+ <span></span><span class="kn">from</span><span class="w"> </span><span class="nn">__future__</span><span class="w"> </span><span class="kn">import</span> <span class="n">annotations</span>
276
+
277
+ <span class="kn">import</span><span class="w"> </span><span class="nn">logging</span>
278
+ <span class="kn">from</span><span class="w"> </span><span class="nn">typing</span><span class="w"> </span><span class="kn">import</span> <span class="n">Any</span><span class="p">,</span> <span class="n">Iterable</span><span class="p">,</span> <span class="n">List</span><span class="p">,</span> <span class="n">Optional</span><span class="p">,</span> <span class="n">TypeVar</span><span class="p">,</span> <span class="n">Union</span>
279
+
280
+ <span class="kn">from</span><span class="w"> </span><span class="nn">bson</span><span class="w"> </span><span class="kn">import</span> <span class="n">ObjectId</span>
281
+ <span class="kn">from</span><span class="w"> </span><span class="nn">motor.motor_asyncio</span><span class="w"> </span><span class="kn">import</span> <span class="n">AsyncIOMotorClient</span>
282
+ <span class="kn">from</span><span class="w"> </span><span class="nn">odmantic</span><span class="w"> </span><span class="kn">import</span> <span class="n">AIOEngine</span>
283
+ <span class="kn">from</span><span class="w"> </span><span class="nn">odmantic</span><span class="w"> </span><span class="kn">import</span> <span class="n">Model</span><span class="p">,</span> <span class="n">EmbeddedModel</span>
284
+
285
+ <span class="kn">from</span><span class="w"> </span><span class="nn">simstack.core.definitions</span><span class="w"> </span><span class="kn">import</span> <span class="n">DBType</span><span class="p">,</span> <span class="n">TaskStatus</span>
286
+ <span class="kn">from</span><span class="w"> </span><span class="nn">simstack.models.node_registry</span><span class="w"> </span><span class="kn">import</span> <span class="n">NodeRegistry</span>
287
+ <span class="kn">from</span><span class="w"> </span><span class="nn">simstack.util.database_information</span><span class="w"> </span><span class="kn">import</span> <span class="n">DatabaseInformation</span>
288
+
289
+ <span class="c1"># from simstack.util.importer import import_class</span>
290
+
291
+ <span class="n">logger</span> <span class="o">=</span> <span class="n">logging</span><span class="o">.</span><span class="n">getLogger</span><span class="p">(</span><span class="vm">__name__</span><span class="p">)</span>
292
+
293
+ <span class="n">T</span> <span class="o">=</span> <span class="n">TypeVar</span><span class="p">(</span><span class="s2">&quot;T&quot;</span><span class="p">,</span> <span class="n">bound</span><span class="o">=</span><span class="n">Model</span><span class="p">)</span>
294
+
295
+
296
+ <div class="viewcode-block" id="Database">
297
+ <a class="viewcode-back" href="../../../simstack/modules/simstack.util.html#simstack.util.db.Database">[docs]</a>
298
+ <span class="k">class</span><span class="w"> </span><span class="nc">Database</span><span class="p">:</span>
299
+ <span class="w"> </span><span class="sd">&quot;&quot;&quot;Server-owned database facade.</span>
300
+
301
+ <span class="sd"> The server should depend on this object instead of reaching into SimStack</span>
302
+ <span class="sd"> core engine abstractions directly. The underlying persistence primitive is</span>
303
+ <span class="sd"> the plain ODMantic engine; SimStack-specific behavior lives in this facade.</span>
304
+ <span class="sd"> &quot;&quot;&quot;</span>
305
+
306
+ <span class="k">def</span><span class="w"> </span><span class="fm">__init__</span><span class="p">(</span>
307
+ <span class="bp">self</span><span class="p">,</span>
308
+ <span class="o">*</span><span class="p">,</span>
309
+ <span class="n">client</span><span class="p">:</span> <span class="n">AsyncIOMotorClient</span> <span class="o">|</span> <span class="kc">None</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
310
+ <span class="n">database_name</span><span class="p">:</span> <span class="nb">str</span> <span class="o">|</span> <span class="kc">None</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
311
+ <span class="n">connection_string</span><span class="p">:</span> <span class="nb">str</span> <span class="o">|</span> <span class="kc">None</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
312
+ <span class="n">engine</span><span class="p">:</span> <span class="n">Any</span> <span class="o">|</span> <span class="kc">None</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
313
+ <span class="n">db_type</span><span class="p">:</span> <span class="n">DBType</span> <span class="o">|</span> <span class="kc">None</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
314
+ <span class="n">server_url</span><span class="p">:</span> <span class="nb">str</span> <span class="o">|</span> <span class="kc">None</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
315
+ <span class="n">server_token</span><span class="p">:</span> <span class="nb">str</span> <span class="o">|</span> <span class="kc">None</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
316
+ <span class="p">)</span> <span class="o">-&gt;</span> <span class="kc">None</span><span class="p">:</span>
317
+ <span class="k">if</span> <span class="n">engine</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
318
+ <span class="k">if</span> <span class="n">client</span> <span class="ow">is</span> <span class="kc">None</span> <span class="ow">or</span> <span class="n">database_name</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
319
+ <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span>
320
+ <span class="s2">&quot;client and database_name are required when engine is not provided&quot;</span>
321
+ <span class="p">)</span>
322
+ <span class="n">engine</span> <span class="o">=</span> <span class="n">AIOEngine</span><span class="p">(</span><span class="n">client</span><span class="o">=</span><span class="n">client</span><span class="p">,</span> <span class="n">database</span><span class="o">=</span><span class="n">database_name</span><span class="p">)</span>
323
+
324
+ <span class="bp">self</span><span class="o">.</span><span class="n">_db_type</span> <span class="o">=</span> <span class="n">db_type</span> <span class="ow">or</span> <span class="nb">getattr</span><span class="p">(</span><span class="n">engine</span><span class="p">,</span> <span class="s2">&quot;db_type&quot;</span><span class="p">,</span> <span class="kc">None</span><span class="p">)</span>
325
+ <span class="bp">self</span><span class="o">.</span><span class="n">_engine</span> <span class="o">=</span> <span class="n">engine</span>
326
+ <span class="bp">self</span><span class="o">.</span><span class="n">_client</span> <span class="o">=</span> <span class="n">client</span> <span class="ow">or</span> <span class="nb">getattr</span><span class="p">(</span><span class="n">engine</span><span class="p">,</span> <span class="s2">&quot;client&quot;</span><span class="p">,</span> <span class="kc">None</span><span class="p">)</span>
327
+ <span class="bp">self</span><span class="o">.</span><span class="n">_database_name</span> <span class="o">=</span> <span class="n">database_name</span> <span class="ow">or</span> <span class="nb">getattr</span><span class="p">(</span><span class="n">engine</span><span class="p">,</span> <span class="s2">&quot;database_name&quot;</span><span class="p">,</span> <span class="kc">None</span><span class="p">)</span>
328
+ <span class="bp">self</span><span class="o">.</span><span class="n">_connection_string</span> <span class="o">=</span> <span class="n">connection_string</span> <span class="ow">or</span> <span class="nb">getattr</span><span class="p">(</span><span class="n">engine</span><span class="p">,</span> <span class="s2">&quot;connection_string&quot;</span><span class="p">,</span> <span class="kc">None</span><span class="p">)</span>
329
+ <span class="bp">self</span><span class="o">.</span><span class="n">_server_url</span> <span class="o">=</span> <span class="n">server_url</span>
330
+ <span class="bp">self</span><span class="o">.</span><span class="n">_server_token</span> <span class="o">=</span> <span class="n">server_token</span>
331
+
332
+ <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">_database_name</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
333
+ <span class="n">database</span> <span class="o">=</span> <span class="nb">getattr</span><span class="p">(</span><span class="n">engine</span><span class="p">,</span> <span class="s2">&quot;database&quot;</span><span class="p">,</span> <span class="kc">None</span><span class="p">)</span>
334
+ <span class="bp">self</span><span class="o">.</span><span class="n">_database_name</span> <span class="o">=</span> <span class="nb">getattr</span><span class="p">(</span><span class="n">database</span><span class="p">,</span> <span class="s2">&quot;name&quot;</span><span class="p">,</span> <span class="kc">None</span><span class="p">)</span>
335
+
336
+ <span class="nd">@property</span>
337
+ <span class="k">def</span><span class="w"> </span><span class="nf">databae_type</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">DBType</span><span class="p">:</span>
338
+ <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_db_type</span>
339
+
340
+ <span class="nd">@property</span>
341
+ <span class="k">def</span><span class="w"> </span><span class="nf">connection_string</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="nb">str</span> <span class="o">|</span> <span class="kc">None</span><span class="p">:</span>
342
+ <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_connection_string</span>
343
+
344
+ <div class="viewcode-block" id="Database.from_db_info">
345
+ <a class="viewcode-back" href="../../../simstack/modules/simstack.util.html#simstack.util.db.Database.from_db_info">[docs]</a>
346
+ <span class="nd">@classmethod</span>
347
+ <span class="k">def</span><span class="w"> </span><span class="nf">from_db_info</span><span class="p">(</span><span class="bp">cls</span><span class="p">,</span> <span class="n">db_info</span><span class="p">:</span> <span class="n">DatabaseInformation</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="s2">&quot;Database&quot;</span><span class="p">:</span>
348
+ <span class="k">if</span> <span class="n">db_info</span><span class="o">.</span><span class="n">db_type</span> <span class="o">==</span> <span class="n">DBType</span><span class="o">.</span><span class="n">IN_MEMORY</span><span class="p">:</span>
349
+ <span class="c1"># For tests, use in-memory MongoDB (mongomock)</span>
350
+ <span class="kn">from</span><span class="w"> </span><span class="nn">mongomock_motor</span><span class="w"> </span><span class="kn">import</span> <span class="n">AsyncMongoMockClient</span>
351
+
352
+ <span class="k">try</span><span class="p">:</span>
353
+ <span class="c1"># import mongomock</span>
354
+ <span class="n">client</span> <span class="o">=</span> <span class="n">AsyncMongoMockClient</span><span class="p">()</span>
355
+ <span class="n">logger</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;Using in-memory MongoDB mock&quot;</span><span class="p">)</span>
356
+ <span class="k">except</span> <span class="ne">ImportError</span><span class="p">:</span>
357
+ <span class="n">logger</span><span class="o">.</span><span class="n">warning</span><span class="p">(</span>
358
+ <span class="s2">&quot;mongomock not installed, falling back to localhost MongoDB&quot;</span>
359
+ <span class="p">)</span>
360
+ <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span>
361
+ <span class="s2">&quot;mongomock not installed, cannot use in-memory MongoDB&quot;</span>
362
+ <span class="p">)</span>
363
+
364
+ <span class="k">elif</span> <span class="n">db_info</span><span class="o">.</span><span class="n">db_type</span> <span class="o">==</span> <span class="n">DBType</span><span class="o">.</span><span class="n">MONGODB</span><span class="p">:</span>
365
+ <span class="n">connection_string</span> <span class="o">=</span> <span class="n">db_info</span><span class="o">.</span><span class="n">connection_string</span>
366
+ <span class="k">if</span> <span class="ow">not</span> <span class="n">connection_string</span><span class="p">:</span>
367
+ <span class="n">connection_string</span> <span class="o">=</span> <span class="s2">&quot;mongodb://localhost:27017&quot;</span>
368
+ <span class="n">client</span> <span class="o">=</span> <span class="n">AsyncIOMotorClient</span><span class="p">(</span><span class="n">connection_string</span><span class="p">)</span>
369
+ <span class="n">logger</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="s2">&quot;Connected to MongoDB&quot;</span><span class="p">)</span>
370
+ <span class="k">else</span><span class="p">:</span>
371
+ <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span>
372
+ <span class="sa">f</span><span class="s2">&quot;Unsupported database type for MongoDB: </span><span class="si">{</span><span class="n">db_info</span><span class="o">.</span><span class="n">db_type</span><span class="si">}</span><span class="s2">&quot;</span>
373
+ <span class="p">)</span>
374
+
375
+ <span class="c1"># Create engine</span>
376
+ <span class="n">engine</span> <span class="o">=</span> <span class="n">AIOEngine</span><span class="p">(</span><span class="n">client</span><span class="o">=</span><span class="n">client</span><span class="p">,</span> <span class="n">database</span><span class="o">=</span><span class="n">db_info</span><span class="o">.</span><span class="n">db_name</span><span class="p">)</span>
377
+ <span class="k">return</span> <span class="bp">cls</span><span class="p">(</span>
378
+ <span class="n">engine</span><span class="o">=</span><span class="n">engine</span><span class="p">,</span>
379
+ <span class="n">client</span><span class="o">=</span><span class="n">client</span><span class="p">,</span>
380
+ <span class="n">database_name</span><span class="o">=</span><span class="n">db_info</span><span class="o">.</span><span class="n">db_name</span><span class="p">,</span>
381
+ <span class="n">connection_string</span><span class="o">=</span><span class="n">db_info</span><span class="o">.</span><span class="n">connection_string</span><span class="p">,</span>
382
+ <span class="n">db_type</span><span class="o">=</span><span class="n">db_info</span><span class="o">.</span><span class="n">db_type</span><span class="p">,</span>
383
+ <span class="n">server_url</span><span class="o">=</span><span class="n">db_info</span><span class="o">.</span><span class="n">server_url</span><span class="p">,</span>
384
+ <span class="n">server_token</span><span class="o">=</span><span class="n">db_info</span><span class="o">.</span><span class="n">server_token</span><span class="p">,</span>
385
+ <span class="p">)</span></div>
386
+
387
+
388
+ <span class="nd">@property</span>
389
+ <span class="k">def</span><span class="w"> </span><span class="nf">core_engine</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Any</span><span class="p">:</span>
390
+ <span class="w"> </span><span class="sd">&quot;&quot;&quot;Compatibility escape hatch for SimStack core internals only.&quot;&quot;&quot;</span>
391
+ <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_engine</span>
392
+
393
+ <span class="nd">@property</span>
394
+ <span class="k">def</span><span class="w"> </span><span class="nf">server_url</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="nb">str</span> <span class="o">|</span> <span class="kc">None</span><span class="p">:</span>
395
+ <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_server_url</span>
396
+
397
+ <span class="nd">@property</span>
398
+ <span class="k">def</span><span class="w"> </span><span class="nf">server_token</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="nb">str</span> <span class="o">|</span> <span class="kc">None</span><span class="p">:</span>
399
+ <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_server_token</span>
400
+
401
+ <span class="nd">@property</span>
402
+ <span class="k">def</span><span class="w"> </span><span class="nf">client</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">AsyncIOMotorClient</span><span class="p">:</span>
403
+ <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">_client</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
404
+ <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="s2">&quot;Database client is not initialized&quot;</span><span class="p">)</span>
405
+ <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_client</span>
406
+
407
+ <span class="nd">@property</span>
408
+ <span class="k">def</span><span class="w"> </span><span class="nf">database_name</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="nb">str</span><span class="p">:</span>
409
+ <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">_database_name</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
410
+ <span class="k">raise</span> <span class="ne">RuntimeError</span><span class="p">(</span><span class="s2">&quot;Database name is not initialized&quot;</span><span class="p">)</span>
411
+ <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_database_name</span>
412
+
413
+ <span class="nd">@property</span>
414
+ <span class="k">def</span><span class="w"> </span><span class="nf">raw_database</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Any</span><span class="p">:</span>
415
+ <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">_client</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">and</span> <span class="bp">self</span><span class="o">.</span><span class="n">_database_name</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
416
+ <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_client</span><span class="p">[</span><span class="bp">self</span><span class="o">.</span><span class="n">_database_name</span><span class="p">]</span>
417
+ <span class="k">return</span> <span class="nb">getattr</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">_engine</span><span class="p">,</span> <span class="s2">&quot;database&quot;</span><span class="p">)</span>
418
+
419
+ <span class="nd">@property</span>
420
+ <span class="k">def</span><span class="w"> </span><span class="nf">database</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Any</span><span class="p">:</span>
421
+ <span class="k">return</span> <span class="nb">getattr</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">_engine</span><span class="p">,</span> <span class="s2">&quot;database&quot;</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">raw_database</span><span class="p">)</span>
422
+
423
+ <div class="viewcode-block" id="Database.collection">
424
+ <a class="viewcode-back" href="../../../simstack/modules/simstack.util.html#simstack.util.db.Database.collection">[docs]</a>
425
+ <span class="k">def</span><span class="w"> </span><span class="nf">collection</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">model_or_name</span><span class="p">:</span> <span class="n">Any</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Any</span><span class="p">:</span>
426
+ <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">model_or_name</span><span class="p">,</span> <span class="nb">str</span><span class="p">):</span>
427
+ <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">raw_database</span><span class="p">[</span><span class="n">model_or_name</span><span class="p">]</span>
428
+ <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_engine</span><span class="o">.</span><span class="n">get_collection</span><span class="p">(</span><span class="n">model_or_name</span><span class="p">)</span></div>
429
+
430
+
431
+ <div class="viewcode-block" id="Database.get_collection">
432
+ <a class="viewcode-back" href="../../../simstack/modules/simstack.util.html#simstack.util.db.Database.get_collection">[docs]</a>
433
+ <span class="k">def</span><span class="w"> </span><span class="nf">get_collection</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">model_or_name</span><span class="p">:</span> <span class="n">Any</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Any</span><span class="p">:</span>
434
+ <span class="w"> </span><span class="sd">&quot;&quot;&quot;Temporary compatibility alias for code still being migrated.&quot;&quot;&quot;</span>
435
+ <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">collection</span><span class="p">(</span><span class="n">model_or_name</span><span class="p">)</span></div>
436
+
437
+
438
+ <div class="viewcode-block" id="Database.find">
439
+ <a class="viewcode-back" href="../../../simstack/modules/simstack.util.html#simstack.util.db.Database.find">[docs]</a>
440
+ <span class="k">async</span> <span class="k">def</span><span class="w"> </span><span class="nf">find</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">:</span> <span class="n">Any</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">:</span> <span class="n">Any</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Any</span><span class="p">:</span>
441
+ <span class="k">if</span> <span class="ow">not</span> <span class="n">args</span><span class="p">:</span>
442
+ <span class="c1"># Re-raise TypeError as expected by tests if no args provided</span>
443
+ <span class="k">raise</span> <span class="ne">TypeError</span><span class="p">(</span>
444
+ <span class="s2">&quot;AIOEngine.find() missing 1 required positional argument: &#39;model&#39;&quot;</span>
445
+ <span class="p">)</span>
446
+
447
+ <span class="k">try</span><span class="p">:</span>
448
+ <span class="n">results</span> <span class="o">=</span> <span class="k">await</span> <span class="bp">self</span><span class="o">.</span><span class="n">_engine</span><span class="o">.</span><span class="n">find</span><span class="p">(</span><span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>
449
+ <span class="k">except</span> <span class="ne">TypeError</span> <span class="k">as</span> <span class="n">e</span><span class="p">:</span>
450
+ <span class="k">if</span> <span class="s2">&quot;Can only call find with a Model class&quot;</span> <span class="ow">in</span> <span class="nb">str</span><span class="p">(</span><span class="n">e</span><span class="p">):</span>
451
+ <span class="n">logger</span><span class="o">.</span><span class="n">warning</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;Engine.find TypeError bypassed for </span><span class="si">{</span><span class="n">args</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">)</span>
452
+ <span class="k">return</span> <span class="p">[]</span>
453
+ <span class="k">raise</span>
454
+
455
+ <span class="c1"># Handle possible wrapped return from mongomock/odmantic</span>
456
+ <span class="kn">import</span><span class="w"> </span><span class="nn">inspect</span>
457
+ <span class="k">if</span> <span class="n">inspect</span><span class="o">.</span><span class="n">isawaitable</span><span class="p">(</span><span class="n">results</span><span class="p">):</span>
458
+ <span class="n">results</span> <span class="o">=</span> <span class="k">await</span> <span class="n">results</span>
459
+
460
+ <span class="k">try</span><span class="p">:</span>
461
+ <span class="n">model_class</span> <span class="o">=</span> <span class="n">args</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
462
+ <span class="kn">from</span><span class="w"> </span><span class="nn">odmantic</span><span class="w"> </span><span class="kn">import</span> <span class="n">Model</span>
463
+ <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">model_class</span><span class="p">,</span> <span class="nb">type</span><span class="p">)</span> <span class="ow">and</span> <span class="nb">issubclass</span><span class="p">(</span><span class="n">model_class</span><span class="p">,</span> <span class="n">Model</span><span class="p">):</span>
464
+ <span class="k">for</span> <span class="n">result</span> <span class="ow">in</span> <span class="n">results</span><span class="p">:</span>
465
+ <span class="k">await</span> <span class="bp">self</span><span class="o">.</span><span class="n">_apply_postprocess</span><span class="p">(</span><span class="n">model_class</span><span class="p">,</span> <span class="n">result</span><span class="p">)</span>
466
+ <span class="k">except</span> <span class="ne">Exception</span><span class="p">:</span>
467
+ <span class="k">pass</span>
468
+ <span class="k">return</span> <span class="n">results</span></div>
469
+
470
+
471
+ <div class="viewcode-block" id="Database.find_all">
472
+ <a class="viewcode-back" href="../../../simstack/modules/simstack.util.html#simstack.util.db.Database.find_all">[docs]</a>
473
+ <span class="k">async</span> <span class="k">def</span><span class="w"> </span><span class="nf">find_all</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">:</span> <span class="n">Any</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">:</span> <span class="n">Any</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Any</span><span class="p">:</span>
474
+ <span class="k">return</span> <span class="k">await</span> <span class="bp">self</span><span class="o">.</span><span class="n">find</span><span class="p">(</span><span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span></div>
475
+
476
+
477
+ <div class="viewcode-block" id="Database.find_one">
478
+ <a class="viewcode-back" href="../../../simstack/modules/simstack.util.html#simstack.util.db.Database.find_one">[docs]</a>
479
+ <span class="k">async</span> <span class="k">def</span><span class="w"> </span><span class="nf">find_one</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">:</span> <span class="n">Any</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">:</span> <span class="n">Any</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Any</span><span class="p">:</span>
480
+ <span class="k">if</span> <span class="ow">not</span> <span class="n">args</span><span class="p">:</span>
481
+ <span class="c1"># Re-raise TypeError as expected by tests if no args provided</span>
482
+ <span class="k">raise</span> <span class="ne">TypeError</span><span class="p">(</span>
483
+ <span class="s2">&quot;Database.find_one() missing 1 required positional argument: &#39;model&#39;&quot;</span>
484
+ <span class="p">)</span>
485
+
486
+ <span class="k">try</span><span class="p">:</span>
487
+ <span class="n">result</span> <span class="o">=</span> <span class="k">await</span> <span class="bp">self</span><span class="o">.</span><span class="n">_engine</span><span class="o">.</span><span class="n">find_one</span><span class="p">(</span><span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>
488
+ <span class="k">except</span> <span class="ne">TypeError</span> <span class="k">as</span> <span class="n">e</span><span class="p">:</span>
489
+ <span class="k">if</span> <span class="s2">&quot;Can only call find with a Model class&quot;</span> <span class="ow">in</span> <span class="nb">str</span><span class="p">(</span><span class="n">e</span><span class="p">):</span>
490
+ <span class="n">logger</span><span class="o">.</span><span class="n">warning</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;Engine.find_one TypeError bypassed for </span><span class="si">{</span><span class="n">args</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">)</span>
491
+ <span class="k">return</span> <span class="kc">None</span>
492
+ <span class="k">raise</span>
493
+
494
+ <span class="kn">import</span><span class="w"> </span><span class="nn">inspect</span>
495
+ <span class="k">if</span> <span class="n">inspect</span><span class="o">.</span><span class="n">isawaitable</span><span class="p">(</span><span class="n">result</span><span class="p">):</span>
496
+ <span class="n">result</span> <span class="o">=</span> <span class="k">await</span> <span class="n">result</span>
497
+
498
+ <span class="k">if</span> <span class="n">result</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
499
+ <span class="k">try</span><span class="p">:</span>
500
+ <span class="n">model_class</span> <span class="o">=</span> <span class="n">args</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
501
+ <span class="kn">from</span><span class="w"> </span><span class="nn">odmantic</span><span class="w"> </span><span class="kn">import</span> <span class="n">Model</span>
502
+ <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">model_class</span><span class="p">,</span> <span class="nb">type</span><span class="p">)</span> <span class="ow">and</span> <span class="nb">issubclass</span><span class="p">(</span><span class="n">model_class</span><span class="p">,</span> <span class="n">Model</span><span class="p">):</span>
503
+ <span class="k">await</span> <span class="bp">self</span><span class="o">.</span><span class="n">_apply_postprocess</span><span class="p">(</span><span class="n">model_class</span><span class="p">,</span> <span class="n">result</span><span class="p">)</span>
504
+ <span class="k">except</span> <span class="ne">Exception</span><span class="p">:</span>
505
+ <span class="k">pass</span>
506
+ <span class="k">return</span> <span class="n">result</span></div>
507
+
508
+
509
+ <div class="viewcode-block" id="Database.find_one_by_model_name">
510
+ <a class="viewcode-back" href="../../../simstack/modules/simstack.util.html#simstack.util.db.Database.find_one_by_model_name">[docs]</a>
511
+ <span class="k">async</span> <span class="k">def</span><span class="w"> </span><span class="nf">find_one_by_model_name</span><span class="p">(</span>
512
+ <span class="bp">self</span><span class="p">,</span> <span class="n">model_mapping</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span> <span class="n">item_id</span><span class="p">:</span> <span class="nb">str</span> <span class="o">|</span> <span class="n">ObjectId</span>
513
+ <span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Any</span><span class="p">]:</span>
514
+ <span class="k">if</span> <span class="s2">&quot;.&quot;</span> <span class="ow">in</span> <span class="n">model_mapping</span><span class="p">:</span>
515
+ <span class="kn">from</span><span class="w"> </span><span class="nn">simstack.util.importer</span><span class="w"> </span><span class="kn">import</span> <span class="n">import_class</span>
516
+
517
+ <span class="n">model_class</span> <span class="o">=</span> <span class="k">await</span> <span class="n">import_class</span><span class="p">(</span><span class="n">model_mapping</span><span class="p">,</span> <span class="bp">self</span><span class="p">)</span>
518
+ <span class="k">else</span><span class="p">:</span>
519
+ <span class="kn">from</span><span class="w"> </span><span class="nn">simstack.util.importer</span><span class="w"> </span><span class="kn">import</span> <span class="n">import_class_by_name</span>
520
+
521
+ <span class="n">model_class</span> <span class="o">=</span> <span class="k">await</span> <span class="n">import_class_by_name</span><span class="p">(</span><span class="n">model_mapping</span><span class="p">,</span> <span class="bp">self</span><span class="p">)</span>
522
+
523
+ <span class="k">if</span> <span class="n">model_class</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
524
+ <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span>
525
+ <span class="sa">f</span><span class="s2">&quot;DB: model class </span><span class="si">{</span><span class="n">model_mapping</span><span class="si">}</span><span class="s2"> not found in the available modules&quot;</span>
526
+ <span class="p">)</span>
527
+
528
+ <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">item_id</span><span class="p">,</span> <span class="nb">str</span><span class="p">):</span>
529
+ <span class="n">item_id</span> <span class="o">=</span> <span class="n">ObjectId</span><span class="p">(</span><span class="n">item_id</span><span class="p">)</span>
530
+
531
+ <span class="n">instance</span> <span class="o">=</span> <span class="k">await</span> <span class="bp">self</span><span class="o">.</span><span class="n">find_one</span><span class="p">(</span><span class="n">model_class</span><span class="p">,</span> <span class="n">model_class</span><span class="o">.</span><span class="n">id</span> <span class="o">==</span> <span class="n">item_id</span><span class="p">)</span>
532
+ <span class="k">if</span> <span class="n">instance</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
533
+ <span class="n">logger</span><span class="o">.</span><span class="n">error</span><span class="p">(</span>
534
+ <span class="sa">f</span><span class="s2">&quot;Instance of &#39;</span><span class="si">{</span><span class="n">model_class</span><span class="o">.</span><span class="vm">__name__</span><span class="si">}</span><span class="s2">&#39; with id &#39;</span><span class="si">{</span><span class="n">item_id</span><span class="si">}</span><span class="s2">&#39; does not exist&quot;</span>
535
+ <span class="p">)</span>
536
+ <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span>
537
+ <span class="sa">f</span><span class="s2">&quot;Instance of &#39;</span><span class="si">{</span><span class="n">model_class</span><span class="o">.</span><span class="vm">__name__</span><span class="si">}</span><span class="s2">&#39; with id &#39;</span><span class="si">{</span><span class="n">item_id</span><span class="si">}</span><span class="s2">&#39; does not exist&quot;</span>
538
+ <span class="p">)</span>
539
+ <span class="k">return</span> <span class="n">instance</span></div>
540
+
541
+
542
+ <span class="k">async</span> <span class="k">def</span><span class="w"> </span><span class="nf">_apply_postprocess</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">model_class</span><span class="p">:</span> <span class="nb">type</span><span class="p">,</span> <span class="n">result</span><span class="p">:</span> <span class="n">Model</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="kc">None</span><span class="p">:</span>
543
+ <span class="w"> </span><span class="sd">&quot;&quot;&quot;Apply db_find_postprocess to results if defined on the model class.&quot;&quot;&quot;</span>
544
+ <span class="n">post_process</span> <span class="o">=</span> <span class="nb">getattr</span><span class="p">(</span><span class="n">model_class</span><span class="p">,</span> <span class="s2">&quot;db_find_postprocess&quot;</span><span class="p">,</span> <span class="kc">None</span><span class="p">)</span>
545
+
546
+ <span class="k">if</span> <span class="n">post_process</span> <span class="ow">and</span> <span class="nb">callable</span><span class="p">(</span><span class="n">post_process</span><span class="p">):</span>
547
+ <span class="kn">import</span><span class="w"> </span><span class="nn">inspect</span>
548
+
549
+ <span class="n">is_async</span> <span class="o">=</span> <span class="n">inspect</span><span class="o">.</span><span class="n">iscoroutinefunction</span><span class="p">(</span><span class="n">post_process</span><span class="p">)</span>
550
+
551
+ <span class="k">try</span><span class="p">:</span>
552
+ <span class="k">if</span> <span class="n">is_async</span><span class="p">:</span>
553
+ <span class="k">await</span> <span class="n">post_process</span><span class="p">(</span><span class="n">result</span><span class="p">,</span> <span class="bp">self</span><span class="p">)</span>
554
+ <span class="k">else</span><span class="p">:</span>
555
+ <span class="n">logger</span><span class="o">.</span><span class="n">warning</span><span class="p">(</span>
556
+ <span class="sa">f</span><span class="s2">&quot;Sync db_find_postprocess on </span><span class="si">{</span><span class="n">model_class</span><span class="o">.</span><span class="vm">__name__</span><span class="si">}</span><span class="s2"> is deprecated. &quot;</span>
557
+ <span class="s2">&quot;Please make it asynchronous.&quot;</span>
558
+ <span class="p">)</span>
559
+ <span class="n">post_process</span><span class="p">(</span><span class="n">result</span><span class="p">,</span> <span class="bp">self</span><span class="p">)</span>
560
+ <span class="k">except</span> <span class="ne">Exception</span> <span class="k">as</span> <span class="n">e</span><span class="p">:</span>
561
+ <span class="n">logger</span><span class="o">.</span><span class="n">error</span><span class="p">(</span>
562
+ <span class="sa">f</span><span class="s2">&quot;Error during post-processing </span><span class="si">{</span><span class="n">model_class</span><span class="o">.</span><span class="vm">__name__</span><span class="si">}</span><span class="s2">: </span><span class="si">{</span><span class="n">e</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">,</span>
563
+ <span class="n">exc_info</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
564
+ <span class="p">)</span>
565
+ <span class="c1"># Should we re-raise or just log? Usually post-processing failure is critical.</span>
566
+ <span class="k">raise</span>
567
+
568
+ <span class="c1"># Recursively apply postprocess to nested Model attributes</span>
569
+ <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">result</span><span class="p">,</span> <span class="n">Model</span><span class="p">)</span> <span class="ow">or</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">result</span><span class="p">,</span> <span class="n">EmbeddedModel</span><span class="p">):</span>
570
+ <span class="k">return</span>
571
+
572
+ <span class="k">try</span><span class="p">:</span>
573
+ <span class="k">for</span> <span class="n">attr_name</span> <span class="ow">in</span> <span class="nb">dir</span><span class="p">(</span><span class="n">result</span><span class="p">):</span>
574
+ <span class="k">if</span> <span class="n">attr_name</span><span class="o">.</span><span class="n">startswith</span><span class="p">(</span><span class="s2">&quot;_&quot;</span><span class="p">):</span>
575
+ <span class="k">continue</span>
576
+
577
+ <span class="k">try</span><span class="p">:</span>
578
+ <span class="n">attr_value</span> <span class="o">=</span> <span class="nb">getattr</span><span class="p">(</span><span class="n">result</span><span class="p">,</span> <span class="n">attr_name</span><span class="p">,</span> <span class="kc">None</span><span class="p">)</span>
579
+ <span class="k">except</span> <span class="ne">Exception</span><span class="p">:</span>
580
+ <span class="k">continue</span>
581
+
582
+ <span class="k">if</span> <span class="n">attr_value</span> <span class="ow">is</span> <span class="kc">None</span> <span class="ow">or</span> <span class="nb">callable</span><span class="p">(</span><span class="n">attr_value</span><span class="p">):</span>
583
+ <span class="k">continue</span>
584
+
585
+ <span class="c1"># Handle direct Model instance</span>
586
+ <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">attr_value</span><span class="p">,</span> <span class="n">Model</span><span class="p">)</span> <span class="ow">or</span> <span class="nb">isinstance</span><span class="p">(</span>
587
+ <span class="n">attr_value</span><span class="p">,</span> <span class="n">EmbeddedModel</span>
588
+ <span class="p">):</span>
589
+ <span class="k">await</span> <span class="bp">self</span><span class="o">.</span><span class="n">_apply_postprocess</span><span class="p">(</span><span class="nb">type</span><span class="p">(</span><span class="n">attr_value</span><span class="p">),</span> <span class="n">attr_value</span><span class="p">)</span>
590
+
591
+ <span class="c1"># Handle list of Models</span>
592
+ <span class="k">elif</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">attr_value</span><span class="p">,</span> <span class="p">(</span><span class="nb">list</span><span class="p">,</span> <span class="nb">tuple</span><span class="p">)):</span>
593
+ <span class="k">for</span> <span class="n">item</span> <span class="ow">in</span> <span class="n">attr_value</span><span class="p">:</span>
594
+ <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">item</span><span class="p">,</span> <span class="n">Model</span><span class="p">)</span> <span class="ow">or</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">item</span><span class="p">,</span> <span class="n">EmbeddedModel</span><span class="p">):</span>
595
+ <span class="k">await</span> <span class="bp">self</span><span class="o">.</span><span class="n">_apply_postprocess</span><span class="p">(</span><span class="nb">type</span><span class="p">(</span><span class="n">item</span><span class="p">),</span> <span class="n">item</span><span class="p">)</span>
596
+
597
+ <span class="c1"># Handle dict with Model values</span>
598
+ <span class="k">elif</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">attr_value</span><span class="p">,</span> <span class="nb">dict</span><span class="p">):</span>
599
+ <span class="k">for</span> <span class="n">value</span> <span class="ow">in</span> <span class="n">attr_value</span><span class="o">.</span><span class="n">values</span><span class="p">():</span>
600
+ <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">value</span><span class="p">,</span> <span class="n">Model</span><span class="p">)</span> <span class="ow">or</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">value</span><span class="p">,</span> <span class="n">EmbeddedModel</span><span class="p">):</span>
601
+ <span class="k">await</span> <span class="bp">self</span><span class="o">.</span><span class="n">_apply_postprocess</span><span class="p">(</span><span class="nb">type</span><span class="p">(</span><span class="n">value</span><span class="p">),</span> <span class="n">value</span><span class="p">)</span>
602
+
603
+ <span class="k">except</span> <span class="ne">Exception</span> <span class="k">as</span> <span class="n">e</span><span class="p">:</span>
604
+ <span class="n">logger</span><span class="o">.</span><span class="n">error</span><span class="p">(</span>
605
+ <span class="sa">f</span><span class="s2">&quot;Error during recursive post-processing traversal for </span><span class="si">{</span><span class="n">model_class</span><span class="o">.</span><span class="vm">__name__</span><span class="si">}</span><span class="s2">: </span><span class="si">{</span><span class="n">e</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">,</span>
606
+ <span class="n">exc_info</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
607
+ <span class="p">)</span>
608
+
609
+ <div class="viewcode-block" id="Database.close">
610
+ <a class="viewcode-back" href="../../../simstack/modules/simstack.util.html#simstack.util.db.Database.close">[docs]</a>
611
+ <span class="k">async</span> <span class="k">def</span><span class="w"> </span><span class="nf">close</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="kc">None</span><span class="p">:</span>
612
+ <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">_client</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
613
+ <span class="bp">self</span><span class="o">.</span><span class="n">_client</span><span class="o">.</span><span class="n">close</span><span class="p">()</span></div>
614
+
615
+
616
+ <div class="viewcode-block" id="Database.save">
617
+ <a class="viewcode-back" href="../../../simstack/modules/simstack.util.html#simstack.util.db.Database.save">[docs]</a>
618
+ <span class="k">async</span> <span class="k">def</span><span class="w"> </span><span class="nf">save</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">:</span> <span class="n">Any</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">:</span> <span class="n">Any</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Any</span><span class="p">:</span>
619
+ <span class="k">if</span> <span class="ow">not</span> <span class="n">args</span><span class="p">:</span>
620
+ <span class="c1"># Re-raise TypeError as expected by tests if no args provided</span>
621
+ <span class="k">raise</span> <span class="ne">TypeError</span><span class="p">(</span>
622
+ <span class="s2">&quot;AIOEngine.save() missing 1 required positional argument: &#39;instance&#39;&quot;</span>
623
+ <span class="p">)</span>
624
+
625
+ <span class="n">obj</span> <span class="o">=</span> <span class="n">args</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
626
+ <span class="n">rest_args</span> <span class="o">=</span> <span class="n">args</span><span class="p">[</span><span class="mi">1</span><span class="p">:]</span>
627
+ <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">obj</span><span class="p">,</span> <span class="p">(</span><span class="nb">list</span><span class="p">,</span> <span class="nb">tuple</span><span class="p">,</span> <span class="nb">set</span><span class="p">)):</span>
628
+ <span class="c1"># Use list comprehension to collect results and ensure they are returned</span>
629
+ <span class="n">results</span> <span class="o">=</span> <span class="p">[]</span>
630
+ <span class="k">for</span> <span class="n">item</span> <span class="ow">in</span> <span class="n">obj</span><span class="p">:</span>
631
+ <span class="c1"># Assign ID if missing for in_memory/mock</span>
632
+ <span class="k">if</span> <span class="nb">hasattr</span><span class="p">(</span><span class="n">item</span><span class="p">,</span> <span class="s2">&quot;id&quot;</span><span class="p">)</span> <span class="ow">and</span> <span class="n">item</span><span class="o">.</span><span class="n">id</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
633
+ <span class="kn">from</span><span class="w"> </span><span class="nn">odmantic</span><span class="w"> </span><span class="kn">import</span> <span class="n">ObjectId</span>
634
+ <span class="n">item</span><span class="o">.</span><span class="n">id</span> <span class="o">=</span> <span class="n">ObjectId</span><span class="p">()</span>
635
+ <span class="n">res</span> <span class="o">=</span> <span class="k">await</span> <span class="bp">self</span><span class="o">.</span><span class="n">_save_one</span><span class="p">(</span><span class="n">item</span><span class="p">,</span> <span class="o">*</span><span class="n">rest_args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>
636
+ <span class="n">results</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">res</span><span class="p">)</span>
637
+ <span class="k">return</span> <span class="n">results</span>
638
+
639
+ <span class="c1"># Assign ID if missing for in_memory/mock</span>
640
+ <span class="k">if</span> <span class="nb">hasattr</span><span class="p">(</span><span class="n">obj</span><span class="p">,</span> <span class="s2">&quot;id&quot;</span><span class="p">)</span> <span class="ow">and</span> <span class="n">obj</span><span class="o">.</span><span class="n">id</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
641
+ <span class="kn">from</span><span class="w"> </span><span class="nn">odmantic</span><span class="w"> </span><span class="kn">import</span> <span class="n">ObjectId</span>
642
+ <span class="n">obj</span><span class="o">.</span><span class="n">id</span> <span class="o">=</span> <span class="n">ObjectId</span><span class="p">()</span>
643
+
644
+ <span class="n">result</span> <span class="o">=</span> <span class="k">await</span> <span class="bp">self</span><span class="o">.</span><span class="n">_save_one</span><span class="p">(</span><span class="n">obj</span><span class="p">,</span> <span class="o">*</span><span class="n">rest_args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>
645
+ <span class="k">return</span> <span class="n">result</span></div>
646
+
647
+
648
+ <div class="viewcode-block" id="Database.save_unchecked">
649
+ <a class="viewcode-back" href="../../../simstack/modules/simstack.util.html#simstack.util.db.Database.save_unchecked">[docs]</a>
650
+ <span class="k">async</span> <span class="k">def</span><span class="w"> </span><span class="nf">save_unchecked</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">:</span> <span class="n">Any</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">:</span> <span class="n">Any</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Any</span><span class="p">:</span>
651
+ <span class="k">return</span> <span class="k">await</span> <span class="bp">self</span><span class="o">.</span><span class="n">_engine</span><span class="o">.</span><span class="n">save</span><span class="p">(</span><span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span></div>
652
+
653
+
654
+ <div class="viewcode-block" id="Database.delete">
655
+ <a class="viewcode-back" href="../../../simstack/modules/simstack.util.html#simstack.util.db.Database.delete">[docs]</a>
656
+ <span class="k">async</span> <span class="k">def</span><span class="w"> </span><span class="nf">delete</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">:</span> <span class="n">Any</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">:</span> <span class="n">Any</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Any</span><span class="p">:</span>
657
+ <span class="k">return</span> <span class="k">await</span> <span class="bp">self</span><span class="o">.</span><span class="n">_engine</span><span class="o">.</span><span class="n">delete</span><span class="p">(</span><span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span></div>
658
+
659
+
660
+ <div class="viewcode-block" id="Database.apply_resource_assignment_to_node_registry">
661
+ <a class="viewcode-back" href="../../../simstack/modules/simstack.util.html#simstack.util.db.Database.apply_resource_assignment_to_node_registry">[docs]</a>
662
+ <span class="k">async</span> <span class="k">def</span><span class="w"> </span><span class="nf">apply_resource_assignment_to_node_registry</span><span class="p">(</span>
663
+ <span class="bp">self</span><span class="p">,</span> <span class="n">node_registry</span><span class="p">:</span> <span class="n">Any</span>
664
+ <span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Any</span><span class="p">:</span>
665
+ <span class="kn">from</span><span class="w"> </span><span class="nn">simstack.core.resource_assignment</span><span class="w"> </span><span class="kn">import</span> <span class="p">(</span>
666
+ <span class="n">apply_resource_assignment_to_node_registry</span><span class="p">,</span>
667
+ <span class="p">)</span>
668
+
669
+ <span class="k">return</span> <span class="k">await</span> <span class="n">apply_resource_assignment_to_node_registry</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">node_registry</span><span class="p">)</span></div>
670
+
671
+
672
+ <div class="viewcode-block" id="Database.find_artifact_mappings">
673
+ <a class="viewcode-back" href="../../../simstack/modules/simstack.util.html#simstack.util.db.Database.find_artifact_mappings">[docs]</a>
674
+ <span class="k">async</span> <span class="k">def</span><span class="w"> </span><span class="nf">find_artifact_mappings</span><span class="p">(</span>
675
+ <span class="bp">self</span><span class="p">,</span>
676
+ <span class="n">node_registry_path</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span>
677
+ <span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Any</span><span class="p">:</span>
678
+ <span class="kn">from</span><span class="w"> </span><span class="nn">simstack.core.artifacts</span><span class="w"> </span><span class="kn">import</span> <span class="n">find_artifact_mappings</span>
679
+
680
+ <span class="k">return</span> <span class="k">await</span> <span class="n">find_artifact_mappings</span><span class="p">(</span><span class="n">node_registry_path</span><span class="p">,</span> <span class="bp">self</span><span class="p">)</span></div>
681
+
682
+
683
+ <div class="viewcode-block" id="Database.find_all_artifacts">
684
+ <a class="viewcode-back" href="../../../simstack/modules/simstack.util.html#simstack.util.db.Database.find_all_artifacts">[docs]</a>
685
+ <span class="k">async</span> <span class="k">def</span><span class="w"> </span><span class="nf">find_all_artifacts</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">node_registry</span><span class="p">:</span> <span class="n">Any</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Any</span><span class="p">:</span>
686
+ <span class="kn">from</span><span class="w"> </span><span class="nn">simstack.core.artifacts</span><span class="w"> </span><span class="kn">import</span> <span class="n">find_all_artifacts</span>
687
+
688
+ <span class="k">return</span> <span class="k">await</span> <span class="n">find_all_artifacts</span><span class="p">(</span><span class="n">node_registry</span><span class="p">,</span> <span class="bp">self</span><span class="p">)</span></div>
689
+
690
+
691
+ <div class="viewcode-block" id="Database.ping">
692
+ <a class="viewcode-back" href="../../../simstack/modules/simstack.util.html#simstack.util.db.Database.ping">[docs]</a>
693
+ <span class="k">async</span> <span class="k">def</span><span class="w"> </span><span class="nf">ping</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Any</span><span class="p">:</span>
694
+ <span class="k">return</span> <span class="k">await</span> <span class="bp">self</span><span class="o">.</span><span class="n">client</span><span class="o">.</span><span class="n">admin</span><span class="o">.</span><span class="n">command</span><span class="p">(</span><span class="s2">&quot;ping&quot;</span><span class="p">)</span></div>
695
+
696
+
697
+ <div class="viewcode-block" id="Database.stats">
698
+ <a class="viewcode-back" href="../../../simstack/modules/simstack.util.html#simstack.util.db.Database.stats">[docs]</a>
699
+ <span class="k">async</span> <span class="k">def</span><span class="w"> </span><span class="nf">stats</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Any</span><span class="p">:</span>
700
+ <span class="k">return</span> <span class="k">await</span> <span class="bp">self</span><span class="o">.</span><span class="n">database</span><span class="o">.</span><span class="n">command</span><span class="p">(</span><span class="s2">&quot;dbStats&quot;</span><span class="p">)</span></div>
701
+
702
+
703
+ <span class="k">async</span> <span class="k">def</span><span class="w"> </span><span class="nf">_save_one</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">model</span><span class="p">:</span> <span class="n">Any</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">:</span> <span class="n">Any</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">:</span> <span class="n">Any</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Any</span><span class="p">:</span>
704
+ <span class="c1"># Avoid recursion if custom save is called</span>
705
+ <span class="k">if</span> <span class="nb">getattr</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="s2">&quot;_currently_saving&quot;</span><span class="p">,</span> <span class="kc">False</span><span class="p">):</span>
706
+ <span class="n">result</span> <span class="o">=</span> <span class="k">await</span> <span class="bp">self</span><span class="o">.</span><span class="n">_engine</span><span class="o">.</span><span class="n">save</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>
707
+ <span class="k">return</span> <span class="n">result</span> <span class="k">if</span> <span class="n">result</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="k">else</span> <span class="n">model</span>
708
+
709
+ <span class="n">custom_save_called</span> <span class="o">=</span> <span class="k">await</span> <span class="bp">self</span><span class="o">.</span><span class="n">_maybe_call_custom_save</span><span class="p">(</span><span class="n">model</span><span class="p">)</span>
710
+
711
+ <span class="c1"># Parts saves should always be called unless the main model has a custom save</span>
712
+ <span class="c1"># that handles everything. But in our current logic, _maybe_call_custom_save</span>
713
+ <span class="c1"># returning True means the custom save WAS called.</span>
714
+ <span class="k">if</span> <span class="ow">not</span> <span class="n">custom_save_called</span><span class="p">:</span>
715
+ <span class="k">await</span> <span class="bp">self</span><span class="o">.</span><span class="n">_call_parts_saves</span><span class="p">(</span><span class="n">model</span><span class="p">)</span>
716
+
717
+ <span class="k">if</span> <span class="n">custom_save_called</span><span class="p">:</span>
718
+ <span class="k">return</span> <span class="n">model</span>
719
+
720
+ <span class="n">result</span> <span class="o">=</span> <span class="k">await</span> <span class="bp">self</span><span class="o">.</span><span class="n">_engine</span><span class="o">.</span><span class="n">save</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">)</span>
721
+ <span class="k">return</span> <span class="n">result</span> <span class="k">if</span> <span class="n">result</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="k">else</span> <span class="n">model</span>
722
+
723
+ <span class="k">async</span> <span class="k">def</span><span class="w"> </span><span class="nf">_maybe_call_custom_save</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">target</span><span class="p">:</span> <span class="n">Any</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="nb">bool</span><span class="p">:</span>
724
+ <span class="n">save_attr</span> <span class="o">=</span> <span class="nb">getattr</span><span class="p">(</span><span class="n">target</span><span class="p">,</span> <span class="s2">&quot;save&quot;</span><span class="p">,</span> <span class="kc">None</span><span class="p">)</span>
725
+ <span class="k">if</span> <span class="ow">not</span> <span class="nb">callable</span><span class="p">(</span><span class="n">save_attr</span><span class="p">):</span>
726
+ <span class="k">return</span> <span class="kc">False</span>
727
+
728
+ <span class="c1"># Avoid recursion if custom save is called</span>
729
+ <span class="k">if</span> <span class="nb">getattr</span><span class="p">(</span><span class="n">target</span><span class="p">,</span> <span class="s2">&quot;_currently_saving&quot;</span><span class="p">,</span> <span class="kc">False</span><span class="p">):</span>
730
+ <span class="k">return</span> <span class="kc">False</span>
731
+
732
+ <span class="c1"># Set a flag to prevent recursion if custom save calls db.save(self)</span>
733
+ <span class="nb">object</span><span class="o">.</span><span class="fm">__setattr__</span><span class="p">(</span><span class="n">target</span><span class="p">,</span> <span class="s2">&quot;_currently_saving&quot;</span><span class="p">,</span> <span class="kc">True</span><span class="p">)</span>
734
+ <span class="k">try</span><span class="p">:</span>
735
+ <span class="k">await</span> <span class="n">save_attr</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span>
736
+ <span class="k">finally</span><span class="p">:</span>
737
+ <span class="nb">object</span><span class="o">.</span><span class="fm">__setattr__</span><span class="p">(</span><span class="n">target</span><span class="p">,</span> <span class="s2">&quot;_currently_saving&quot;</span><span class="p">,</span> <span class="kc">False</span><span class="p">)</span>
738
+ <span class="k">return</span> <span class="kc">True</span>
739
+
740
+ <span class="k">async</span> <span class="k">def</span><span class="w"> </span><span class="nf">_call_parts_saves</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">model</span><span class="p">:</span> <span class="n">Any</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="nb">bool</span><span class="p">:</span>
741
+ <span class="n">any_saved</span> <span class="o">=</span> <span class="kc">False</span>
742
+ <span class="n">seen_ids</span><span class="p">:</span> <span class="nb">set</span><span class="p">[</span><span class="nb">int</span><span class="p">]</span> <span class="o">=</span> <span class="nb">set</span><span class="p">()</span>
743
+
744
+ <span class="k">for</span> <span class="n">part</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">_iter_save_parts</span><span class="p">(</span><span class="n">model</span><span class="p">):</span>
745
+ <span class="n">part_id</span> <span class="o">=</span> <span class="nb">id</span><span class="p">(</span><span class="n">part</span><span class="p">)</span>
746
+ <span class="k">if</span> <span class="n">part_id</span> <span class="ow">in</span> <span class="n">seen_ids</span><span class="p">:</span>
747
+ <span class="k">continue</span>
748
+ <span class="n">seen_ids</span><span class="o">.</span><span class="n">add</span><span class="p">(</span><span class="n">part_id</span><span class="p">)</span>
749
+
750
+ <span class="k">if</span> <span class="k">await</span> <span class="bp">self</span><span class="o">.</span><span class="n">_maybe_call_custom_save</span><span class="p">(</span><span class="n">part</span><span class="p">):</span>
751
+ <span class="n">any_saved</span> <span class="o">=</span> <span class="kc">True</span>
752
+
753
+ <span class="k">return</span> <span class="n">any_saved</span>
754
+
755
+ <span class="nd">@staticmethod</span>
756
+ <span class="k">def</span><span class="w"> </span><span class="nf">_iter_save_parts</span><span class="p">(</span><span class="n">root</span><span class="p">:</span> <span class="n">Any</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Iterable</span><span class="p">[</span><span class="n">Any</span><span class="p">]:</span>
757
+ <span class="k">try</span><span class="p">:</span>
758
+ <span class="n">values</span> <span class="o">=</span> <span class="nb">list</span><span class="p">(</span><span class="nb">vars</span><span class="p">(</span><span class="n">root</span><span class="p">)</span><span class="o">.</span><span class="n">values</span><span class="p">())</span>
759
+ <span class="k">except</span> <span class="ne">TypeError</span><span class="p">:</span>
760
+ <span class="k">return</span> <span class="p">[]</span>
761
+
762
+ <span class="n">parts</span><span class="p">:</span> <span class="nb">list</span><span class="p">[</span><span class="n">Any</span><span class="p">]</span> <span class="o">=</span> <span class="p">[]</span>
763
+ <span class="k">for</span> <span class="n">value</span> <span class="ow">in</span> <span class="n">values</span><span class="p">:</span>
764
+ <span class="n">parts</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">value</span><span class="p">)</span>
765
+ <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">value</span><span class="p">,</span> <span class="p">(</span><span class="nb">list</span><span class="p">,</span> <span class="nb">tuple</span><span class="p">,</span> <span class="nb">set</span><span class="p">)):</span>
766
+ <span class="n">parts</span><span class="o">.</span><span class="n">extend</span><span class="p">(</span><span class="n">value</span><span class="p">)</span>
767
+ <span class="k">elif</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">value</span><span class="p">,</span> <span class="nb">dict</span><span class="p">):</span>
768
+ <span class="n">parts</span><span class="o">.</span><span class="n">extend</span><span class="p">(</span><span class="n">value</span><span class="o">.</span><span class="n">values</span><span class="p">())</span>
769
+ <span class="k">return</span> <span class="n">parts</span>
770
+
771
+ <div class="viewcode-block" id="Database.load_task">
772
+ <a class="viewcode-back" href="../../../simstack/modules/simstack.util.html#simstack.util.db.Database.load_task">[docs]</a>
773
+ <span class="k">async</span> <span class="k">def</span><span class="w"> </span><span class="nf">load_task</span><span class="p">(</span>
774
+ <span class="bp">self</span><span class="p">,</span> <span class="n">name</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span> <span class="n">arg_hash</span><span class="p">:</span> <span class="nb">str</span><span class="p">,</span> <span class="n">function_hash</span><span class="p">:</span> <span class="nb">str</span>
775
+ <span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Optional</span><span class="p">[</span><span class="s2">&quot;NodeRegistry&quot;</span><span class="p">]:</span>
776
+ <span class="w"> </span><span class="sd">&quot;&quot;&quot;</span>
777
+ <span class="sd"> Load a task based on name, arg_hash and function_hash</span>
778
+
779
+ <span class="sd"> Args:</span>
780
+ <span class="sd"> name: Node name</span>
781
+ <span class="sd"> arg_hash: Hash of the arguments</span>
782
+ <span class="sd"> function_hash: Hash of the function</span>
783
+
784
+ <span class="sd"> Returns:</span>
785
+ <span class="sd"> The found NodeRegistry instance or None</span>
786
+ <span class="sd"> &quot;&quot;&quot;</span>
787
+ <span class="n">result</span> <span class="o">=</span> <span class="k">await</span> <span class="bp">self</span><span class="o">.</span><span class="n">find_one</span><span class="p">(</span>
788
+ <span class="n">NodeRegistry</span><span class="p">,</span>
789
+ <span class="p">(</span><span class="n">NodeRegistry</span><span class="o">.</span><span class="n">name</span> <span class="o">==</span> <span class="n">name</span><span class="p">)</span>
790
+ <span class="o">&amp;</span> <span class="p">(</span><span class="n">NodeRegistry</span><span class="o">.</span><span class="n">arg_hash</span> <span class="o">==</span> <span class="n">arg_hash</span><span class="p">)</span>
791
+ <span class="o">&amp;</span> <span class="p">(</span><span class="n">NodeRegistry</span><span class="o">.</span><span class="n">function_hash</span> <span class="o">==</span> <span class="n">function_hash</span><span class="p">),</span>
792
+ <span class="p">)</span>
793
+ <span class="k">return</span> <span class="n">result</span></div>
794
+
795
+
796
+ <span class="c1"># legacy functions ... these are functions in the old database class which we do not want to migrate if possible</span>
797
+ <span class="c1">#</span>
798
+ <span class="c1"># load_waiting_tasks_for_resource DONE</span>
799
+ <span class="c1"># reset_database DONE</span>
800
+ <span class="c1"># the rest is hopefully not needed anymore</span>
801
+ <span class="c1"># list_collections</span>
802
+ <span class="c1"># upsert</span>
803
+ <span class="c1"># _save_references</span>
804
+ <span class="c1"># find_one_by_model_name</span>
805
+ <span class="c1"># find_all</span>
806
+ <span class="c1"># find_many</span>
807
+ <span class="c1"># delete_by_id</span>
808
+ <span class="c1"># drop_collection</span>
809
+ <span class="c1"># load_from_collection</span>
810
+ <span class="c1"># load_node_model_by_name</span>
811
+ <span class="c1"># load_task_by_id</span>
812
+ <span class="c1"># count</span>
813
+ <span class="c1"># aggregate</span>
814
+
815
+ <div class="viewcode-block" id="Database.load_waiting_tasks_for_resource">
816
+ <a class="viewcode-back" href="../../../simstack/modules/simstack.util.html#simstack.util.db.Database.load_waiting_tasks_for_resource">[docs]</a>
817
+ <span class="k">async</span> <span class="k">def</span><span class="w"> </span><span class="nf">load_waiting_tasks_for_resource</span><span class="p">(</span>
818
+ <span class="bp">self</span><span class="p">,</span> <span class="n">resource</span><span class="p">:</span> <span class="nb">str</span>
819
+ <span class="p">)</span> <span class="o">-&gt;</span> <span class="n">List</span><span class="p">[</span><span class="s2">&quot;NodeRegistry&quot;</span><span class="p">]:</span>
820
+ <span class="w"> </span><span class="sd">&quot;&quot;&quot;</span>
821
+ <span class="sd"> Load all waiting tasks for a specific resource</span>
822
+
823
+ <span class="sd"> Args:</span>
824
+ <span class="sd"> resource: The resource name</span>
825
+
826
+ <span class="sd"> Returns:</span>
827
+ <span class="sd"> List of matching NodeRegistry instances</span>
828
+ <span class="sd"> &quot;&quot;&quot;</span>
829
+ <span class="c1"># Try to use the engine directly if find is failing in tests</span>
830
+ <span class="k">try</span><span class="p">:</span>
831
+ <span class="n">submitted_tasks</span> <span class="o">=</span> <span class="k">await</span> <span class="bp">self</span><span class="o">.</span><span class="n">find</span><span class="p">(</span>
832
+ <span class="n">NodeRegistry</span><span class="p">,</span> <span class="n">NodeRegistry</span><span class="o">.</span><span class="n">status</span> <span class="o">==</span> <span class="n">TaskStatus</span><span class="o">.</span><span class="n">SUBMITTED</span>
833
+ <span class="p">)</span>
834
+ <span class="k">except</span> <span class="ne">Exception</span> <span class="k">as</span> <span class="n">e</span><span class="p">:</span>
835
+ <span class="n">logger</span><span class="o">.</span><span class="n">warning</span><span class="p">(</span>
836
+ <span class="sa">f</span><span class="s2">&quot;Error calling self.find in load_waiting_tasks_for_resource: </span><span class="si">{</span><span class="n">e</span><span class="si">}</span><span class="s2">. Falling back to engine.find.&quot;</span>
837
+ <span class="p">)</span>
838
+ <span class="n">submitted_tasks</span> <span class="o">=</span> <span class="k">await</span> <span class="bp">self</span><span class="o">.</span><span class="n">_engine</span><span class="o">.</span><span class="n">find</span><span class="p">(</span>
839
+ <span class="n">NodeRegistry</span><span class="p">,</span> <span class="n">NodeRegistry</span><span class="o">.</span><span class="n">status</span> <span class="o">==</span> <span class="n">TaskStatus</span><span class="o">.</span><span class="n">SUBMITTED</span>
840
+ <span class="p">)</span>
841
+
842
+ <span class="c1"># Then filter them in Python by checking the resource field</span>
843
+ <span class="n">matching_tasks</span> <span class="o">=</span> <span class="p">[]</span>
844
+ <span class="k">for</span> <span class="n">task</span> <span class="ow">in</span> <span class="n">submitted_tasks</span><span class="p">:</span>
845
+ <span class="c1"># Check if parameters has a resource attribute and if it matches our resource</span>
846
+ <span class="c1"># the local runner will also do the immidiate tasks</span>
847
+ <span class="c1"># TODO not the local runner the default resource should pick up self but there should be no seld</span>
848
+ <span class="k">if</span> <span class="p">(</span>
849
+ <span class="nb">hasattr</span><span class="p">(</span><span class="n">task</span><span class="o">.</span><span class="n">parameters</span><span class="p">,</span> <span class="s2">&quot;resource&quot;</span><span class="p">)</span>
850
+ <span class="ow">and</span> <span class="n">task</span><span class="o">.</span><span class="n">parameters</span><span class="o">.</span><span class="n">resource</span> <span class="o">==</span> <span class="n">resource</span>
851
+ <span class="p">):</span>
852
+ <span class="n">matching_tasks</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">task</span><span class="p">)</span>
853
+ <span class="k">if</span> <span class="p">(</span>
854
+ <span class="nb">hasattr</span><span class="p">(</span><span class="n">task</span><span class="o">.</span><span class="n">parameters</span><span class="p">,</span> <span class="s2">&quot;resource&quot;</span><span class="p">)</span>
855
+ <span class="ow">and</span> <span class="n">task</span><span class="o">.</span><span class="n">parameters</span><span class="o">.</span><span class="n">resource</span> <span class="o">==</span> <span class="s2">&quot;self&quot;</span>
856
+ <span class="p">):</span>
857
+ <span class="n">logger</span><span class="o">.</span><span class="n">error</span><span class="p">(</span><span class="s2">&quot;There should be tasks submitted to self&quot;</span><span class="p">)</span>
858
+ <span class="k">return</span> <span class="n">matching_tasks</span></div>
859
+
860
+
861
+ <div class="viewcode-block" id="Database.reset_database">
862
+ <a class="viewcode-back" href="../../../simstack/modules/simstack.util.html#simstack.util.db.Database.reset_database">[docs]</a>
863
+ <span class="k">async</span> <span class="k">def</span><span class="w"> </span><span class="nf">reset_database</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="kc">None</span><span class="p">:</span>
864
+ <span class="w"> </span><span class="sd">&quot;&quot;&quot;</span>
865
+ <span class="sd"> Reset the database by dropping all collections and recreating them</span>
866
+ <span class="sd"> &quot;&quot;&quot;</span>
867
+ <span class="n">db</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">client</span><span class="p">[</span><span class="bp">self</span><span class="o">.</span><span class="n">database_name</span><span class="p">]</span>
868
+ <span class="n">collections</span> <span class="o">=</span> <span class="k">await</span> <span class="n">db</span><span class="o">.</span><span class="n">list_collection_names</span><span class="p">()</span>
869
+
870
+ <span class="k">for</span> <span class="n">collection</span> <span class="ow">in</span> <span class="n">collections</span><span class="p">:</span>
871
+ <span class="k">await</span> <span class="n">db</span><span class="p">[</span><span class="n">collection</span><span class="p">]</span><span class="o">.</span><span class="n">drop</span><span class="p">()</span>
872
+
873
+ <span class="n">logger</span><span class="o">.</span><span class="n">info</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;Database </span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">database_name</span><span class="si">}</span><span class="s2"> has been reset&quot;</span><span class="p">)</span></div>
874
+
875
+
876
+ <div class="viewcode-block" id="Database.load_task_by_id">
877
+ <a class="viewcode-back" href="../../../simstack/modules/simstack.util.html#simstack.util.db.Database.load_task_by_id">[docs]</a>
878
+ <span class="k">async</span> <span class="k">def</span><span class="w"> </span><span class="nf">load_task_by_id</span><span class="p">(</span>
879
+ <span class="bp">self</span><span class="p">,</span> <span class="n">task_id</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">ObjectId</span><span class="p">]</span>
880
+ <span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Optional</span><span class="p">[</span><span class="n">NodeRegistry</span><span class="p">]:</span>
881
+ <span class="w"> </span><span class="sd">&quot;&quot;&quot;</span>
882
+ <span class="sd"> Load a task based on its ID</span>
883
+
884
+ <span class="sd"> Args:</span>
885
+ <span class="sd"> task_id: The task ID</span>
886
+
887
+ <span class="sd"> Returns:</span>
888
+ <span class="sd"> The found NodeRegistry instance or None</span>
889
+ <span class="sd"> &quot;&quot;&quot;</span>
890
+ <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">task_id</span><span class="p">,</span> <span class="nb">str</span><span class="p">):</span>
891
+ <span class="n">task_id</span> <span class="o">=</span> <span class="n">ObjectId</span><span class="p">(</span><span class="n">task_id</span><span class="p">)</span>
892
+
893
+ <span class="k">return</span> <span class="k">await</span> <span class="bp">self</span><span class="o">.</span><span class="n">find_one</span><span class="p">(</span><span class="n">NodeRegistry</span><span class="p">,</span> <span class="n">NodeRegistry</span><span class="o">.</span><span class="n">id</span> <span class="o">==</span> <span class="n">task_id</span><span class="p">)</span></div>
894
+ </div>
895
+
896
+
897
+
898
+ <span class="c1">#</span>
899
+ <span class="c1"># def set_database_core_context(database: Any):</span>
900
+ <span class="c1"># set_context = getattr(database, &quot;set_core_context&quot;, None)</span>
901
+ <span class="c1"># if callable(set_context):</span>
902
+ <span class="c1"># return set_context()</span>
903
+ <span class="c1"># return current_engine_context.set(database)</span>
904
+ <span class="c1">#</span>
905
+ <span class="c1">#</span>
906
+ <span class="c1"># def reset_database_core_context(database: Any, token) -&gt; None:</span>
907
+ <span class="c1"># reset_context = getattr(database, &quot;reset_core_context&quot;, None)</span>
908
+ <span class="c1"># if callable(reset_context):</span>
909
+ <span class="c1"># reset_context(token)</span>
910
+ <span class="c1"># return</span>
911
+ <span class="c1"># current_engine_context.reset(token)</span>
912
+ <span class="c1">#</span>
913
+
914
+
915
+ <span class="c1"># TODO engines</span>
916
+ <div class="viewcode-block" id="find_all_artifacts_for_database">
917
+ <a class="viewcode-back" href="../../../simstack/modules/simstack.util.html#simstack.util.db.find_all_artifacts_for_database">[docs]</a>
918
+ <span class="k">async</span> <span class="k">def</span><span class="w"> </span><span class="nf">find_all_artifacts_for_database</span><span class="p">(</span>
919
+ <span class="n">database</span><span class="p">:</span> <span class="n">Database</span><span class="p">,</span> <span class="n">node_registry</span><span class="p">:</span> <span class="n">Any</span>
920
+ <span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Any</span><span class="p">:</span>
921
+ <span class="n">find_all</span> <span class="o">=</span> <span class="nb">getattr</span><span class="p">(</span><span class="n">database</span><span class="p">,</span> <span class="s2">&quot;find_all_artifacts&quot;</span><span class="p">,</span> <span class="kc">None</span><span class="p">)</span>
922
+ <span class="k">if</span> <span class="nb">callable</span><span class="p">(</span><span class="n">find_all</span><span class="p">):</span>
923
+ <span class="k">return</span> <span class="k">await</span> <span class="n">database</span><span class="o">.</span><span class="n">find</span><span class="p">(</span><span class="n">node_registry</span><span class="p">)</span>
924
+
925
+ <span class="kn">from</span><span class="w"> </span><span class="nn">simstack.core.artifacts</span><span class="w"> </span><span class="kn">import</span> <span class="n">find_all_artifacts</span>
926
+
927
+ <span class="k">return</span> <span class="k">await</span> <span class="n">find_all_artifacts</span><span class="p">(</span><span class="n">node_registry</span><span class="p">,</span> <span class="n">database</span><span class="p">)</span></div>
928
+
929
+
930
+
931
+ <span class="c1">#</span>
932
+ <span class="c1"># class DatabaseOld(DatabaseInformation):</span>
933
+ <span class="c1"># &quot;&quot;&quot;</span>
934
+ <span class="c1"># Asynchronous MongoDB database access class using ODMantic ORM.</span>
935
+ <span class="c1"># Provides a cleaner interface for database operations.</span>
936
+ <span class="c1"># &quot;&quot;&quot;</span>
937
+ <span class="c1">#</span>
938
+ <span class="c1"># def __init__(self, db_type: DBType, db_name: str = &quot;simstack&quot;, connection_string: str = &quot;&quot;):</span>
939
+ <span class="c1"># super().__init__(db_name, connection_string, db_type)</span>
940
+ <span class="c1"># &quot;&quot;&quot;</span>
941
+ <span class="c1"># Initialize the MongoDB connection</span>
942
+ <span class="c1">#</span>
943
+ <span class="c1"># Args:</span>
944
+ <span class="c1"># db_type: Type of database configuration</span>
945
+ <span class="c1"># connection_string: MongoDB connection string (if not using default)</span>
946
+ <span class="c1"># db_name: Name of the MongoDB database</span>
947
+ <span class="c1"># &quot;&quot;&quot;</span>
948
+ <span class="c1">#</span>
949
+ <span class="c1"># if db_type == DBType.IN_MEMORY:</span>
950
+ <span class="c1"># # For tests, use in-memory MongoDB (mongomock)</span>
951
+ <span class="c1"># try:</span>
952
+ <span class="c1"># # import mongomock</span>
953
+ <span class="c1"># from mongomock_motor import AsyncMongoMockClient</span>
954
+ <span class="c1">#</span>
955
+ <span class="c1"># self.client = AsyncMongoMockClient()</span>
956
+ <span class="c1"># logger.info(&quot;Using in-memory MongoDB mock&quot;)</span>
957
+ <span class="c1"># except ImportError:</span>
958
+ <span class="c1"># logger.warning(</span>
959
+ <span class="c1"># &quot;mongomock not installed, falling back to localhost MongoDB&quot;</span>
960
+ <span class="c1"># )</span>
961
+ <span class="c1"># self.client = AsyncIOMotorClient(&quot;mongodb://localhost:27017&quot;)</span>
962
+ <span class="c1">#</span>
963
+ <span class="c1"># elif db_type == DBType.MONGODB:</span>
964
+ <span class="c1"># if not connection_string:</span>
965
+ <span class="c1"># connection_string = &quot;mongodb://localhost:27017&quot;</span>
966
+ <span class="c1"># self.client = AsyncIOMotorClient(connection_string)</span>
967
+ <span class="c1"># logger.info(&quot;Connected to MongoDB&quot;)</span>
968
+ <span class="c1">#</span>
969
+ <span class="c1"># else:</span>
970
+ <span class="c1"># raise ValueError(f&quot;Unsupported database type for MongoDB: {db_type}&quot;)</span>
971
+ <span class="c1">#</span>
972
+ <span class="c1"># # Create engine</span>
973
+ <span class="c1"># self.engine = AIOEngineProxy(client=self.client, database=db_name)</span>
974
+ <span class="c1"># # this will set the engine for all functions that are either called from the core package or the server</span>
975
+ <span class="c1"># # current_engine_context.set(self.engine)</span>
976
+ <span class="c1">#</span>
977
+ <span class="c1">#</span>
978
+ <span class="c1"># @classmethod</span>
979
+ <span class="c1"># def from_db_info(cls, db_info: DatabaseInformation):</span>
980
+ <span class="c1"># return cls(db_info.db_type, db_info.db_name, db_info.connection_string)</span>
981
+ <span class="c1">#</span>
982
+ <span class="c1"># async def list_collections(self):</span>
983
+ <span class="c1"># &quot;&quot;&quot;</span>
984
+ <span class="c1"># List all collections in the database</span>
985
+ <span class="c1"># &quot;&quot;&quot;</span>
986
+ <span class="c1"># db = self.client[self.db_name]</span>
987
+ <span class="c1"># return await db.list_collection_names()</span>
988
+ <span class="c1">#</span>
989
+ <span class="c1"># async def upsert(self, model: Model) -&gt; Model:</span>
990
+ <span class="c1"># &quot;&quot;&quot;</span>
991
+ <span class="c1"># Save or update a model instance including all references and list fields</span>
992
+ <span class="c1">#</span>
993
+ <span class="c1"># Args:</span>
994
+ <span class="c1"># model: The ODMantic model instance to save</span>
995
+ <span class="c1">#</span>
996
+ <span class="c1"># Returns:</span>
997
+ <span class="c1"># The saved model with updated fields</span>
998
+ <span class="c1"># &quot;&quot;&quot;</span>
999
+ <span class="c1"># # First, handle all references to ensure they exist in the database</span>
1000
+ <span class="c1"># if isinstance(Model, NodeRegistry):</span>
1001
+ <span class="c1"># if not hasattr(model, &quot;name&quot;):</span>
1002
+ <span class="c1"># logger.exception(</span>
1003
+ <span class="c1"># f&quot;Fatal Error A trying to save node_registry: {model.model_dump()} for task_id: {model.id}&quot;</span>
1004
+ <span class="c1"># )</span>
1005
+ <span class="c1"># return model</span>
1006
+ <span class="c1"># if model.name is None:</span>
1007
+ <span class="c1"># logger.exception(</span>
1008
+ <span class="c1"># f&quot;Fatal Error B trying to save node_registry: {model.model_dump()} for task_id: {model.id}&quot;</span>
1009
+ <span class="c1"># )</span>
1010
+ <span class="c1"># return model</span>
1011
+ <span class="c1">#</span>
1012
+ <span class="c1"># await self._save_references(model)</span>
1013
+ <span class="c1">#</span>
1014
+ <span class="c1"># # Then save the model itself</span>
1015
+ <span class="c1"># return await self.engine.save(model)</span>
1016
+ <span class="c1">#</span>
1017
+ <span class="c1"># async def _save_references(self, model: Model, visited=None):</span>
1018
+ <span class="c1"># &quot;&quot;&quot;</span>
1019
+ <span class="c1"># Recursively save all references within a model</span>
1020
+ <span class="c1">#</span>
1021
+ <span class="c1"># Args:</span>
1022
+ <span class="c1"># model: The model containing references to save</span>
1023
+ <span class="c1"># visited: Set of object IDs already processed to prevent infinite recursion</span>
1024
+ <span class="c1"># &quot;&quot;&quot;</span>
1025
+ <span class="c1"># if visited is None:</span>
1026
+ <span class="c1"># visited = set()</span>
1027
+ <span class="c1">#</span>
1028
+ <span class="c1"># # Skip if we&#39;ve already processed this object (prevents circular references)</span>
1029
+ <span class="c1"># model_id = id(model)</span>
1030
+ <span class="c1"># if model_id in visited:</span>
1031
+ <span class="c1"># return</span>
1032
+ <span class="c1"># visited.add(model_id)</span>
1033
+ <span class="c1">#</span>
1034
+ <span class="c1"># # Get all model fields, including those with default_factory</span>
1035
+ <span class="c1"># model_fields = getattr(model.__class__, &quot;model_fields&quot;, {})</span>
1036
+ <span class="c1">#</span>
1037
+ <span class="c1"># # Process each field in the model</span>
1038
+ <span class="c1"># for field_name, field_value in model.__dict__.items():</span>
1039
+ <span class="c1"># if field_value is None:</span>
1040
+ <span class="c1"># continue</span>
1041
+ <span class="c1">#</span>
1042
+ <span class="c1"># # Get field info if available</span>
1043
+ <span class="c1"># field_info = model_fields.get(field_name)</span>
1044
+ <span class="c1"># if not field_info:</span>
1045
+ <span class="c1"># continue</span>
1046
+ <span class="c1">#</span>
1047
+ <span class="c1"># # Handle different field types</span>
1048
+ <span class="c1">#</span>
1049
+ <span class="c1"># # Case 1: Direct Reference fields</span>
1050
+ <span class="c1"># if hasattr(field_info, &quot;annotation&quot;) and &quot;Reference&quot; in str(</span>
1051
+ <span class="c1"># field_info.annotation</span>
1052
+ <span class="c1"># ):</span>
1053
+ <span class="c1"># if field_value is not None:</span>
1054
+ <span class="c1"># logger.info(</span>
1055
+ <span class="c1"># f&quot;Saving reference field {field_name} of type {type(field_value).__name__}&quot;</span>
1056
+ <span class="c1"># )</span>
1057
+ <span class="c1"># await self._save_references(field_value, visited)</span>
1058
+ <span class="c1"># await self.engine.save(field_value)</span>
1059
+ <span class="c1">#</span>
1060
+ <span class="c1"># # Case 2: List fields that might contain models</span>
1061
+ <span class="c1"># elif isinstance(field_value, list):</span>
1062
+ <span class="c1"># for item in field_value:</span>
1063
+ <span class="c1"># if isinstance(item, Model):</span>
1064
+ <span class="c1"># logger.info(</span>
1065
+ <span class="c1"># f&quot;Saving list item of type {type(item).__name__} in field {field_name}&quot;</span>
1066
+ <span class="c1"># )</span>
1067
+ <span class="c1"># await self._save_references(item, visited)</span>
1068
+ <span class="c1"># await self.engine.save(item)</span>
1069
+ <span class="c1">#</span>
1070
+ <span class="c1"># # Case 3: Embedded models (like in FileInstance within FileStack)</span>
1071
+ <span class="c1"># elif isinstance(field_value, Model):</span>
1072
+ <span class="c1"># logger.info(</span>
1073
+ <span class="c1"># f&quot;Saving embedded model of type {type(field_value).__name__} in field {field_name}&quot;</span>
1074
+ <span class="c1"># )</span>
1075
+ <span class="c1"># await self._save_references(field_value, visited)</span>
1076
+ <span class="c1">#</span>
1077
+ <span class="c1"># async def save(self, model: Model) -&gt; Model:</span>
1078
+ <span class="c1"># return await self.upsert(model)</span>
1079
+ <span class="c1">#</span>
1080
+ <span class="c1"># async def find_one(self, model_class: Type[T], query=None, **kwargs) -&gt; Optional[T]:</span>
1081
+ <span class="c1"># &quot;&quot;&quot;</span>
1082
+ <span class="c1"># Find a single document matching the query</span>
1083
+ <span class="c1"># :param model_class:</span>
1084
+ <span class="c1"># :param query:</span>
1085
+ <span class="c1"># Returns:</span>
1086
+ <span class="c1"># The found model instance or None</span>
1087
+ <span class="c1">#</span>
1088
+ <span class="c1"># &quot;&quot;&quot;</span>
1089
+ <span class="c1"># return await self.engine.find_one(model_class, query, **kwargs)</span>
1090
+ <span class="c1">#</span>
1091
+ <span class="c1"># async def find_one_by_model_name(</span>
1092
+ <span class="c1"># self, model_mapping: str, item_id: str</span>
1093
+ <span class="c1"># ) -&gt; Optional[Any]:</span>
1094
+ <span class="c1"># &quot;&quot;&quot;</span>
1095
+ <span class="c1"># Find a single document matching the query by model name</span>
1096
+ <span class="c1">#</span>
1097
+ <span class="c1"># Args:</span>
1098
+ <span class="c1"># model_name: The name of the ODMantic model class as a string</span>
1099
+ <span class="c1">#</span>
1100
+ <span class="c1"># Returns:</span>
1101
+ <span class="c1"># The found model instance or None</span>
1102
+ <span class="c1">#</span>
1103
+ <span class="c1"># Raises:</span>
1104
+ <span class="c1"># ValueError: If model name is not found in the global namespace</span>
1105
+ <span class="c1"># &quot;&quot;&quot;</span>
1106
+ <span class="c1"># # Import common models that might be used</span>
1107
+ <span class="c1">#</span>
1108
+ <span class="c1"># # Find the model class based on its name</span>
1109
+ <span class="c1"># # model_class = import_class(model_mapping)</span>
1110
+ <span class="c1"># # model_elements = model_mapping.split(&quot;.&quot;)</span>
1111
+ <span class="c1"># # if len(model_elements) &gt; 1:</span>
1112
+ <span class="c1"># # model_name = model_mapping.split(&quot;.&quot;)[-1]</span>
1113
+ <span class="c1"># # else:</span>
1114
+ <span class="c1"># # model_name = model_mapping</span>
1115
+ <span class="c1"># # # Search through modules in current namespace</span>
1116
+ <span class="c1"># # for module_name, module in sys.modules.items():</span>
1117
+ <span class="c1"># # if hasattr(module, model_name):</span>
1118
+ <span class="c1"># # potential_class = getattr(module, model_name)</span>
1119
+ <span class="c1"># # # Check if it&#39;s likely a model class (has attributes like id, __collection__)</span>
1120
+ <span class="c1"># # if hasattr(potential_class, &quot;id&quot;) and hasattr(</span>
1121
+ <span class="c1"># # potential_class, &quot;__collection__&quot;</span>
1122
+ <span class="c1"># # ):</span>
1123
+ <span class="c1"># # model_class = potential_class</span>
1124
+ <span class="c1"># # break</span>
1125
+ <span class="c1"># # if model_class is None:</span>
1126
+ <span class="c1"># # logger.info(f&quot;Trying to import model: {model_mapping}&quot;)</span>
1127
+ <span class="c1">#</span>
1128
+ <span class="c1"># from simstack.util.importer import import_class</span>
1129
+ <span class="c1"># model_class = await import_class(model_mapping, self)</span>
1130
+ <span class="c1"># if model_class is None:</span>
1131
+ <span class="c1"># raise ValueError(</span>
1132
+ <span class="c1"># f&quot;DB: model class {model_mapping} not found in the available modules&quot;</span>
1133
+ <span class="c1"># )</span>
1134
+ <span class="c1">#</span>
1135
+ <span class="c1"># if isinstance(item_id, str):</span>
1136
+ <span class="c1"># item_id = ObjectId(item_id)</span>
1137
+ <span class="c1">#</span>
1138
+ <span class="c1"># instance = await self.engine.find_one(model_class, model_class.id == item_id)</span>
1139
+ <span class="c1"># if not instance:</span>
1140
+ <span class="c1"># logger.error(</span>
1141
+ <span class="c1"># f&quot;Instance of &#39;{model_class.__name__}&#39; with id &#39;{item_id}&#39; does not exist&quot;</span>
1142
+ <span class="c1"># )</span>
1143
+ <span class="c1"># raise ValueError(</span>
1144
+ <span class="c1"># f&quot;Instance of &#39;{model_class.__name__}&#39; with id &#39;{item_id}&#39; does not exist&quot;</span>
1145
+ <span class="c1"># )</span>
1146
+ <span class="c1"># return instance</span>
1147
+ <span class="c1">#</span>
1148
+ <span class="c1"># async def find_all(self, model_class: Type[T], **kwargs) -&gt; List[T]:</span>
1149
+ <span class="c1"># &quot;&quot;&quot;</span>
1150
+ <span class="c1"># Find all documents of a given model class</span>
1151
+ <span class="c1">#</span>
1152
+ <span class="c1"># Args:</span>
1153
+ <span class="c1"># model_class: The ODMantic model class</span>
1154
+ <span class="c1"># **kwargs: Query filters</span>
1155
+ <span class="c1"># &quot;&quot;&quot;</span>
1156
+ <span class="c1"># return await self.engine.find(model_class, **kwargs)</span>
1157
+ <span class="c1">#</span>
1158
+ <span class="c1"># async def find_many(self, model_class: Type[T], query, **kwargs) -&gt; List[T]:</span>
1159
+ <span class="c1"># &quot;&quot;&quot;</span>
1160
+ <span class="c1"># Find multiple documents matching the query</span>
1161
+ <span class="c1">#</span>
1162
+ <span class="c1"># Args:</span>
1163
+ <span class="c1"># model_class: The ODMantic model class</span>
1164
+ <span class="c1"># :param query:</span>
1165
+ <span class="c1"># **kwargs: Query filters</span>
1166
+ <span class="c1">#</span>
1167
+ <span class="c1"># Returns:</span>
1168
+ <span class="c1"># List of matching model instances</span>
1169
+ <span class="c1">#</span>
1170
+ <span class="c1"># &quot;&quot;&quot;</span>
1171
+ <span class="c1"># return await self.engine.find(model_class, query, **kwargs)</span>
1172
+ <span class="c1">#</span>
1173
+ <span class="c1"># async def delete(self, model: Model) -&gt; None:</span>
1174
+ <span class="c1"># &quot;&quot;&quot;</span>
1175
+ <span class="c1"># Delete a model instance</span>
1176
+ <span class="c1">#</span>
1177
+ <span class="c1"># Args:</span>
1178
+ <span class="c1"># model: The model instance to delete</span>
1179
+ <span class="c1"># &quot;&quot;&quot;</span>
1180
+ <span class="c1"># await self.engine.delete(model)</span>
1181
+ <span class="c1">#</span>
1182
+ <span class="c1"># async def delete_by_id(</span>
1183
+ <span class="c1"># self, model_class: Type[T], id: Union[str, ObjectId]</span>
1184
+ <span class="c1"># ) -&gt; None:</span>
1185
+ <span class="c1"># &quot;&quot;&quot;</span>
1186
+ <span class="c1"># Delete a document by its ID</span>
1187
+ <span class="c1">#</span>
1188
+ <span class="c1"># Args:</span>
1189
+ <span class="c1"># model_class: The ODMantic model class</span>
1190
+ <span class="c1"># id: The document ID (either string or ObjectId)</span>
1191
+ <span class="c1"># &quot;&quot;&quot;</span>
1192
+ <span class="c1"># # Convert string ID to ObjectId if needed</span>
1193
+ <span class="c1"># if isinstance(id, str):</span>
1194
+ <span class="c1"># id = ObjectId(id)</span>
1195
+ <span class="c1">#</span>
1196
+ <span class="c1"># instance = await self.engine.find_one(model_class, model_class.id == id)</span>
1197
+ <span class="c1"># if instance:</span>
1198
+ <span class="c1"># await self.engine.delete(instance)</span>
1199
+ <span class="c1"># else:</span>
1200
+ <span class="c1"># logger.error(f&quot;No data found in &#39;{model_class.__name__}&#39; with id &#39;{id}&#39;&quot;)</span>
1201
+ <span class="c1">#</span>
1202
+ <span class="c1"># async def drop_collection(self, model_class: Type[T]) -&gt; None:</span>
1203
+ <span class="c1"># &quot;&quot;&quot;</span>
1204
+ <span class="c1"># Drop the collection for the given model class</span>
1205
+ <span class="c1">#</span>
1206
+ <span class="c1"># Args:</span>
1207
+ <span class="c1"># model_class: The ODMantic model class</span>
1208
+ <span class="c1"># &quot;&quot;&quot;</span>
1209
+ <span class="c1"># collection = self.engine.get_collection(model_class)</span>
1210
+ <span class="c1"># if collection is None:</span>
1211
+ <span class="c1"># logger.error(f&quot;Could not drop collection {model_class.__name__}&quot;)</span>
1212
+ <span class="c1"># await collection.drop()</span>
1213
+ <span class="c1">#</span>
1214
+ <span class="c1"># # collection_name = model_class.__collection__</span>
1215
+ <span class="c1"># # db = self.client[self.db_name]</span>
1216
+ <span class="c1"># # await db[collection_name].drop()</span>
1217
+ <span class="c1">#</span>
1218
+ <span class="c1"># async def load_from_collection(</span>
1219
+ <span class="c1"># self, model_class: Type[T], id: Union[str, ObjectId]</span>
1220
+ <span class="c1"># ) -&gt; Optional[T]:</span>
1221
+ <span class="c1"># &quot;&quot;&quot;</span>
1222
+ <span class="c1"># Load a document by its ID</span>
1223
+ <span class="c1">#</span>
1224
+ <span class="c1"># Args:</span>
1225
+ <span class="c1"># model_class: The ODMantic model class</span>
1226
+ <span class="c1"># id: The document ID (either string or ObjectId)</span>
1227
+ <span class="c1">#</span>
1228
+ <span class="c1"># Returns:</span>
1229
+ <span class="c1"># The found model instance or None</span>
1230
+ <span class="c1">#</span>
1231
+ <span class="c1"># Raises:</span>
1232
+ <span class="c1"># ValueError: If document is not found</span>
1233
+ <span class="c1"># &quot;&quot;&quot;</span>
1234
+ <span class="c1"># # Convert string ID to ObjectId if needed</span>
1235
+ <span class="c1"># if isinstance(id, str):</span>
1236
+ <span class="c1"># id = ObjectId(id)</span>
1237
+ <span class="c1">#</span>
1238
+ <span class="c1"># instance = await self.engine.find_one(model_class, model_class.id == id)</span>
1239
+ <span class="c1"># if not instance:</span>
1240
+ <span class="c1"># logger.error(f&quot;No data found in &#39;{model_class.__name__}&#39; with id &#39;{id}&#39;&quot;)</span>
1241
+ <span class="c1"># raise ValueError(</span>
1242
+ <span class="c1"># f&quot;No data found in &#39;{model_class.__name__}&#39; with id &#39;{id}&#39;&quot;</span>
1243
+ <span class="c1"># )</span>
1244
+ <span class="c1"># return instance</span>
1245
+ <span class="c1">#</span>
1246
+ <span class="c1"># async def load_node_model_by_name(</span>
1247
+ <span class="c1"># self, node_model_name: str</span>
1248
+ <span class="c1"># ) -&gt; Optional[&quot;NodeModel&quot;]:</span>
1249
+ <span class="c1"># &quot;&quot;&quot;</span>
1250
+ <span class="c1"># Load a node based on its name</span>
1251
+ <span class="c1">#</span>
1252
+ <span class="c1"># Args:</span>
1253
+ <span class="c1"># node_model_name: The node name</span>
1254
+ <span class="c1">#</span>
1255
+ <span class="c1"># Returns:</span>
1256
+ <span class="c1"># The found NodeRegistry instance or None</span>
1257
+ <span class="c1"># &quot;&quot;&quot;</span>
1258
+ <span class="c1"># from simstack.models import NodeModel</span>
1259
+ <span class="c1"># return await self.engine.find_one(NodeModel, NodeModel.name == node_model_name)</span>
1260
+ <span class="c1">#</span>
1261
+ <span class="c1"># async def load_task_by_id(</span>
1262
+ <span class="c1"># self, task_id: Union[str, ObjectId]</span>
1263
+ <span class="c1"># ) -&gt; Optional[NodeRegistry]:</span>
1264
+ <span class="c1"># &quot;&quot;&quot;</span>
1265
+ <span class="c1"># Load a task based on its ID</span>
1266
+ <span class="c1">#</span>
1267
+ <span class="c1"># Args:</span>
1268
+ <span class="c1"># task_id: The task ID</span>
1269
+ <span class="c1">#</span>
1270
+ <span class="c1"># Returns:</span>
1271
+ <span class="c1"># The found NodeRegistry instance or None</span>
1272
+ <span class="c1"># &quot;&quot;&quot;</span>
1273
+ <span class="c1"># if isinstance(task_id, str):</span>
1274
+ <span class="c1"># task_id = ObjectId(task_id)</span>
1275
+ <span class="c1">#</span>
1276
+ <span class="c1"># return await self.engine.find_one(NodeRegistry, NodeRegistry.id == task_id)</span>
1277
+ <span class="c1">#</span>
1278
+ <span class="c1"># async def load_waiting_tasks_for_resource(</span>
1279
+ <span class="c1"># self, resource: str</span>
1280
+ <span class="c1"># ) -&gt; List[NodeRegistry]:</span>
1281
+ <span class="c1"># &quot;&quot;&quot;</span>
1282
+ <span class="c1"># Load all waiting tasks for a specific resource</span>
1283
+ <span class="c1">#</span>
1284
+ <span class="c1"># Args:</span>
1285
+ <span class="c1"># resource: The resource name</span>
1286
+ <span class="c1">#</span>
1287
+ <span class="c1"># Returns:</span>
1288
+ <span class="c1"># List of matching NodeRegistry instances</span>
1289
+ <span class="c1"># &quot;&quot;&quot;</span>
1290
+ <span class="c1"># submitted_tasks = await self.engine.find(</span>
1291
+ <span class="c1"># NodeRegistry, NodeRegistry.status == TaskStatus.SUBMITTED</span>
1292
+ <span class="c1"># )</span>
1293
+ <span class="c1"># # Then filter them in Python by checking the resource field</span>
1294
+ <span class="c1"># matching_tasks = []</span>
1295
+ <span class="c1"># for task in submitted_tasks:</span>
1296
+ <span class="c1"># # Check if parameters has a resource attribute and if it matches our resource</span>
1297
+ <span class="c1"># # the local runner will also do the immidiate tasks</span>
1298
+ <span class="c1"># if hasattr(task.parameters, &quot;resource&quot;) and (</span>
1299
+ <span class="c1"># task.parameters.resource == resource</span>
1300
+ <span class="c1"># or (resource == &quot;local&quot; and task.parameters.resource == &quot;self&quot;)</span>
1301
+ <span class="c1"># ):</span>
1302
+ <span class="c1"># if resource == &quot;local&quot; and task.parameters.resource == &quot;self&quot;:</span>
1303
+ <span class="c1"># logger.info(f&quot;local runner taking job for &#39;self&#39; with {task.id}&quot;)</span>
1304
+ <span class="c1"># matching_tasks.append(task)</span>
1305
+ <span class="c1"># return matching_tasks</span>
1306
+ <span class="c1">#</span>
1307
+ <span class="c1"># async def reset_database(self) -&gt; None:</span>
1308
+ <span class="c1"># &quot;&quot;&quot;</span>
1309
+ <span class="c1"># Reset the database by dropping all collections and recreating them</span>
1310
+ <span class="c1"># &quot;&quot;&quot;</span>
1311
+ <span class="c1"># db = self.client[self.db_name]</span>
1312
+ <span class="c1"># collections = await db.list_collection_names()</span>
1313
+ <span class="c1">#</span>
1314
+ <span class="c1"># for collection in collections:</span>
1315
+ <span class="c1"># await db[collection].drop()</span>
1316
+ <span class="c1">#</span>
1317
+ <span class="c1"># logger.info(f&quot;Database {self.db_name} has been reset&quot;)</span>
1318
+ <span class="c1">#</span>
1319
+ <span class="c1"># async def count(self, model_class: Type[T], **kwargs) -&gt; int:</span>
1320
+ <span class="c1"># &quot;&quot;&quot;</span>
1321
+ <span class="c1"># Count documents matching the query</span>
1322
+ <span class="c1">#</span>
1323
+ <span class="c1"># Args:</span>
1324
+ <span class="c1"># model_class: The ODMantic model class</span>
1325
+ <span class="c1"># **kwargs: Query filters</span>
1326
+ <span class="c1">#</span>
1327
+ <span class="c1"># Returns:</span>
1328
+ <span class="c1"># Number of matching documents</span>
1329
+ <span class="c1"># &quot;&quot;&quot;</span>
1330
+ <span class="c1"># return await self.engine.count(model_class, **kwargs)</span>
1331
+ <span class="c1">#</span>
1332
+ <span class="c1"># async def aggregate(</span>
1333
+ <span class="c1"># self, model_class: Type[Model], pipeline: List[Dict]</span>
1334
+ <span class="c1"># ) -&gt; List[Dict]:</span>
1335
+ <span class="c1"># &quot;&quot;&quot;</span>
1336
+ <span class="c1"># Perform an aggregation operation</span>
1337
+ <span class="c1">#</span>
1338
+ <span class="c1"># Args:</span>
1339
+ <span class="c1"># model_class: The ODMantic model class</span>
1340
+ <span class="c1"># pipeline: MongoDB aggregation pipeline</span>
1341
+ <span class="c1">#</span>
1342
+ <span class="c1"># Returns:</span>
1343
+ <span class="c1"># List of aggregation results</span>
1344
+ <span class="c1"># &quot;&quot;&quot;</span>
1345
+ <span class="c1"># collection = self.engine.get_collection(model_class)</span>
1346
+ <span class="c1"># cursor = collection.aggregate(pipeline)</span>
1347
+ <span class="c1"># return await cursor.to_list(length=None)</span>
1348
+ <span class="c1">#</span>
1349
+ <span class="c1"># async def close(self) -&gt; None:</span>
1350
+ <span class="c1"># &quot;&quot;&quot;</span>
1351
+ <span class="c1"># Close database connections</span>
1352
+ <span class="c1"># &quot;&quot;&quot;</span>
1353
+ <span class="c1"># self.client.close()</span>
1354
+ </pre></div>
1355
+ </article>
1356
+ </div>
1357
+ <footer>
1358
+
1359
+ <div class="related-pages">
1360
+
1361
+
1362
+ </div>
1363
+ <div class="bottom-of-page">
1364
+ <div class="left-details">
1365
+ <div class="copyright">
1366
+ Copyright &#169; 2026, Wolfgang Wenzel
1367
+ </div>
1368
+ Made with <a href="https://www.sphinx-doc.org/">Sphinx</a> and <a class="muted-link" href="https://pradyunsg.me">@pradyunsg</a>'s
1369
+
1370
+ <a href="https://github.com/pradyunsg/furo">Furo</a>
1371
+
1372
+ </div>
1373
+ <div class="right-details">
1374
+
1375
+ </div>
1376
+ </div>
1377
+
1378
+ </footer>
1379
+ </div>
1380
+ <aside class="toc-drawer no-toc">
1381
+
1382
+
1383
+
1384
+ </aside>
1385
+ </div>
1386
+ </div><script src="../../../_static/documentation_options.js?v=5929fcd5"></script>
1387
+ <script src="../../../_static/doctools.js?v=fd6eb6e6"></script>
1388
+ <script src="../../../_static/sphinx_highlight.js?v=6ffebe34"></script>
1389
+ <script src="../../../_static/scripts/furo.js?v=46bd48cc"></script>
1390
+ <script data-docs-built-at="2026-07-17 10:50:10 UTC" src="../../../_static/docs-build-badge.js?v=57cb9fc6"></script>
1391
+ </body>
1392
+ </html>