UncountablePythonSDK 0.0.69__py3-none-any.whl → 0.0.71__py3-none-any.whl

{UncountablePythonSDK-0.0.69.dist-info → UncountablePythonSDK-0.0.71.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: UncountablePythonSDK
-Version: 0.0.69
+Version: 0.0.71
 Summary: Uncountable SDK
 Project-URL: Homepage, https://github.com/uncountableinc/uncountable-python-sdk
 Project-URL: Repository, https://github.com/uncountableinc/uncountable-python-sdk.git
@@ -35,6 +35,8 @@ Requires-Dist: paramiko ==3.*
 Requires-Dist: boto3 ==1.*
 Requires-Dist: flask ==3.*
 Requires-Dist: simplejson ==3.*
+Requires-Dist: grpcio ==1.*
+Requires-Dist: protobuf ==4.*
 Provides-Extra: test
 Requires-Dist: mypy ==1.* ; extra == 'test'
 Requires-Dist: ruff ==0.* ; extra == 'test'

{UncountablePythonSDK-0.0.69.dist-info → UncountablePythonSDK-0.0.71.dist-info}/RECORD

@@ -22,7 +22,7 @@ examples/set_recipe_metadata_file.py,sha256=oPBhMo9T17zj4YpJRkmVH9lknYcT8livph0K
 examples/set_recipe_output_file_sdk.py,sha256=Lz1amqppnWTX83z-C090wCJ4hcKmCD3kb-4v0uBRi0Y,782
 examples/upload_files.py,sha256=tUfKFqiqwnw08OL5Y8_e4j5pSRhp94cFex8XTuVa_ig,487
 examples/integration-server/pyproject.toml,sha256=yVYpx5VG3bKQCYjmCSKN_MRjxoOoF2xF2KzkCcvDHh4,9134
-examples/integration-server/jobs/materials_auto/example_cron.py,sha256=ITXXIDLrFhwS_YMrcuEBTciGoC8ozjB9SXL-qTCZT2k,707
+examples/integration-server/jobs/materials_auto/example_cron.py,sha256=7VVQ-UJsq3DbGpN3XPnorRVZYo-vCwbfSU3VVDluIzA,699
 examples/integration-server/jobs/materials_auto/profile.yaml,sha256=WZTNlgc8uH412oRmWHWWuenYo_6LnrIL11V0HMkGsB0,352
 pkgs/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 pkgs/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -81,29 +81,48 @@ uncountable/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 uncountable/core/__init__.py,sha256=RFv0kO6rKFf1PtBPu83hCGmxqkJamRtsgQ9_-ztw7tA,341
 uncountable/core/async_batch.py,sha256=Gur0VOS0AH2ugwvk65hwoX-iqwQAAyJaejY_LyAZZPo,1210
 uncountable/core/client.py,sha256=KUJN3XcbawMg_GuJ5DvmDsmhRzsnafCjyq27vOD4jC4,10640
-uncountable/core/environment.py,sha256=4rLWeUIDYB6w1iqrAiSEaD_OqMoo_ed3xcjz5jM4HGU,572
+uncountable/core/environment.py,sha256=n46mWvG5uBOy6H_aetu_iuaYmO_d23uSbQBRgb6xStw,677
 uncountable/core/file_upload.py,sha256=qR7BBBWVxFNrb1_WICreo3dkZygE9lcE1fmZCQrDZU0,3469
 uncountable/core/types.py,sha256=s2CjqYJpsmbC7xMwxxT7kJ_V9bwokrjjWVVjpMcQpKI,333
 uncountable/integration/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-uncountable/integration/cli.py,sha256=R_z84CejxV8QW8pXQWX1RxVgMiA3Sgt_X-zMFVNtXLw,2856
+uncountable/integration/cli.py,sha256=wv9XHXujO0Zop-R4jmjNQttPOjEDQlXPfvmfZ5HW9QE,2898
 uncountable/integration/construct_client.py,sha256=I2XTamht13vs-JYkV4PpNS_Pc4FJm-KVYqNNvxI4qNk,1916
-uncountable/integration/cron.py,sha256=Gwu1Mq4lZZvtvx5tlnWJHnXJp948Xk9CXdS1ULVCVws,1629
-uncountable/integration/entrypoint.py,sha256=wgOXhTzErttRjOzV4rS4psZW5qUKIa5ez89QndQl61k,785
-uncountable/integration/job.py,sha256=UrsZHWXKE2wD5M3lFKJCXvDdWj7QMDtAREKv6RBnT3Q,1548
-uncountable/integration/scan_profiles.py,sha256=fHRS9Su6JlYWKBGZkA3KgUSqbT_PrB91Xe2rkdvv1Ew,1396
-uncountable/integration/scheduler.py,sha256=1YwY2f26Wr1kIlJfUthuRIXsUJZfynf9X3Zu76NK3ts,3262
+uncountable/integration/cron.py,sha256=6eH-kIs3sdYPCyb62_L2M7U_uQTdMTdwY5hreEJb0hw,887
+uncountable/integration/entrypoint.py,sha256=EjJdTmED80MrVcNoQ5BhFP3bAYHdU6CzRGiPSZF9y8w,802
+uncountable/integration/job.py,sha256=af197JUceIKzpIN5C2z8zeZOPhIQ16ipyC6qVt1WXv0,2386
+uncountable/integration/scan_profiles.py,sha256=ro4u5_1ClUeyW4vcWJ9epf2UQaJrjtc7S1R5bZpuHTM,1430
+uncountable/integration/scheduler.py,sha256=sVe7V5zlUbSzSMyC30rAlWR30w0jurhrYYCgK8-lVRo,4546
 uncountable/integration/server.py,sha256=-ALtXNlBDnDbeSJcqAHHJzl46mGuq43aIM9RLzG-ZWA,4701
 uncountable/integration/telemetry.py,sha256=MwQLmgCoxpmA_UTp3e2ZB37wcDzKM0qzm5MrmaiJWhU,7142
 uncountable/integration/db/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-uncountable/integration/db/connect.py,sha256=YtQHJ1DBGPhxKFRCfiXqohOYUceKSxMVOJ88aPI48Ug,181
+uncountable/integration/db/connect.py,sha256=mE3bdV0huclH2iT_dXCQdRL4LkjIuf_myAR64RTWXEs,498
+uncountable/integration/db/session.py,sha256=96cGQXpe6IugBTdSsjdP0S5yhJ6toSmbVB6qhc3FJzE,693
 uncountable/integration/executors/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 uncountable/integration/executors/executors.py,sha256=CbwatKkHrLhnqYr_nsBjr0KYeOn8KqijxrZPHbxChBY,1961
-uncountable/integration/executors/generic_upload_executor.py,sha256=PpxZmGUac4PtwlEJi7-S7scSa3y6HrdBPs_p2QDyqZs,10286
+uncountable/integration/executors/generic_upload_executor.py,sha256=NlW5WcYePPA7_fwp5uW_2afAiQLKK7rCkKF06wQ948E,10375
 uncountable/integration/executors/script_executor.py,sha256=OmSBOtU48G3mqza9c2lCm84pGGyaDk-ZBJCx3RsdJXc,846
+uncountable/integration/queue_runner/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+uncountable/integration/queue_runner/job_scheduler.py,sha256=n6bM6ZqVOPD0PoJuZV5Y5tuhmw2gI-_p6JbnVlK42uI,5016
+uncountable/integration/queue_runner/queue_runner.py,sha256=0BmYu5zHdothTevGsB-nXg6MBd1UD-WkP3h1WCKMdQg,710
+uncountable/integration/queue_runner/types.py,sha256=8qTq29BTSa5rmW6CBlBntP0pNIiDcwu1wHa78pjroS0,219
+uncountable/integration/queue_runner/worker.py,sha256=GNz8PmPUBMLA-R2VWkC_86WOaeHjHFI2mvBlNRcI3wE,4238
+uncountable/integration/queue_runner/command_server/__init__.py,sha256=gQPVILGpWzCr2i5GJyoqna7AOSFvtn4tav69gB78mTQ,571
+uncountable/integration/queue_runner/command_server/command_client.py,sha256=DJb0TUVFkiiLBEQzHSN94sTRnuEbutNEgdN39XmnOXI,2046
+uncountable/integration/queue_runner/command_server/command_server.py,sha256=yyXryhiEC2eGS0yFElLGsVzSKwOuYvj-zp22jQorkv0,2138
+uncountable/integration/queue_runner/command_server/types.py,sha256=A9FpGplHxoSkmi8dn4oGFN7bkstQjY1DoTXypbicHpk,991
+uncountable/integration/queue_runner/command_server/protocol/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+uncountable/integration/queue_runner/command_server/protocol/command_server.proto,sha256=pf7FAT2eGuao0VYCFrgTAsM-tiPi1Bhz19XN5So1WFk,439
+uncountable/integration/queue_runner/command_server/protocol/command_server_pb2.py,sha256=-lBTc5Tz48agqNSeOSpBE69e2kRmWF59sUaowCl8p7U,2207
+uncountable/integration/queue_runner/command_server/protocol/command_server_pb2.pyi,sha256=9viBn6PHvtfMSRwam57ke5O2D_k8LapWYVfBRjknIYg,1281
+uncountable/integration/queue_runner/command_server/protocol/command_server_pb2_grpc.py,sha256=ZVHkuLDjEbXMCxBsw1UrRhT3EEF8CDDqEvmE3Kbp1H4,5359
+uncountable/integration/queue_runner/datastore/__init__.py,sha256=6BefApqN8D2zlVOH14QAeVzwQ8j5NIb41-njT02Za0k,88
+uncountable/integration/queue_runner/datastore/datastore_sqlite.py,sha256=UZQABTrM3HxVaojCnI1pTU3v1GsbE3G-OCgXP5ekcxI,3801
+uncountable/integration/queue_runner/datastore/interface.py,sha256=j4D-zVvLq-48VTVwHVei82UVUJ_P3cxiseyiTl0MoNw,534
+uncountable/integration/queue_runner/datastore/model.py,sha256=8-RI5A2yPZVGBLWINVmMd6VOl_oHtqGtnaNXcapAChw,577
 uncountable/integration/secret_retrieval/__init__.py,sha256=3QXVj35w8rRMxVvmmsViFYDi3lcb3g70incfalOEm6o,87
 uncountable/integration/secret_retrieval/retrieve_secret.py,sha256=eoPWbkUtCn_63A4TFlK_nvEDvfm4u2fiOoglmAkBG3U,3004
-uncountable/integration/webhook_server/entrypoint.py,sha256=01yx5I_4e380oGqaceCcsGhV1JPHcY5lXIFJoghFmwc,6179
-uncountable/types/__init__.py,sha256=MDzXxVKMd9OY6WGXqTRRO_XxgHeHy9OOig9L__BlF9s,8498
+uncountable/integration/webhook_server/entrypoint.py,sha256=hgbEtdVo3QU3odlHSygxmrsxR9g7rgU0yKt-o9nVAHE,5686
+uncountable/types/__init__.py,sha256=KSsSEBnGhn88NPex5q9MZtY4kj8PRqVTbaKtko4hxUU,8561
 uncountable/types/async_batch.py,sha256=_OhT25_dEVts_z_n1kqfJH3xlZg3btLqR6TNkfFLlXE,609
 uncountable/types/async_batch_processor.py,sha256=obVzN-PcYLV2pHScszfCGjSq6-Xc34WM1ysx6Fv6tZk,11293
 uncountable/types/async_batch_t.py,sha256=ipSGz93O1KB-WE2dvlvflTKS51rJrf3bJkUojyxos7I,2193
@@ -148,6 +167,8 @@ uncountable/types/phases.py,sha256=YCsU77DdjRJJWdLTwLuOZNG4e9ML82NIBI1xTWr3ggA,2
 uncountable/types/phases_t.py,sha256=21_4-O5MlO1AzsZijSlSUeEHknY0T9KCPilFMG8GMEo,544
 uncountable/types/post_base.py,sha256=GES5_IhXFAjpzI6ChbVP4zTkX6f3tPUIHST1sxsRN6Q,279
 uncountable/types/post_base_t.py,sha256=2et3TDnFChZcx0RWU-18RJHw9Yiyj2TJz-2tByHXwaI,770
+uncountable/types/queued_job.py,sha256=67exX5vfpewKzvFKxuxTLsSJTDKlE3JqKfWnnYx2Jos,842
+uncountable/types/queued_job_t.py,sha256=yd3FhtcbfMxPjBq_ykLY2ASx_EJvQ4kW3MKIL2gdPWc,3049
 uncountable/types/recipe_identifiers.py,sha256=Pi5KX64kzoBp_t_tjb3uVk9Ef1WPIIXGtUdINXi5vcY,654
 uncountable/types/recipe_identifiers_t.py,sha256=suumLdp_8SZD7oqIxSEaV9D1irMeLJ5_-pJqqUrQ5k8,1786
 uncountable/types/recipe_inputs.py,sha256=5ThNFEOGbADqqfj1V3hNvZUcO5pn1Gm17LmzQkWDrtI,351
@@ -241,14 +262,14 @@ uncountable/types/api/recipes/get_recipe_calculations.py,sha256=Hq9hmIUBdZ5KRb9F
 uncountable/types/api/recipes/get_recipe_links.py,sha256=zD4dEIQVAV5wv0NirBldwMEROgvOf7mXT-pWlaSiuN4,975
 uncountable/types/api/recipes/get_recipe_names.py,sha256=abBQoZ5JRek0UDOE6ETn6NLUQp9stZ5XhCNFdc5kM3U,1004
 uncountable/types/api/recipes/get_recipe_output_metadata.py,sha256=ORokvA6vvLkkuUHZjY1ym2voxovljMrA48qjt8RD7mU,1456
-uncountable/types/api/recipes/get_recipes_data.py,sha256=RIQ_Q7Yb-bZgwrrn36MZAOQ7GUVpsDjDYSia6Tb8Sjc,5454
+uncountable/types/api/recipes/get_recipes_data.py,sha256=s8LE8VQSYgGXEUYpbVJYCOBelrrwMCwXKqiBvvLx9pE,5511
 uncountable/types/api/recipes/lock_recipes.py,sha256=hdxqn5gOk80Rbw05J9RWi1dPMiW-4wJwkVT_fmKiDUw,1356
 uncountable/types/api/recipes/remove_recipe_from_project.py,sha256=3lqSU9UQIPMmOIvrdD9H--6QIdolYD0SadIuKxtObaE,872
 uncountable/types/api/recipes/set_recipe_inputs.py,sha256=yCGFQoguzsSK_weOK4kcvqnUSkRJV1rzKxWnanIZbko,1444
 uncountable/types/api/recipes/set_recipe_metadata.py,sha256=J2F1FDgFLKjlbGw29TW1Xefy38Zt-C7V1WDK70gRiuY,914
 uncountable/types/api/recipes/set_recipe_output_annotations.py,sha256=ayG9vnCroYGQ3mmNBALLlQhf9otrR-q0QiNLnWgmIsw,2930
 uncountable/types/api/recipes/set_recipe_output_file.py,sha256=8s0oxgjhv_hIyV3EoAjIBb970TrJJbKTFaSdNvbgU78,1225
-uncountable/types/api/recipes/set_recipe_outputs.py,sha256=I7NUDO_QrKt1rZB7zHr6WWS2FdXCK8Bqc0BxFkDfjyQ,1729
+uncountable/types/api/recipes/set_recipe_outputs.py,sha256=bLR-FtncB6hzqvzQhCts6ERdHJWgPMne4bQi6-xxu_k,1890
 uncountable/types/api/recipes/set_recipe_tags.py,sha256=kYX13RvHyUOlzL6LbBKDQ8yNvpvNobbnF2WDtgZsKRc,2493
 uncountable/types/api/recipes/unarchive_recipes.py,sha256=KZKzw0fbQdOBUR3YBunFrHFcex4qPMiA8mDgsB6n24k,814
 uncountable/types/api/recipes/unlock_recipes.py,sha256=AvzQeZCLs9i7CuhMs3Xltdi4n0q4Td55yHYyBGQgWag,1103
@@ -256,7 +277,7 @@ uncountable/types/api/triggers/__init__.py,sha256=gCgbynxG3jA8FQHzercKtrHKHkiIKr
 uncountable/types/api/triggers/run_trigger.py,sha256=_Rpha9nxXI3Xr17CrGDtofg4HZ81x2lt0rMZ6As0qfE,893
 uncountable/types/api/uploader/__init__.py,sha256=gCgbynxG3jA8FQHzercKtrHKHkiIKr8APdZYUniAor8,55
 uncountable/types/api/uploader/invoke_uploader.py,sha256=Rc77y5q-3R9-SNQgm8P35zKaW2D1Hbtm7PDixnOn1G0,1025
-UncountablePythonSDK-0.0.69.dist-info/METADATA,sha256=8gI9naRGRIPuEbO1GWAD8E-WBs1Ya09vjlx4omQj5y8,1993
-UncountablePythonSDK-0.0.69.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
-UncountablePythonSDK-0.0.69.dist-info/top_level.txt,sha256=1UVGjAU-6hJY9qw2iJ7nCBeEwZ793AEN5ZfKX9A1uj4,31
-UncountablePythonSDK-0.0.69.dist-info/RECORD,,
+UncountablePythonSDK-0.0.71.dist-info/METADATA,sha256=-WFnnxQlZ9qZ37_W7tekpeHnbqw0uyYzNqqAbd8BL_4,2051
+UncountablePythonSDK-0.0.71.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
+UncountablePythonSDK-0.0.71.dist-info/top_level.txt,sha256=1UVGjAU-6hJY9qw2iJ7nCBeEwZ793AEN5ZfKX9A1uj4,31
+UncountablePythonSDK-0.0.71.dist-info/RECORD,,

examples/integration-server/jobs/materials_auto/example_cron.py

@@ -1,11 +1,11 @@
-from uncountable.integration.job import CronJob, CronJobArguments, register_job
+from uncountable.integration.job import CronJob, JobArguments, register_job
 from uncountable.types import entity_t
 from uncountable.types.job_definition_t import JobResult
 
 
 @register_job
 class MyCronJob(CronJob):
-    def run(self, args: CronJobArguments) -> JobResult:
+    def run(self, args: JobArguments) -> JobResult:
         matfam = args.client.get_entities_data(
             entity_ids=[1],
             entity_type=entity_t.EntityType.MATERIAL_FAMILY,

uncountable/core/environment.py

@@ -17,7 +17,11 @@ def get_integration_env() -> str | None:
 
 
 def get_webhook_server_port() -> int:
-    return int(os.environ.get("UNC_WEBHOOK_SERVER_PORT") or 5001)
+    return int(os.environ.get("UNC_WEBHOOK_SERVER_PORT", 5001))
+
+
+def get_local_admin_server_port() -> int:
+    return int(os.environ.get("UNC_ADMIN_SERVER_PORT", 50051))
 
 
 def get_otel_enabled() -> bool:

uncountable/integration/cli.py

@@ -58,6 +58,7 @@ def main() -> None:
                             client=client,
                             batch_processor=batch_processor,
                             logger=job_logger,
+                            payload=None,
                         )
                         job_result = execute_job(
                             job_definition=job,

uncountable/integration/cron.py

@@ -1,11 +1,11 @@
 from dataclasses import dataclass
 
 from pkgs.argument_parser import CachedParser
-from uncountable.core.async_batch import AsyncBatchProcessor
-from uncountable.integration.construct_client import construct_uncountable_client
-from uncountable.integration.executors.executors import execute_job
-from uncountable.integration.job import CronJobArguments
-from uncountable.integration.telemetry import JobLogger, get_otel_tracer
+from uncountable.core.environment import get_local_admin_server_port
+from uncountable.integration.queue_runner.command_server.command_client import (
+    send_job_queue_message,
+)
+from uncountable.types import queued_job_t
 from uncountable.types.job_definition_t import JobDefinition, ProfileMetadata
 
 
@@ -20,26 +20,10 @@ cron_args_parser = CachedParser(CronJobArgs)
 
 def cron_job_executor(**kwargs: dict) -> None:
     args_passed = cron_args_parser.parse_storage(kwargs)
-    with get_otel_tracer().start_as_current_span(name="cron_executor") as span:
-        job_logger = JobLogger(
-            base_span=span,
-            profile_metadata=args_passed.profile_metadata,
-            job_definition=args_passed.definition,
-        )
-        client = construct_uncountable_client(
-            profile_meta=args_passed.profile_metadata, job_logger=job_logger
-        )
-        batch_processor = AsyncBatchProcessor(client=client)
-        args = CronJobArguments(
-            job_definition=args_passed.definition,
-            client=client,
-            batch_processor=batch_processor,
-            profile_metadata=args_passed.profile_metadata,
-            logger=job_logger,
-        )
-
-        execute_job(
-            args=args,
-            profile_metadata=args_passed.profile_metadata,
-            job_definition=args_passed.definition,
-        )
+    send_job_queue_message(
+        job_ref_name=args_passed.definition.id,
+        payload=queued_job_t.QueuedJobPayload(
+            invocation_context=queued_job_t.InvocationContextCron()
+        ),
+        port=get_local_admin_server_port(),
+    )

uncountable/integration/db/connect.py

@@ -1,8 +1,18 @@
 import os
+from enum import StrEnum
 
 from sqlalchemy import create_engine
 from sqlalchemy.engine.base import Engine
 
 
-def create_db_engine() -> Engine:
-    return create_engine(os.environ["UNC_SQLITE_URI"])
+class IntegrationDBService(StrEnum):
+    CRON = "cron"
+    RUNNER = "runner"
+
+
+def create_db_engine(service: IntegrationDBService) -> Engine:
+    match service:
+        case IntegrationDBService.CRON:
+            return create_engine(os.environ["UNC_CRON_SQLITE_URI"])
+        case IntegrationDBService.RUNNER:
+            return create_engine(os.environ["UNC_RUNNER_SQLITE_URI"])

uncountable/integration/db/session.py

@@ -0,0 +1,25 @@
+from contextlib import _GeneratorContextManager, contextmanager
+from typing import Callable, Generator
+
+from sqlalchemy.engine import Engine
+from sqlalchemy.orm import Session, sessionmaker
+
+DBSessionMaker = Callable[[], _GeneratorContextManager[Session]]
+
+
+def get_session_maker(engine: Engine) -> DBSessionMaker:
+    session_maker = sessionmaker(bind=engine)
+
+    @contextmanager
+    def session_manager() -> Generator[Session, None, None]:
+        session = session_maker()
+        try:
+            yield session
+            session.commit()
+        except Exception:
+            session.rollback()
+            raise
+        finally:
+            session.close()
+
+    return session_manager

uncountable/integration/entrypoint.py

@@ -1,10 +1,10 @@
-from uncountable.integration.db.connect import create_db_engine
+from uncountable.integration.db.connect import IntegrationDBService, create_db_engine
 from uncountable.integration.scan_profiles import load_profiles
 from uncountable.integration.server import IntegrationServer
 
 
-def main(blocking: bool) -> None:
-    with IntegrationServer(create_db_engine()) as server:
+def main() -> None:
+    with IntegrationServer(create_db_engine(IntegrationDBService.CRON)) as server:
         for profile_details in load_profiles():
             server.register_profile(
                 profile_name=profile_details.name,
@@ -14,8 +14,8 @@ def main(blocking: bool) -> None:
                 client_options=profile_details.definition.client_options,
             )
 
-        if blocking:
-            server.serve_forever()
+        server.serve_forever()
 
 
-main(__name__ == "__main__")
+if __name__ == "__main__":
+    main()

uncountable/integration/executors/generic_upload_executor.py

@@ -165,7 +165,7 @@ def _move_files_post_upload(
     filesystem_session.move_files([*success_file_transfers, *failed_file_transfers])
 
 
-class GenericUploadJob(Job):
+class GenericUploadJob(Job[None]):
     def __init__(
         self,
         data_source: GenericUploadDataSource,
@@ -177,6 +177,10 @@ class GenericUploadJob(Job):
         self.upload_strategy = upload_strategy
         self.data_source = data_source
 
+    @property
+    def payload_type(self) -> type[None]:
+        return type(None)
+
     def _construct_filesystem_session(self, args: JobArguments) -> FileSystemSession:
         match self.data_source:
             case GenericUploadDataSourceSFTP():

uncountable/integration/job.py

@@ -1,58 +1,85 @@
+import functools
+import typing
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
 
+from pkgs.argument_parser import CachedParser
 from uncountable.core.async_batch import AsyncBatchProcessor
 from uncountable.core.client import Client
 from uncountable.integration.telemetry import JobLogger
-from uncountable.types import webhook_job_t
+from uncountable.types import base_t, webhook_job_t
 from uncountable.types.job_definition_t import JobDefinition, JobResult, ProfileMetadata
 
 
 @dataclass(kw_only=True)
-class JobArgumentsBase:
+class JobArguments:
     job_definition: JobDefinition
     profile_metadata: ProfileMetadata
     client: Client
     batch_processor: AsyncBatchProcessor
     logger: JobLogger
+    payload: base_t.JsonValue
 
 
-@dataclass(kw_only=True)
-class CronJobArguments(JobArgumentsBase):
-    pass
-
-
-@dataclass(kw_only=True)
-class WebhookJobArguments(JobArgumentsBase):
-    payload: webhook_job_t.WebhookEventBody
+# only for compatibility:
+CronJobArguments = JobArguments
 
 
-JobArguments = CronJobArguments | WebhookJobArguments
+PT = typing.TypeVar("PT")
 
 
-class Job(ABC):
+class Job(ABC, typing.Generic[PT]):
     _unc_job_registered: bool = False
 
+    @property
+    @abstractmethod
+    def payload_type(self) -> type[PT]: ...
+
     @abstractmethod
     def run_outer(self, args: JobArguments) -> JobResult: ...
 
+    @functools.cached_property
+    def _cached_payload_parser(self) -> CachedParser[PT]:
+        return CachedParser(self.payload_type)
+
+    def get_payload(self, payload: base_t.JsonValue) -> PT:
+        return self._cached_payload_parser.parse_storage(payload)
+
 
 class CronJob(Job):
+    @property
+    def payload_type(self) -> type[None]:
+        return type(None)
+
     def run_outer(self, args: JobArguments) -> JobResult:
         assert isinstance(args, CronJobArguments)
         return self.run(args)
 
     @abstractmethod
-    def run(self, args: CronJobArguments) -> JobResult: ...
+    def run(self, args: JobArguments) -> JobResult: ...
+
+
+WPT = typing.TypeVar("WPT")
+
 
+class WebhookJob(Job[webhook_job_t.WebhookEventPayload], typing.Generic[WPT]):
+    @property
+    def payload_type(self) -> type[webhook_job_t.WebhookEventPayload]:
+        return webhook_job_t.WebhookEventPayload
+
+    @property
+    @abstractmethod
+    def webhook_payload_type(self) -> type[WPT]: ...
 
-class WebhookJob(Job):
     def run_outer(self, args: JobArguments) -> JobResult:
-        assert isinstance(args, WebhookJobArguments)
-        return self.run(args)
+        webhook_body = self.get_payload(args.payload)
+        inner_payload = CachedParser(self.webhook_payload_type).parse_api(
+            webhook_body.data
+        )
+        return self.run(args, inner_payload)
 
     @abstractmethod
-    def run(self, args: WebhookJobArguments) -> JobResult: ...
+    def run(self, args: JobArguments, payload: WPT) -> JobResult: ...
 
 
 def register_job(cls: type[Job]) -> type[Job]:

uncountable/integration/queue_runner/command_server/__init__.py

@@ -0,0 +1,24 @@
+from .command_client import check_health, send_job_queue_message
+from .command_server import serve
+from .types import (
+    CommandEnqueueJob,
+    CommandEnqueueJobResponse,
+    CommandQueue,
+    CommandServerBadResponse,
+    CommandServerException,
+    CommandServerTimeout,
+    CommandTask,
+)
+
+__all__: list[str] = [
+    "serve",
+    "check_health",
+    "send_job_queue_message",
+    "CommandEnqueueJob",
+    "CommandEnqueueJobResponse",
+    "CommandTask",
+    "CommandQueue",
+    "CommandServerTimeout",
+    "CommandServerException",
+    "CommandServerBadResponse",
+]

uncountable/integration/queue_runner/command_server/command_client.py

@@ -0,0 +1,68 @@
+from contextlib import contextmanager
+from typing import Generator
+
+import grpc
+import simplejson as json
+
+from pkgs.serialization_util import serialize_for_api
+from uncountable.integration.queue_runner.command_server.protocol.command_server_pb2 import (
+    CheckHealthRequest,
+    CheckHealthResult,
+    EnqueueJobRequest,
+    EnqueueJobResult,
+)
+from uncountable.integration.queue_runner.command_server.types import (
+    CommandServerBadResponse,
+    CommandServerTimeout,
+)
+from uncountable.types import queued_job_t
+
+from .protocol.command_server_pb2_grpc import CommandServerStub
+
+_LOCAL_RPC_HOST = "localhost"
+_DEFAULT_MESSAGE_TIMEOUT_SECS = 2
+
+
+@contextmanager
+def command_server_connection(
+    host: str, port: int
+) -> Generator[CommandServerStub, None, None]:
+    try:
+        with grpc.insecure_channel(f"{host}:{port}") as channel:
+            stub = CommandServerStub(channel)
+            yield stub
+    except grpc._channel._InactiveRpcError as e:
+        raise CommandServerTimeout() from e
+
+
+def send_job_queue_message(
+    *,
+    job_ref_name: str,
+    payload: queued_job_t.QueuedJobPayload,
+    host: str = "localhost",
+    port: int,
+) -> str:
+    with command_server_connection(host=host, port=port) as stub:
+        request = EnqueueJobRequest(
+            job_ref_name=job_ref_name,
+            serialized_payload=json.dumps(serialize_for_api(payload)),
+        )
+
+        response = stub.EnqueueJob(request, timeout=_DEFAULT_MESSAGE_TIMEOUT_SECS)
+
+        assert isinstance(response, EnqueueJobResult)
+        if not response.successfully_queued:
+            raise CommandServerBadResponse("queue operation was not successful")
+
+        return response.queued_job_uuid
+
+
+def check_health(*, host: str = _LOCAL_RPC_HOST, port: int) -> bool:
+    with command_server_connection(host=host, port=port) as stub:
+        request = CheckHealthRequest()
+
+        response = stub.CheckHealth(request, timeout=_DEFAULT_MESSAGE_TIMEOUT_SECS)
+
+        assert isinstance(response, CheckHealthResult)
+
+        return response.success

uncountable/integration/queue_runner/command_server/command_server.py

@@ -0,0 +1,64 @@
+import asyncio
+
+import simplejson as json
+from grpc import aio
+
+from pkgs.argument_parser import CachedParser
+from uncountable.core.environment import get_local_admin_server_port
+from uncountable.integration.queue_runner.command_server.protocol.command_server_pb2 import (
+    CheckHealthRequest,
+    CheckHealthResult,
+    EnqueueJobRequest,
+    EnqueueJobResult,
+)
+from uncountable.integration.queue_runner.command_server.types import (
+    CommandEnqueueJob,
+    CommandEnqueueJobResponse,
+    CommandQueue,
+)
+from uncountable.types import queued_job_t
+
+from .protocol.command_server_pb2_grpc import (
+    CommandServerServicer,
+    add_CommandServerServicer_to_server,
+)
+
+queued_job_payload_parser = CachedParser(queued_job_t.QueuedJobPayload)
+
+
+async def serve(command_queue: CommandQueue) -> None:
+    server = aio.server()
+
+    class CommandServerHandler(CommandServerServicer):
+        async def EnqueueJob(
+            self, request: EnqueueJobRequest, context: aio.ServicerContext
+        ) -> EnqueueJobResult:
+            payload_json = json.loads(request.serialized_payload)
+            payload = queued_job_payload_parser.parse_api(payload_json)
+            response_queue: asyncio.Queue[CommandEnqueueJobResponse] = asyncio.Queue()
+            await command_queue.put(
+                CommandEnqueueJob(
+                    job_ref_name=request.job_ref_name,
+                    payload=payload,
+                    response_queue=response_queue,
+                )
+            )
+            response = await response_queue.get()
+            result = EnqueueJobResult(
+                successfully_queued=True, queued_job_uuid=response.queued_job_uuid
+            )
+            return result
+
+        async def CheckHealth(
+            self, request: CheckHealthRequest, context: aio.ServicerContext
+        ) -> CheckHealthResult:
+            return CheckHealthResult(success=True)
+
+    add_CommandServerServicer_to_server(CommandServerHandler(), server)
+
+    listen_addr = f"[::]:{get_local_admin_server_port()}"
+
+    server.add_insecure_port(listen_addr)
+
+    await server.start()
+    await server.wait_for_termination()

uncountable/integration/queue_runner/command_server/protocol/command_server.proto

@@ -0,0 +1,22 @@
+syntax = "proto3";
+
+service CommandServer {
+	rpc EnqueueJob(EnqueueJobRequest) returns (EnqueueJobResult) {}
+	rpc CheckHealth(CheckHealthRequest) returns (CheckHealthResult) {}
+}
+
+message EnqueueJobRequest {
+	string job_ref_name = 1;
+	string serialized_payload = 2;
+}
+
+message EnqueueJobResult {
+	bool successfully_queued = 1;
+	string queued_job_uuid = 2;
+}
+
+message CheckHealthRequest {}
+
+message CheckHealthResult {
+	bool success = 1;
+}

uncountable/integration/queue_runner/command_server/protocol/command_server_pb2.py

@@ -0,0 +1,40 @@
+# ruff: noqa
+# -*- coding: utf-8 -*-
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# source: uncountable/integration/queue_runner/command_server/protocol/command_server.proto
+# Protobuf Python Version: 4.25.1
+"""Generated protocol buffer code."""
+
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import descriptor_pool as _descriptor_pool
+from google.protobuf import symbol_database as _symbol_database
+from google.protobuf.internal import builder as _builder
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(
+    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'
+)
+
+_globals = globals()
+_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
+_builder.BuildTopDescriptorsAndMessages(
+    DESCRIPTOR,
+    "uncountable.integration.queue_runner.command_server.protocol.command_server_pb2",
+    _globals,
+)
+if _descriptor._USE_C_DESCRIPTORS == False:
+    DESCRIPTOR._options = None
+    _globals["_ENQUEUEJOBREQUEST"]._serialized_start = 85
+    _globals["_ENQUEUEJOBREQUEST"]._serialized_end = 154
+    _globals["_ENQUEUEJOBRESULT"]._serialized_start = 156
+    _globals["_ENQUEUEJOBRESULT"]._serialized_end = 228
+    _globals["_CHECKHEALTHREQUEST"]._serialized_start = 230
+    _globals["_CHECKHEALTHREQUEST"]._serialized_end = 250
+    _globals["_CHECKHEALTHRESULT"]._serialized_start = 252
+    _globals["_CHECKHEALTHRESULT"]._serialized_end = 288
+    _globals["_COMMANDSERVER"]._serialized_start = 291
+    _globals["_COMMANDSERVER"]._serialized_end = 419
+# @@protoc_insertion_point(module_scope)