UncountablePythonSDK 0.0.69__py3-none-any.whl → 0.0.70__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 (38) hide show
  1. {UncountablePythonSDK-0.0.69.dist-info → UncountablePythonSDK-0.0.70.dist-info}/METADATA +3 -1
  2. {UncountablePythonSDK-0.0.69.dist-info → UncountablePythonSDK-0.0.70.dist-info}/RECORD +38 -17
  3. examples/integration-server/jobs/materials_auto/example_cron.py +2 -2
  4. uncountable/core/environment.py +5 -1
  5. uncountable/integration/cli.py +1 -0
  6. uncountable/integration/cron.py +12 -28
  7. uncountable/integration/db/connect.py +12 -2
  8. uncountable/integration/db/session.py +25 -0
  9. uncountable/integration/entrypoint.py +6 -6
  10. uncountable/integration/executors/generic_upload_executor.py +5 -1
  11. uncountable/integration/job.py +44 -17
  12. uncountable/integration/queue_runner/__init__.py +0 -0
  13. uncountable/integration/queue_runner/command_server/__init__.py +24 -0
  14. uncountable/integration/queue_runner/command_server/command_client.py +68 -0
  15. uncountable/integration/queue_runner/command_server/command_server.py +64 -0
  16. uncountable/integration/queue_runner/command_server/protocol/__init__.py +0 -0
  17. uncountable/integration/queue_runner/command_server/protocol/command_server.proto +22 -0
  18. uncountable/integration/queue_runner/command_server/protocol/command_server_pb2.py +40 -0
  19. uncountable/integration/queue_runner/command_server/protocol/command_server_pb2.pyi +38 -0
  20. uncountable/integration/queue_runner/command_server/protocol/command_server_pb2_grpc.py +129 -0
  21. uncountable/integration/queue_runner/command_server/types.py +52 -0
  22. uncountable/integration/queue_runner/datastore/__init__.py +3 -0
  23. uncountable/integration/queue_runner/datastore/datastore_sqlite.py +93 -0
  24. uncountable/integration/queue_runner/datastore/interface.py +19 -0
  25. uncountable/integration/queue_runner/datastore/model.py +17 -0
  26. uncountable/integration/queue_runner/job_scheduler.py +119 -0
  27. uncountable/integration/queue_runner/queue_runner.py +26 -0
  28. uncountable/integration/queue_runner/types.py +7 -0
  29. uncountable/integration/queue_runner/worker.py +109 -0
  30. uncountable/integration/scan_profiles.py +2 -0
  31. uncountable/integration/scheduler.py +40 -3
  32. uncountable/integration/webhook_server/entrypoint.py +27 -31
  33. uncountable/types/__init__.py +2 -0
  34. uncountable/types/api/recipes/get_recipes_data.py +1 -0
  35. uncountable/types/queued_job.py +16 -0
  36. uncountable/types/queued_job_t.py +107 -0
  37. {UncountablePythonSDK-0.0.69.dist-info → UncountablePythonSDK-0.0.70.dist-info}/WHEEL +0 -0
  38. {UncountablePythonSDK-0.0.69.dist-info → UncountablePythonSDK-0.0.70.dist-info}/top_level.txt +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: UncountablePythonSDK
3
- Version: 0.0.69
3
+ Version: 0.0.70
4
4
  Summary: Uncountable SDK
5
5
  Project-URL: Homepage, https://github.com/uncountableinc/uncountable-python-sdk
6
6
  Project-URL: Repository, https://github.com/uncountableinc/uncountable-python-sdk.git
@@ -35,6 +35,8 @@ Requires-Dist: paramiko ==3.*
35
35
  Requires-Dist: boto3 ==1.*
36
36
  Requires-Dist: flask ==3.*
37
37
  Requires-Dist: simplejson ==3.*
38
+ Requires-Dist: grpcio ==1.*
39
+ Requires-Dist: protobuf ==4.*
38
40
  Provides-Extra: test
39
41
  Requires-Dist: mypy ==1.* ; extra == 'test'
40
42
  Requires-Dist: ruff ==0.* ; extra == 'test'
@@ -22,7 +22,7 @@ examples/set_recipe_metadata_file.py,sha256=oPBhMo9T17zj4YpJRkmVH9lknYcT8livph0K
22
22
  examples/set_recipe_output_file_sdk.py,sha256=Lz1amqppnWTX83z-C090wCJ4hcKmCD3kb-4v0uBRi0Y,782
23
23
  examples/upload_files.py,sha256=tUfKFqiqwnw08OL5Y8_e4j5pSRhp94cFex8XTuVa_ig,487
24
24
  examples/integration-server/pyproject.toml,sha256=yVYpx5VG3bKQCYjmCSKN_MRjxoOoF2xF2KzkCcvDHh4,9134
25
- examples/integration-server/jobs/materials_auto/example_cron.py,sha256=ITXXIDLrFhwS_YMrcuEBTciGoC8ozjB9SXL-qTCZT2k,707
25
+ examples/integration-server/jobs/materials_auto/example_cron.py,sha256=7VVQ-UJsq3DbGpN3XPnorRVZYo-vCwbfSU3VVDluIzA,699
26
26
  examples/integration-server/jobs/materials_auto/profile.yaml,sha256=WZTNlgc8uH412oRmWHWWuenYo_6LnrIL11V0HMkGsB0,352
