feat: Very rough first pass at client libraries

.pre-commit-config.yaml

@@ -52,8 +52,17 @@ repos:
       require_serial: true
 
     - id: mypy
-      name: mypy
-      entry: mypy
+      name: mypy-tests
+      entry: "mypy"
+      exclude: "clients/"
+      language: system
+      types: [python]
+      require_serial: true
+
+    - id: mypy
+      name: mypy-client
+      entry: "mypy"
+      exclude: "integration_tests/"
       language: system
       types: [python]
       require_serial: true

Makefile

@@ -46,33 +46,33 @@ reset-kafka: setup ## Reset kafka
 .PHONY: reset-kafka
 
 test-rebalance: build reset-kafka ## Run the rebalance integration test
-	python -m pytest python/integration_tests/test_consumer_rebalancing.py -s
-	rm -r python/integration_tests/.tests_output/test_consumer_rebalancing
+	python -m pytest integration_tests/test_consumer_rebalancing.py -s
+	rm -r integration_tests/.tests_output/test_consumer_rebalancing
 .PHONY: test-rebalance
 
 test-worker-processing: build reset-kafka ## Run the worker processing integration test
-	python -m pytest python/integration_tests/test_task_worker_processing.py -s
-	rm -r python/integration_tests/.tests_output/test_task_worker_processing
+	python -m pytest integration_tests/test_task_worker_processing.py -s
+	rm -r integration_tests/.tests_output/test_task_worker_processing
 .PHONY: test-worker-processing
 
 test-upkeep-retry: build reset-kafka ## Run the upkeep retry integration test
-	python -m pytest python/integration_tests/test_upkeep_retry.py -s
-	rm -r python/integration_tests/.tests_output/test_upkeep_retry
+	python -m pytest integration_tests/test_upkeep_retry.py -s
+	rm -r integration_tests/.tests_output/test_upkeep_retry
 .PHONY: test-upkeep-retry
 
 test-upkeep-expiry: build reset-kafka ## Run the upkeep expiry integration test
-	python -m pytest python/integration_tests/test_upkeep_expiry.py -s
-	rm -r python/integration_tests/.tests_output/test_upkeep_expiry
+	python -m pytest integration_tests/test_upkeep_expiry.py -s
+	rm -r integration_tests/.tests_output/test_upkeep_expiry
 .PHONY: test-upkeep-expiry
 
 test-upkeep-delay: build reset-kafka ## Run the upkeep delay integration test
-	python -m pytest python/integration_tests/test_upkeep_delay.py -s
-	rm -r python/integration_tests/.tests_output/test_upkeep_delay
+	python -m pytest integration_tests/test_upkeep_delay.py -s
+	rm -r integration_tests/.tests_output/test_upkeep_delay
 .PHONY: test-upkeep-delay
 
 test-failed-tasks: build reset-kafka ## Run the failed tasks integration test
-	python -m pytest python/integration_tests/test_failed_tasks.py -s
-	rm -r python/integration_tests/.tests_output/test_failed_tasks
+	python -m pytest integration_tests/test_failed_tasks.py -s
+	rm -r integration_tests/.tests_output/test_failed_tasks
 .PHONY: test-failed-tasks
 
 integration-test: test-rebalance test-worker-processing test-upkeep-retry test-upkeep-expiry test-upkeep-delay test-failed-tasks ## Run all integration tests

clients/python/.python-version

@@ -0,0 +1 @@
+3.12.11

clients/python/README.md

@@ -0,0 +1,39 @@
+# Taskbroker python client
+
+This package provides python client libraries for taskbroker. The client libraries help with:
+
+- Defining tasks
+- Spawning tasks
+- Running periodic schedulers
+- Running workers
+
+## Example application
+
+The `src/examples` directory contains a sample application that is used in tests, and can be run locally. First
+install the required dependencies:
+
+```bash
+# Install the required and development dependencies
+uv sync --dev
+
+# Install the optional dependencies for the example application
+uv sync --extra examples
+```
+Before running the examples, make sure you have:
+
+- Kafka running
+- Redis running (for the scheduler)
+- Taskbroker running. Use `cargo run` in the repository root for this.
+
+With all of those pre-requisites complete, you can run the example application:
+
+```bash
+# Generate 5 tasks
+python src/examples/cli.py generate --count 5
+
+# Run the scheduler which emits a task every 1m
+python src/examples/cli.py scheduler
+
+# Run the worker
+python src/examples/cli.py scheduler
+```

clients/python/pyproject.toml