27
27
  pkgs/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
28
28
  pkgs/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -81,29 +81,48 @@ uncountable/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
81
81
  uncountable/core/__init__.py,sha256=RFv0kO6rKFf1PtBPu83hCGmxqkJamRtsgQ9_-ztw7tA,341
82
82
  uncountable/core/async_batch.py,sha256=Gur0VOS0AH2ugwvk65hwoX-iqwQAAyJaejY_LyAZZPo,1210
83
83
  uncountable/core/client.py,sha256=KUJN3XcbawMg_GuJ5DvmDsmhRzsnafCjyq27vOD4jC4,10640
84
- uncountable/core/environment.py,sha256=4rLWeUIDYB6w1iqrAiSEaD_OqMoo_ed3xcjz5jM4HGU,572
84
+ uncountable/core/environment.py,sha256=n46mWvG5uBOy6H_aetu_iuaYmO_d23uSbQBRgb6xStw,677
85
85
  uncountable/core/file_upload.py,sha256=qR7BBBWVxFNrb1_WICreo3dkZygE9lcE1fmZCQrDZU0,3469
86
86
  uncountable/core/types.py,sha256=s2CjqYJpsmbC7xMwxxT7kJ_V9bwokrjjWVVjpMcQpKI,333
87
87
  uncountable/integration/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
88
- uncountable/integration/cli.py,sha256=R_z84CejxV8QW8pXQWX1RxVgMiA3Sgt_X-zMFVNtXLw,2856
88
+ uncountable/integration/cli.py,sha256=wv9XHXujO0Zop-R4jmjNQttPOjEDQlXPfvmfZ5HW9QE,2898
89
89
  uncountable/integration/construct_client.py,sha256=I2XTamht13vs-JYkV4PpNS_Pc4FJm-KVYqNNvxI4qNk,1916
90
- uncountable/integration/cron.py,sha256=Gwu1Mq4lZZvtvx5tlnWJHnXJp948Xk9CXdS1ULVCVws,1629
91
- uncountable/integration/entrypoint.py,sha256=wgOXhTzErttRjOzV4rS4psZW5qUKIa5ez89QndQl61k,785
92
- uncountable/integration/job.py,sha256=UrsZHWXKE2wD5M3lFKJCXvDdWj7QMDtAREKv6RBnT3Q,1548
93
- uncountable/integration/scan_profiles.py,sha256=fHRS9Su6JlYWKBGZkA3KgUSqbT_PrB91Xe2rkdvv1Ew,1396
94
- uncountable/integration/scheduler.py,sha256=1YwY2f26Wr1kIlJfUthuRIXsUJZfynf9X3Zu76NK3ts,3262
90
+ uncountable/integration/cron.py,sha256=6eH-kIs3sdYPCyb62_L2M7U_uQTdMTdwY5hreEJb0hw,887
91
+ uncountable/integration/entrypoint.py,sha256=EjJdTmED80MrVcNoQ5BhFP3bAYHdU6CzRGiPSZF9y8w,802
92
+ uncountable/integration/job.py,sha256=af197JUceIKzpIN5C2z8zeZOPhIQ16ipyC6qVt1WXv0,2386
93
+ uncountable/integration/scan_profiles.py,sha256=ro4u5_1ClUeyW4vcWJ9epf2UQaJrjtc7S1R5bZpuHTM,1430
94
+ uncountable/integration/scheduler.py,sha256=sVe7V5zlUbSzSMyC30rAlWR30w0jurhrYYCgK8-lVRo,4546
95
95
  uncountable/integration/server.py,sha256=-ALtXNlBDnDbeSJcqAHHJzl46mGuq43aIM9RLzG-ZWA,4701
96
96
  uncountable/integration/telemetry.py,sha256=MwQLmgCoxpmA_UTp3e2ZB37wcDzKM0qzm5MrmaiJWhU,7142
97
97
  uncountable/integration/db/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
98
- uncountable/integration/db/connect.py,sha256=YtQHJ1DBGPhxKFRCfiXqohOYUceKSxMVOJ88aPI48Ug,181
98
+ uncountable/integration/db/connect.py,sha256=mE3bdV0huclH2iT_dXCQdRL4LkjIuf_myAR64RTWXEs,498
99
+ uncountable/integration/db/session.py,sha256=96cGQXpe6IugBTdSsjdP0S5yhJ6toSmbVB6qhc3FJzE,693
99
100
  uncountable/integration/executors/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
100
101
  uncountable/integration/executors/executors.py,sha256=CbwatKkHrLhnqYr_nsBjr0KYeOn8KqijxrZPHbxChBY,1961
101
- uncountable/integration/executors/generic_upload_executor.py,sha256=PpxZmGUac4PtwlEJi7-S7scSa3y6HrdBPs_p2QDyqZs,10286
102
+ uncountable/integration/executors/generic_upload_executor.py,sha256=NlW5WcYePPA7_fwp5uW_2afAiQLKK7rCkKF06wQ948E,10375
102
103
  uncountable/integration/executors/script_executor.py,sha256=OmSBOtU48G3mqza9c2lCm84pGGyaDk-ZBJCx3RsdJXc,846
104
+ uncountable/integration/queue_runner/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
105
+ uncountable/integration/queue_runner/job_scheduler.py,sha256=4AdI_89mPNAgO-WSJvqRaN3HbKlzlVRh9aIkwfQvL7o,4679
106
+ uncountable/integration/queue_runner/queue_runner.py,sha256=0BmYu5zHdothTevGsB-nXg6MBd1UD-WkP3h1WCKMdQg,710
107
+ uncountable/integration/queue_runner/types.py,sha256=8qTq29BTSa5rmW6CBlBntP0pNIiDcwu1wHa78pjroS0,219
108
+ uncountable/integration/queue_runner/worker.py,sha256=GNz8PmPUBMLA-R2VWkC_86WOaeHjHFI2mvBlNRcI3wE,4238
109
+ uncountable/integration/queue_runner/command_server/__init__.py,sha256=gQPVILGpWzCr2i5GJyoqna7AOSFvtn4tav69gB78mTQ,571
110
+ uncountable/integration/queue_runner/command_server/command_client.py,sha256=DJb0TUVFkiiLBEQzHSN94sTRnuEbutNEgdN39XmnOXI,2046
111
+ uncountable/integration/queue_runner/command_server/command_server.py,sha256=yyXryhiEC2eGS0yFElLGsVzSKwOuYvj-zp22jQorkv0,2138
112
+ uncountable/integration/queue_runner/command_server/types.py,sha256=A9FpGplHxoSkmi8dn4oGFN7bkstQjY1DoTXypbicHpk,991
113
+ uncountable/integration/queue_runner/command_server/protocol/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
114
+ uncountable/integration/queue_runner/command_server/protocol/command_server.proto,sha256=pf7FAT2eGuao0VYCFrgTAsM-tiPi1Bhz19XN5So1WFk,439
115
+ uncountable/integration/queue_runner/command_server/protocol/command_server_pb2.py,sha256=-lBTc5Tz48agqNSeOSpBE69e2kRmWF59sUaowCl8p7U,2207
116
+ uncountable/integration/queue_runner/command_server/protocol/command_server_pb2.pyi,sha256=9viBn6PHvtfMSRwam57ke5O2D_k8LapWYVfBRjknIYg,1281
117
+ uncountable/integration/queue_runner/command_server/protocol/command_server_pb2_grpc.py,sha256=ZVHkuLDjEbXMCxBsw1UrRhT3EEF8CDDqEvmE3Kbp1H4,5359
118
+ uncountable/integration/queue_runner/datastore/__init__.py,sha256=6BefApqN8D2zlVOH14QAeVzwQ8j5NIb41-njT02Za0k,88
119
+ uncountable/integration/queue_runner/datastore/datastore_sqlite.py,sha256=UZQABTrM3HxVaojCnI1pTU3v1GsbE3G-OCgXP5ekcxI,3801
120
+ uncountable/integration/queue_runner/datastore/interface.py,sha256=j4D-zVvLq-48VTVwHVei82UVUJ_P3cxiseyiTl0MoNw,534
121
+ uncountable/integration/queue_runner/datastore/model.py,sha256=8-RI5A2yPZVGBLWINVmMd6VOl_oHtqGtnaNXcapAChw,577
103
122
  uncountable/integration/secret_retrieval/__init__.py,sha256=3QXVj35w8rRMxVvmmsViFYDi3lcb3g70incfalOEm6o,87
104
123
  uncountable/integration/secret_retrieval/retrieve_secret.py,sha256=eoPWbkUtCn_63A4TFlK_nvEDvfm4u2fiOoglmAkBG3U,3004
105
- uncountable/integration/webhook_server/entrypoint.py,sha256=01yx5I_4e380oGqaceCcsGhV1JPHcY5lXIFJoghFmwc,6179
106
- uncountable/types/__init__.py,sha256=MDzXxVKMd9OY6WGXqTRRO_XxgHeHy9OOig9L__BlF9s,8498
124
+ uncountable/integration/webhook_server/entrypoint.py,sha256=hgbEtdVo3QU3odlHSygxmrsxR9g7rgU0yKt-o9nVAHE,5686
125
+ uncountable/types/__init__.py,sha256=KSsSEBnGhn88NPex5q9MZtY4kj8PRqVTbaKtko4hxUU,8561
107
126
  uncountable/types/async_batch.py,sha256=_OhT25_dEVts_z_n1kqfJH3xlZg3btLqR6TNkfFLlXE,609
108
127
  uncountable/types/async_batch_processor.py,sha256=obVzN-PcYLV2pHScszfCGjSq6-Xc34WM1ysx6Fv6tZk,11293
109
128
  uncountable/types/async_batch_t.py,sha256=ipSGz93O1KB-WE2dvlvflTKS51rJrf3bJkUojyxos7I,2193
@@ -148,6 +167,8 @@ uncountable/types/phases.py,sha256=YCsU77DdjRJJWdLTwLuOZNG4e9ML82NIBI1xTWr3ggA,2
148
167
  uncountable/types/phases_t.py,sha256=21_4-O5MlO1AzsZijSlSUeEHknY0T9KCPilFMG8GMEo,544
149
168
  uncountable/types/post_base.py,sha256=GES5_IhXFAjpzI6ChbVP4zTkX6f3tPUIHST1sxsRN6Q,279