@@ -0,0 +1,90 @@
+[project]
+name = "taskbroker-client"
+version = "0.1.0"
+description = "Taskbroker python client and worker runtime"
+readme = "README.md"
+dependencies = [
+    "sentry-arroyo>=2.33.1",
+    "sentry-sdk[http2]>=2.43.0",
+    "sentry-protos>=0.4.11",
+    "confluent_kafka>=2.3.0",
+    "cronsim>=2.6",
+    "grpcio==1.66.1",
+    "orjson>=3.10.10",
+    "protobuf>=5.28.3",
+    "types-protobuf>=6.30.2.20250703",
+    "redis>=3.4.1",
+    "redis-py-cluster>=2.1.0",
+    "zstandard>=0.18.0",
+]
+
+[dependency-groups]
+dev = [
+    "devservices>=1.2.1",
+    "sentry-devenv>=1.22.2",
+    "black==24.10.0",
+    "pre-commit>=4.2.0",
+    "pytest>=8.3.3",
+    "flake8>=7.3.0",
+    "isort>=5.13.2",
+    "mypy>=1.17.1",
+    "time-machine>=2.16.0",
+]
+[project.optional-dependencies]
+examples = [
+    "click>=8.3"
+]
+
+[build-system]
+requires = ["uv_build>=0.8.2,<0.9.0"]
+build-backend = "uv_build"
+
+[tool.uv]
+environments = ["sys_platform == 'darwin' or sys_platform == 'linux'"]
+
+[[tool.uv.index]]
+url = "https://pypi.devinfra.sentry.io/simple"
+default = true
+
+[tool.pytest.ini_options]
+pythonpath = ["python"]
+testpaths = ["tests"]
+python_files = ["test_*.py"]
+python_functions = ["test_*"]
+
+[tool.mypy]
+files = ["."]
+mypy_path = ["src"]
+explicit_package_bases = true
+# minimal strictness settings
+check_untyped_defs = true
+no_implicit_reexport = true
+warn_unreachable = true
+warn_unused_configs = true
+warn_unused_ignores = true
+warn_redundant_casts = true
+enable_error_code = ["ignore-without-code", "redundant-self"]
+local_partial_types = true  # compat with dmypy
+disallow_any_generics = true
+disallow_untyped_defs = true
+
+# begin: missing 3rd party stubs
+[[tool.mypy.overrides]]
+module = [
+    ".conftest",
+    "redis.*",
+    "rediscluster.*",
+    "confluent_kafka.*",
+]
+ignore_missing_imports = true
+# end: missing 3rd party stubs
+
+[tool.black]
+# File filtering is taken care of in pre-commit.
+line-length = 100
+target-version = ['py311']
+
+[tool.isort]
+profile = "black"
+line_length = 100
+lines_between_sections = 1

clients/python/src/examples/app.py

@@ -0,0 +1,22 @@
+from arroyo.backends.kafka import KafkaProducer
+
+from examples.store import StubAtMostOnce
+from taskbroker_client.app import TaskbrokerApp
+
+
+def producer_factory(topic: str) -> KafkaProducer:
+    # TODO use env vars for kafka host/port
+    config = {
+        "bootstrap.servers": "127.0.0.1:9092",
+        "compression.type": "lz4",
+        "message.max.bytes": 50000000,  # 50MB
+    }
+    return KafkaProducer(config)
+
+
+app = TaskbrokerApp(
+    name="example-app",
+    producer_factory=producer_factory,
+    at_most_once_store=StubAtMostOnce(),
+)
+app.set_modules(["examples.tasks"])

clients/python/src/examples/cli.py

@@ -0,0 +1,96 @@
+import logging
+import os
+import time
+
+import click
+
+logging.basicConfig(
+    level=logging.DEBUG,
+    format="%(asctime)s %(message)s",
+    handlers=[logging.StreamHandler()],
+)
+
+
+@click.group()
+def main() -> None:
+    click.echo("Example application")
+
+
+@main.command()
+@click.option(
+    "--count",
+    help="The number of tasks to generate",
+    default=1,
+)
+def spawn(count: int = 1) -> None:
+    from examples.tasks import timed_task
+
+    click.echo(f"Spawning {count} tasks")
+    for _ in range(0, count):
+        timed_task.delay(sleep_seconds=0.1)
+    click.echo("Complete")
+
+
+@main.command()
+def scheduler() -> None:
+    from redis import StrictRedis
+
+    from examples.app import app
+    from taskbroker_client.metrics import NoOpMetricsBackend
+    from taskbroker_client.scheduler import RunStorage, ScheduleRunner, crontab
+
+    redis_host = os.getenv("REDIS_HOST") or "localhost"
+    redis_port = int(os.getenv("REDIS_PORT") or 6379)
+
+    # Ensure all task modules are loaded.
+    app.load_modules()
+
+    redis = StrictRedis(host=redis_host, port=redis_port, decode_responses=True)
+    metrics = NoOpMetricsBackend()
+    run_storage = RunStorage(metrics=metrics, redis=redis)
+    scheduler = ScheduleRunner(app, run_storage)
+
+    # Define a scheduled task
+    scheduler.add(
+        "simple-task", {"task": "examples:examples.simple_task", "schedule": crontab(minute="*/1")}
+    )
+
+    click.echo("Starting scheduler")
+    scheduler.log_startup()
+    while True:
+        sleep_time = scheduler.tick()
+        time.sleep(sleep_time)
+
+
+@main.command()
+@click.option(
+    "--rpc-host",
+    help="The address of the taskbroker this worker connects to.",
+    default="127.0.0.1:50051",
+)
+@click.option(
+    "--concurrency",
+    help="The number of child processes to start.",
+    default=2,
+)
+def worker(rpc_host: str, concurrency: int) -> None:
+    from taskbroker_client.worker import TaskWorker
+
+    click.echo("Starting worker")
+    worker = TaskWorker(
+        app_module="examples.app:app",
+        broker_hosts=[rpc_host],
+        max_child_task_count=100,
+        concurrency=concurrency,
+        child_tasks_queue_maxsize=concurrency * 2,
+        result_queue_maxsize=concurrency * 2,
+        rebalance_after=32,
+        processing_pool_name="examples",
+        process_type="forkserver",
+    )
+    exitcode = worker.start()
+    raise SystemExit(exitcode)
+
+
+if __name__ == "__main__":
+    main()

clients/python/src/examples/store.py

@@ -0,0 +1,12 @@
+from taskbroker_client.types import AtMostOnceStore
+
+
+class StubAtMostOnce(AtMostOnceStore):
+    def __init__(self) -> None:
+        self._keys: dict[str, str] = {}
+
+    def add(self, key: str, value: str, timeout: int) -> bool:
+        if key in self._keys:
+            return False
+        self._keys[key] = value
+        return True