150
169
  uncountable/types/post_base_t.py,sha256=2et3TDnFChZcx0RWU-18RJHw9Yiyj2TJz-2tByHXwaI,770
170
+ uncountable/types/queued_job.py,sha256=67exX5vfpewKzvFKxuxTLsSJTDKlE3JqKfWnnYx2Jos,842
171
+ uncountable/types/queued_job_t.py,sha256=yd3FhtcbfMxPjBq_ykLY2ASx_EJvQ4kW3MKIL2gdPWc,3049
151
172
  uncountable/types/recipe_identifiers.py,sha256=Pi5KX64kzoBp_t_tjb3uVk9Ef1WPIIXGtUdINXi5vcY,654
152
173
  uncountable/types/recipe_identifiers_t.py,sha256=suumLdp_8SZD7oqIxSEaV9D1irMeLJ5_-pJqqUrQ5k8,1786
153
174
  uncountable/types/recipe_inputs.py,sha256=5ThNFEOGbADqqfj1V3hNvZUcO5pn1Gm17LmzQkWDrtI,351
@@ -241,7 +262,7 @@ uncountable/types/api/recipes/get_recipe_calculations.py,sha256=Hq9hmIUBdZ5KRb9F
241
262
  uncountable/types/api/recipes/get_recipe_links.py,sha256=zD4dEIQVAV5wv0NirBldwMEROgvOf7mXT-pWlaSiuN4,975
242
263
  uncountable/types/api/recipes/get_recipe_names.py,sha256=abBQoZ5JRek0UDOE6ETn6NLUQp9stZ5XhCNFdc5kM3U,1004
243
264
  uncountable/types/api/recipes/get_recipe_output_metadata.py,sha256=ORokvA6vvLkkuUHZjY1ym2voxovljMrA48qjt8RD7mU,1456
244
- uncountable/types/api/recipes/get_recipes_data.py,sha256=RIQ_Q7Yb-bZgwrrn36MZAOQ7GUVpsDjDYSia6Tb8Sjc,5454
265
+ uncountable/types/api/recipes/get_recipes_data.py,sha256=s8LE8VQSYgGXEUYpbVJYCOBelrrwMCwXKqiBvvLx9pE,5511
245
266
  uncountable/types/api/recipes/lock_recipes.py,sha256=hdxqn5gOk80Rbw05J9RWi1dPMiW-4wJwkVT_fmKiDUw,1356
246
267
  uncountable/types/api/recipes/remove_recipe_from_project.py,sha256=3lqSU9UQIPMmOIvrdD9H--6QIdolYD0SadIuKxtObaE,872
247
268
  uncountable/types/api/recipes/set_recipe_inputs.py,sha256=yCGFQoguzsSK_weOK4kcvqnUSkRJV1rzKxWnanIZbko,1444
@@ -256,7 +277,7 @@ uncountable/types/api/triggers/__init__.py,sha256=gCgbynxG3jA8FQHzercKtrHKHkiIKr
256
277
  uncountable/types/api/triggers/run_trigger.py,sha256=_Rpha9nxXI3Xr17CrGDtofg4HZ81x2lt0rMZ6As0qfE,893
257
278
  uncountable/types/api/uploader/__init__.py,sha256=gCgbynxG3jA8FQHzercKtrHKHkiIKr8APdZYUniAor8,55
258
279
  uncountable/types/api/uploader/invoke_uploader.py,sha256=Rc77y5q-3R9-SNQgm8P35zKaW2D1Hbtm7PDixnOn1G0,1025
259
- UncountablePythonSDK-0.0.69.dist-info/METADATA,sha256=8gI9naRGRIPuEbO1GWAD8E-WBs1Ya09vjlx4omQj5y8,1993
260
- UncountablePythonSDK-0.0.69.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
261
- UncountablePythonSDK-0.0.69.dist-info/top_level.txt,sha256=1UVGjAU-6hJY9qw2iJ7nCBeEwZ793AEN5ZfKX9A1uj4,31
262
- UncountablePythonSDK-0.0.69.dist-info/RECORD,,
280
+ UncountablePythonSDK-0.0.70.dist-info/METADATA,sha256=_GI0ixHy91cpiMAMQfwyE-Sfvowp982PEBAw6-emnM0,2051
281
+ UncountablePythonSDK-0.0.70.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
282
+ UncountablePythonSDK-0.0.70.dist-info/top_level.txt,sha256=1UVGjAU-6hJY9qw2iJ7nCBeEwZ793AEN5ZfKX9A1uj4,31
283
+ UncountablePythonSDK-0.0.70.dist-info/RECORD,,
@@ -1,11 +1,11 @@
1
- from uncountable.integration.job import CronJob, CronJobArguments, register_job
1
+ from uncountable.integration.job import CronJob, JobArguments, register_job
2
2
  from uncountable.types import entity_t
3
3
  from uncountable.types.job_definition_t import JobResult
4
4
 
5
5
 
6
6
  @register_job
7
7
  class MyCronJob(CronJob):
8
- def run(self, args: CronJobArguments) -> JobResult:
8
+ def run(self, args: JobArguments) -> JobResult:
9
9
  matfam = args.client.get_entities_data(
10
10
  entity_ids=[1],
11
11
  entity_type=entity_t.EntityType.MATERIAL_FAMILY,
@@ -17,7 +17,11 @@ def get_integration_env() -> str | None:
17
17
 
18
18
 
19
19
  def get_webhook_server_port() -> int:
20
- return int(os.environ.get("UNC_WEBHOOK_SERVER_PORT") or 5001)
20
+ return int(os.environ.get("UNC_WEBHOOK_SERVER_PORT", 5001))
21
+
22
+
23
+ def get_local_admin_server_port() -> int:
24
+ return int(os.environ.get("UNC_ADMIN_SERVER_PORT", 50051))
21
25
 
22
26
 
23
27
  def get_otel_enabled() -> bool:
@@ -58,6 +58,7 @@ def main() -> None:
58
58
  client=client,
59
59
  batch_processor=batch_processor,
60
60
  logger=job_logger,
61
+ payload=None,
61
62
  )
62
63
  job_result = execute_job(
63
64
  job_definition=job,
@@ -1,11 +1,11 @@
1
1
  from dataclasses import dataclass
2
2
 
3
3
  from pkgs.argument_parser import CachedParser
4
- from uncountable.core.async_batch import AsyncBatchProcessor
5
- from uncountable.integration.construct_client import construct_uncountable_client
6
- from uncountable.integration.executors.executors import execute_job
7
- from uncountable.integration.job import CronJobArguments
8
- from uncountable.integration.telemetry import JobLogger, get_otel_tracer
4
+ from uncountable.core.environment import get_local_admin_server_port
5
+ from uncountable.integration.queue_runner.command_server.command_client import (
6
+ send_job_queue_message,
7
+ )
8
+ from uncountable.types import queued_job_t
9
9
  from uncountable.types.job_definition_t import JobDefinition, ProfileMetadata
10
10
 
11
11
 
@@ -20,26 +20,10 @@ cron_args_parser = CachedParser(CronJobArgs)
20
20
 
21
21
  def cron_job_executor(**kwargs: dict) -> None:
22
22
  args_passed = cron_args_parser.parse_storage(kwargs)
23
- with get_otel_tracer().start_as_current_span(name="cron_executor") as span:
24
- job_logger = JobLogger(
25
- base_span=span,
26
- profile_metadata=args_passed.profile_metadata,
27
- job_definition=args_passed.definition,
28
- )
29
- client = construct_uncountable_client(
30
- profile_meta=args_passed.profile_metadata, job_logger=job_logger
31
- )
32
- batch_processor = AsyncBatchProcessor(client=client)
33
- args = CronJobArguments(
34
- job_definition=args_passed.definition,
35
- client=client,
36
- batch_processor=batch_processor,
37
- profile_metadata=args_passed.profile_metadata,
38
- logger=job_logger,
39
- )
40
-
41
- execute_job(
42
- args=args,
43
- profile_metadata=args_passed.profile_metadata,
44
- job_definition=args_passed.definition,
45
- )
23
+ send_job_queue_message(
24
+ job_ref_name=args_passed.definition.id,
25
+ payload=queued_job_t.QueuedJobPayload(
26
+ invocation_context=queued_job_t.InvocationContextCron()
27
+ ),
28
+ port=get_local_admin_server_port(),
29
+ )
@@ -1,8 +1,18 @@
1
1
  import os
2
+ from enum import StrEnum
2
3
 
3
4
  from sqlalchemy import create_engine
4
5
  from sqlalchemy.engine.base import Engine
5
6
 
6
7
 
7
- def create_db_engine() -> Engine:
8
- return create_engine(os.environ["UNC_SQLITE_URI"])
8
+ class IntegrationDBService(StrEnum):
9
+ CRON = "cron"
10
+ RUNNER = "runner"
11
+
12
+
13
+ def create_db_engine(service: IntegrationDBService) -> Engine:
14
+ match service:
15
+ case IntegrationDBService.CRON:
16
+ return create_engine(os.environ["UNC_CRON_SQLITE_URI"])
17
+ case IntegrationDBService.RUNNER:
18
+ return create_engine(os.environ["UNC_RUNNER_SQLITE_URI"])
@@ -0,0 +1,25 @@
1
+ from contextlib import _GeneratorContextManager, contextmanager
2
+ from typing import Callable, Generator
3
+
4
+ from sqlalchemy.engine import Engine
5
+ from sqlalchemy.orm import Session, sessionmaker
6
+
7
+ DBSessionMaker = Callable[[], _GeneratorContextManager[Session]]
8
+
9
+
10
+ def get_session_maker(engine: Engine) -> DBSessionMaker:
11
+ session_maker = sessionmaker(bind=engine)
12
+
13
+ @contextmanager
14
+ def session_manager() -> Generator[Session, None, None]:
15
+ session = session_maker()
16
+ try:
17
+ yield session
18
+ session.commit()
19
+ except Exception:
20
+ session.rollback()
21
+ raise
22
+ finally:
23
+ session.close()
24
+
25
+ return session_manager
@@ -1,10 +1,10 @@
1
- from uncountable.integration.db.connect import create_db_engine
1
+ from uncountable.integration.db.connect import IntegrationDBService, create_db_engine
2
2
  from uncountable.integration.scan_profiles import load_profiles
3
3
  from uncountable.integration.server import IntegrationServer
4
4
 
5
5
 
6
- def main(blocking: bool) -> None:
7
- with IntegrationServer(create_db_engine()) as server:
6
+ def main() -> None:
7
+ with IntegrationServer(create_db_engine(IntegrationDBService.CRON)) as server:
8
8
  for profile_details in load_profiles():
9
9
  server.register_profile(
10
10
  profile_name=profile_details.name,
@@ -14,8 +14,8 @@ def main(blocking: bool) -> None:
14
14
  client_options=profile_details.definition.client_options,
15
15
  )
16
16
 
17
- if blocking:
18
- server.serve_forever()
17
+ server.serve_forever()
19
18
 
20
19
 
21
- main(__name__ == "__main__")
20
+ if __name__ == "__main__":
21
+ main()
@@ -165,7 +165,7 @@ def _move_files_post_upload(
165
165
  filesystem_session.move_files([*success_file_transfers, *failed_file_transfers])
166
166
 
167
167
 
168
- class GenericUploadJob(Job):
168
+ class GenericUploadJob(Job[None]):
169
169
  def __init__(
170
170
  self,
171
171
  data_source: GenericUploadDataSource,
@@ -177,6 +177,10 @@ class GenericUploadJob(Job):
177
177
  self.upload_strategy = upload_strategy
178
178
  self.data_source = data_source
179
179
 
180
+ @property
181
+ def payload_type(self) -> type[None]:
182
+ return type(None)
183
+
180
184
  def _construct_filesystem_session(self, args: JobArguments) -> FileSystemSession:
181
185
  match self.data_source:
182
186
  case GenericUploadDataSourceSFTP():
@@ -1,58 +1,85 @@
1
+ import functools
2
+ import typing
1
3
  from abc import ABC, abstractmethod
2
4
  from dataclasses import dataclass
3
5
 
6
+ from pkgs.argument_parser import CachedParser
4
7
  from uncountable.core.async_batch import AsyncBatchProcessor
5
8
  from uncountable.core.client import Client
6
9
  from uncountable.integration.telemetry import JobLogger
7
- from uncountable.types import webhook_job_t
10
+ from uncountable.types import base_t, webhook_job_t
8
11
  from uncountable.types.job_definition_t import JobDefinition, JobResult, ProfileMetadata
9
12
 
10
13
 
11
14
  @dataclass(kw_only=True)
12
- class JobArgumentsBase:
15
+ class JobArguments:
13
16
  job_definition: JobDefinition
14
17
  profile_metadata: ProfileMetadata
15
18
  client: Client
16
19
  batch_processor: AsyncBatchProcessor
17
20
  logger: JobLogger
21
+ payload: base_t.JsonValue
18
22
 
19
23
 
20
- @dataclass(kw_only=True)
21
- class CronJobArguments(JobArgumentsBase):
22
- pass
23
-
24
-
25
- @dataclass(kw_only=True)
26
- class WebhookJobArguments(JobArgumentsBase):
27
- payload: webhook_job_t.WebhookEventBody
24
+ # only for compatibility:
25
+ CronJobArguments = JobArguments
28
26
 
29
27
 
30
- JobArguments = CronJobArguments | WebhookJobArguments
28
+ PT = typing.TypeVar("PT")
31
29
 
32
30
 
33
- class Job(ABC):
31
+ class Job(ABC, typing.Generic[PT]):
34
32
  _unc_job_registered: bool = False
35
33
 
34
+ @property
35
+ @abstractmethod
36
+ def payload_type(self) -> type[PT]: ...
37
+
36
38
  @abstractmethod
37
39
  def run_outer(self, args: JobArguments) -> JobResult: ...
38
40
 
41
+ @functools.cached_property
42
+ def _cached_payload_parser(self) -> CachedParser[PT]:
43
+ return CachedParser(self.payload_type)
44
+
45
+ def get_payload(self, payload: base_t.JsonValue) -> PT:
46
+ return self._cached_payload_parser.parse_storage(payload)
47
+
39
48
 
40
49
  class CronJob(Job):
50
+ @property
51
+ def payload_type(self) -> type[None]:
52
+ return type(None)
53
+
41
54
  def run_outer(self, args: JobArguments) -> JobResult:
42
55
  assert isinstance(args, CronJobArguments)
43
56
  return self.run(args)
44
57
 
45
58
  @abstractmethod
46
- def run(self, args: CronJobArguments) -> JobResult: ...
59
+ def run(self, args: JobArguments) -> JobResult: ...
60
+
61
+
62
+ WPT = typing.TypeVar("WPT")
63
+
47
64
 
65
+ class WebhookJob(Job[webhook_job_t.WebhookEventPayload], typing.Generic[WPT]):
66
+ @property
67
+ def payload_type(self) -> type[webhook_job_t.WebhookEventPayload]:
68
+ return webhook_job_t.WebhookEventPayload
69
+
70
+ @property
71
+ @abstractmethod
72
+ def webhook_payload_type(self) -> type[WPT]: ...
48
73
 
49
- class WebhookJob(Job):
50
74
  def run_outer(self, args: JobArguments) -> JobResult:
51
- assert isinstance(args, WebhookJobArguments)
52
- return self.run(args)
75
+ webhook_body = self.get_payload(args.payload)
76
+ inner_payload = CachedParser(self.webhook_payload_type).parse_api(
77
+ webhook_body.data
78
+ )
79
+ return self.run(args, inner_payload)
53
80
 
54
81
  @abstractmethod
55
- def run(self, args: WebhookJobArguments) -> JobResult: ...
82
+ def run(self, args: JobArguments, payload: WPT) -> JobResult: ...
56
83
 
57
84
 
58
85
  def register_job(cls: type[Job]) -> type[Job]:
File without changes
@@ -0,0 +1,24 @@
1
+ from .command_client import check_health, send_job_queue_message
2
+ from .command_server import serve
3
+ from .types import (
4
+ CommandEnqueueJob,
5
+ CommandEnqueueJobResponse,
6
+ CommandQueue,
7
+ CommandServerBadResponse,
8
+ CommandServerException,
9
+ CommandServerTimeout,
10
+ CommandTask,
11
+ )
12
+
13
+ __all__: list[str] = [
14
+ "serve",
15
+ "check_health",
16
+ "send_job_queue_message",
17
+ "CommandEnqueueJob",
18
+ "CommandEnqueueJobResponse",
19
+ "CommandTask",
20
+ "CommandQueue",
21
+ "CommandServerTimeout",
22
+ "CommandServerException",
23
+ "CommandServerBadResponse",
24
+ ]
@@ -0,0 +1,68 @@
1
+ from contextlib import contextmanager
2
+ from typing import Generator
3
+
4
+ import grpc
5
+ import simplejson as json
6
+
7
+ from pkgs.serialization_util import serialize_for_api
8
+ from uncountable.integration.queue_runner.command_server.protocol.command_server_pb2 import (
9
+ CheckHealthRequest,
10
+ CheckHealthResult,
11
+ EnqueueJobRequest,
12
+ EnqueueJobResult,
13
+ )
14
+ from uncountable.integration.queue_runner.command_server.types import (
15
+ CommandServerBadResponse,
16
+ CommandServerTimeout,
17
+ )
18
+ from uncountable.types import queued_job_t
19
+
20
+ from .protocol.command_server_pb2_grpc import CommandServerStub
21
+
22
+ _LOCAL_RPC_HOST = "localhost"
23
+ _DEFAULT_MESSAGE_TIMEOUT_SECS = 2
24
+
25
+
26
+ @contextmanager
27
+ def command_server_connection(
28
+ host: str, port: int
29
+ ) -> Generator[CommandServerStub, None, None]:
30
+ try:
31
+ with grpc.insecure_channel(f"{host}:{port}") as channel:
32
+ stub = CommandServerStub(channel)
33
+ yield stub
34
+ except grpc._channel._InactiveRpcError as e:
35
+ raise CommandServerTimeout() from e
36
+
37
+
38
+ def send_job_queue_message(
39
+ *,
40
+ job_ref_name: str,
41
+ payload: queued_job_t.QueuedJobPayload,
42
+ host: str = "localhost",
43
+ port: int,
44
+ ) -> str:
45
+ with command_server_connection(host=host, port=port) as stub:
46
+ request = EnqueueJobRequest(
47
+ job_ref_name=job_ref_name,
48
+ serialized_payload=json.dumps(serialize_for_api(payload)),
49
+ )
50
+
51
+ response = stub.EnqueueJob(request, timeout=_DEFAULT_MESSAGE_TIMEOUT_SECS)
52
+
53
+ assert isinstance(response, EnqueueJobResult)
54
+ if not response.successfully_queued:
55
+ raise CommandServerBadResponse("queue operation was not successful")
56
+
57
+ return response.queued_job_uuid
58
+
59
+
60
+ def check_health(*, host: str = _LOCAL_RPC_HOST, port: int) -> bool:
61
+ with command_server_connection(host=host, port=port) as stub:
62
+ request = CheckHealthRequest()
63
+
64
+ response = stub.CheckHealth(request, timeout=_DEFAULT_MESSAGE_TIMEOUT_SECS)
65
+
66
+ assert isinstance(response, CheckHealthResult)
67
+
68
+ return response.success
@@ -0,0 +1,64 @@
1
+ import asyncio
2
+
3
+ import simplejson as json
4
+ from grpc import aio
5
+
6
+ from pkgs.argument_parser import CachedParser
7
+ from uncountable.core.environment import get_local_admin_server_port
8
+ from uncountable.integration.queue_runner.command_server.protocol.command_server_pb2 import (
9
+ CheckHealthRequest,
10
+ CheckHealthResult,
11
+ EnqueueJobRequest,
12
+ EnqueueJobResult,
13
+ )
14
+ from uncountable.integration.queue_runner.command_server.types import (
15
+ CommandEnqueueJob,
16
+ CommandEnqueueJobResponse,
17
+ CommandQueue,
18
+ )
19
+ from uncountable.types import queued_job_t
20
+
21
+ from .protocol.command_server_pb2_grpc import (
22
+ CommandServerServicer,
23
+ add_CommandServerServicer_to_server,
24
+ )
25
+
26
+ queued_job_payload_parser = CachedParser(queued_job_t.QueuedJobPayload)
27
+
28
+
29
+ async def serve(command_queue: CommandQueue) -> None:
30
+ server = aio.server()
31
+
32
+ class CommandServerHandler(CommandServerServicer):
33
+ async def EnqueueJob(
34
+ self, request: EnqueueJobRequest, context: aio.ServicerContext
35
+ ) -> EnqueueJobResult:
36
+ payload_json = json.loads(request.serialized_payload)
37
+ payload = queued_job_payload_parser.parse_api(payload_json)
38
+ response_queue: asyncio.Queue[CommandEnqueueJobResponse] = asyncio.Queue()
39
+ await command_queue.put(
40
+ CommandEnqueueJob(
41
+ job_ref_name=request.job_ref_name,
42
+ payload=payload,
43
+ response_queue=response_queue,
44
+ )
45
+ )
46
+ response = await response_queue.get()
47
+ result = EnqueueJobResult(
48
+ successfully_queued=True, queued_job_uuid=response.queued_job_uuid
49
+ )
50
+ return result
51
+
52
+ async def CheckHealth(
53
+ self, request: CheckHealthRequest, context: aio.ServicerContext
54
+ ) -> CheckHealthResult:
55
+ return CheckHealthResult(success=True)
56
+
57
+ add_CommandServerServicer_to_server(CommandServerHandler(), server)
58
+
59
+ listen_addr = f"[::]:{get_local_admin_server_port()}"
60
+
61
+ server.add_insecure_port(listen_addr)
62
+
63
+ await server.start()
64
+ await server.wait_for_termination()
@@ -0,0 +1,22 @@
1
+ syntax = "proto3";
2
+
3
+ service CommandServer {
4
+ rpc EnqueueJob(EnqueueJobRequest) returns (EnqueueJobResult) {}
5
+ rpc CheckHealth(CheckHealthRequest) returns (CheckHealthResult) {}
6
+ }
7
+
8
+ message EnqueueJobRequest {
9
+ string job_ref_name = 1;
10
+ string serialized_payload = 2;
11
+ }
12
+
13
+ message EnqueueJobResult {
14
+ bool successfully_queued = 1;
15
+ string queued_job_uuid = 2;
16
+ }
17
+
18
+ message CheckHealthRequest {}
19
+
20
+ message CheckHealthResult {
21
+ bool success = 1;
22
+ }
@@ -0,0 +1,40 @@
1
+ # ruff: noqa
2
+ # -*- coding: utf-8 -*-
3
+ # Generated by the protocol buffer compiler. DO NOT EDIT!
4
+ # source: uncountable/integration/queue_runner/command_server/protocol/command_server.proto
5
+ # Protobuf Python Version: 4.25.1
6
+ """Generated protocol buffer code."""
7
+
8
+ from google.protobuf import descriptor as _descriptor
9
+ from google.protobuf import descriptor_pool as _descriptor_pool
10
+ from google.protobuf import symbol_database as _symbol_database
11
+ from google.protobuf.internal import builder as _builder
12
+ # @@protoc_insertion_point(imports)
13
+
14
+ _sym_db = _symbol_database.Default()
15
+
16
+
17
+ DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(
18
+ b'\nQuncountable/integration/queue_runner/command_server/protocol/command_server.proto"E\n\x11\x45nqueueJobRequest\x12\x14\n\x0cjob_ref_name\x18\x01 \x01(\t\x12\x1a\n\x12serialized_payload\x18\x02 \x01(\t"H\n\x10\x45nqueueJobResult\x12\x1b\n\x13successfully_queued\x18\x01 \x01(\x08\x12\x17\n\x0fqueued_job_uuid\x18\x02 \x01(\t"\x14\n\x12\x43heckHealthRequest"$\n\x11\x43heckHealthResult\x12\x0f\n\x07success\x18\x01 \x01(\x08\x32\x80\x01\n\rCommandServer\x12\x35\n\nEnqueueJob\x12\x12.EnqueueJobRequest\x1a\x11.EnqueueJobResult"\x00\x12\x38\n\x0b\x43heckHealth\x12\x13.CheckHealthRequest\x1a\x12.CheckHealthResult"\x00\x62\x06proto3'
19
+ )
20
+
21
+ _globals = globals()
22
+ _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
23
+ _builder.BuildTopDescriptorsAndMessages(
24
+ DESCRIPTOR,
25
+ "uncountable.integration.queue_runner.command_server.protocol.command_server_pb2",
26
+ _globals,
27
+ )
28
+ if _descriptor._USE_C_DESCRIPTORS == False:
29
+ DESCRIPTOR._options = None
30
+ _globals["_ENQUEUEJOBREQUEST"]._serialized_start = 85
31
+ _globals["_ENQUEUEJOBREQUEST"]._serialized_end = 154
32
+ _globals["_ENQUEUEJOBRESULT"]._serialized_start = 156
33
+ _globals["_ENQUEUEJOBRESULT"]._serialized_end = 228
34
+ _globals["_CHECKHEALTHREQUEST"]._serialized_start = 230
35
+ _globals["_CHECKHEALTHREQUEST"]._serialized_end = 250
36
+ _globals["_CHECKHEALTHRESULT"]._serialized_start = 252
37
+ _globals["_CHECKHEALTHRESULT"]._serialized_end = 288
38
+ _globals["_COMMANDSERVER"]._serialized_start = 291
39
+ _globals["_COMMANDSERVER"]._serialized_end = 419
40
+ # @@protoc_insertion_point(module_scope)