UncountablePythonSDK 0.0.24__py3-none-any.whl → 0.0.131__py3-none-any.whl

uncountable/integration/cli.py

@@ -0,0 +1,142 @@
+import argparse
+
+from dateutil import tz
+from opentelemetry.trace import get_current_span
+from tabulate import tabulate
+
+from uncountable.core.environment import get_local_admin_server_port
+from uncountable.integration.queue_runner.command_server.command_client import (
+    send_job_queue_message,
+    send_list_queued_jobs_message,
+    send_retry_job_message,
+)
+from uncountable.integration.telemetry import Logger
+from uncountable.types import queued_job_t
+
+
+def register_enqueue_job_parser(
+    sub_parser_manager: argparse._SubParsersAction,
+    parents: list[argparse.ArgumentParser],
+) -> None:
+    run_parser = sub_parser_manager.add_parser(
+        "run",
+        parents=parents,
+        help="Process a job with a given host and job ID",
+        description="Process a job with a given host and job ID",
+    )
+    run_parser.add_argument("job_id", type=str, help="The ID of the job to process")
+
+    def _handle_enqueue_job(args: argparse.Namespace) -> None:
+        send_job_queue_message(
+            job_ref_name=args.job_id,
+            payload=queued_job_t.QueuedJobPayload(
+                invocation_context=queued_job_t.InvocationContextManual()
+            ),
+            host=args.host,
+            port=get_local_admin_server_port(),
+        )
+
+    run_parser.set_defaults(func=_handle_enqueue_job)
+
+
+def register_list_queued_jobs(
+    sub_parser_manager: argparse._SubParsersAction,
+    parents: list[argparse.ArgumentParser],
+) -> None:
+    list_queued_jobs_parser = sub_parser_manager.add_parser(
+        "list-queued-jobs",
+        parents=parents,
+        help="List all jobs queued on the integration server",
+        description="List all jobs queued on the integration server",
+    )
+
+    list_queued_jobs_parser.add_argument(
+        "--offset",
+        type=int,
+        default=0,
+        help="Number of jobs to skip. Should be non-negative.",
+    )
+    list_queued_jobs_parser.add_argument(
+        "--limit",
+        type=int,
+        default=100,
+        help="A number between 1 and 100 specifying the number of jobs to return in the result set.",
+    )
+
+    def _handle_list_queued_jobs(args: argparse.Namespace) -> None:
+        queued_jobs = send_list_queued_jobs_message(
+            offset=args.offset,
+            limit=args.limit,
+            host=args.host,
+            port=get_local_admin_server_port(),
+        )
+
+        headers = ["UUID", "Job Ref Name", "Attempts", "Status", "Submitted At"]
+        rows = [
+            [
+                job.uuid,
+                job.job_ref_name,
+                job.num_attempts,
+                job.status,
+                job.submitted_at.ToDatetime(tz.UTC).astimezone(tz.tzlocal()),
+            ]
+            for job in queued_jobs
+        ]
+        print(tabulate(rows, headers=headers, tablefmt="grid"))
+
+    list_queued_jobs_parser.set_defaults(func=_handle_list_queued_jobs)
+
+
+def register_retry_job_parser(
+    sub_parser_manager: argparse._SubParsersAction,
+    parents: list[argparse.ArgumentParser],
+) -> None:
+    retry_failed_jobs_parser = sub_parser_manager.add_parser(
+        "retry-job",
+        parents=parents,
+        help="Retry failed job on the integration server",
+        description="Retry failed job on the integration server",
+    )
+
+    retry_failed_jobs_parser.add_argument(
+        "job_uuid", type=str, help="The uuid of the job to retry"
+    )
+
+    def _handle_retry_job(args: argparse.Namespace) -> None:
+        send_retry_job_message(
+            job_uuid=args.job_uuid,
+            host=args.host,
+            port=get_local_admin_server_port(),
+        )
+
+    retry_failed_jobs_parser.set_defaults(func=_handle_retry_job)
+
+
+def main() -> None:
+    logger = Logger(get_current_span())
+
+    main_parser = argparse.ArgumentParser(
+        description="Execute a given integrations server command."
+    )
+
+    base_parser = argparse.ArgumentParser(add_help=False)
+    base_parser.add_argument(
+        "--host", type=str, default="localhost", nargs="?", help="The host to run on"
+    )
+
+    subparser_action = main_parser.add_subparsers(
+        dest="command",
+        required=True,
+        help="The command to execute (e.g., 'run')",
+    )
+
+    register_enqueue_job_parser(subparser_action, parents=[base_parser])
+    register_retry_job_parser(subparser_action, parents=[base_parser])
+    register_list_queued_jobs(subparser_action, parents=[base_parser])
+
+    args = main_parser.parse_args()
+    with logger.push_scope(args.command):
+        args.func(args)
+
+
+main()

uncountable/integration/construct_client.py

@@ -1,35 +1,51 @@
-import os
-from typing import assert_never
-
-from uncountable.core.client import AuthDetailsApiKey, Client
-from uncountable.integration.types import (
-    AuthRetrievalEnv,
+from uncountable.core import AuthDetailsApiKey, Client
+from uncountable.core.client import ClientConfig
+from uncountable.core.types import AuthDetailsAll, AuthDetailsOAuth
+from uncountable.integration.secret_retrieval.retrieve_secret import retrieve_secret
+from uncountable.integration.telemetry import JobLogger
+from uncountable.types import auth_retrieval_t
+from uncountable.types.job_definition_t import (
     ProfileMetadata,
 )
 
 
-def _get_env_var(name: str) -> str:
-    value = os.getenv(name)
-    if value is None:
-        raise Exception(f"environment variable {name} is missing")
-    return value
-
-
-def construct_uncountable_client(profile_meta: ProfileMetadata) -> Client:
+def _construct_auth_details(profile_meta: ProfileMetadata) -> AuthDetailsAll:
     match profile_meta.auth_retrieval:
-        case AuthRetrievalEnv():
-            api_id = _get_env_var(f"UNC_PROFILE_{profile_meta.name.upper()}_API_ID")
-            api_secret_key = _get_env_var(
-                f"UNC_PROFILE_{profile_meta.name.upper()}_API_SECRET_KEY"
+        case auth_retrieval_t.AuthRetrievalOAuth():
+            refresh_token = retrieve_secret(
+                profile_meta.auth_retrieval.refresh_token_secret,
+                profile_metadata=profile_meta,
+            )
+            return AuthDetailsOAuth(refresh_token=refresh_token)
+        case auth_retrieval_t.AuthRetrievalBasic():
+            api_id = retrieve_secret(
+                profile_meta.auth_retrieval.api_id_secret, profile_metadata=profile_meta
+            )
+            api_key = retrieve_secret(
+                profile_meta.auth_retrieval.api_key_secret,
+                profile_metadata=profile_meta,
             )
 
-            assert api_id is not None
-            assert api_secret_key is not None
+            return AuthDetailsApiKey(api_id=api_id, api_secret_key=api_key)
 
-            return Client(
-                base_url=profile_meta.base_url,
-                auth_details=AuthDetailsApiKey(
-                    api_id=api_id, api_secret_key=api_secret_key
-                ),
-            )
-    assert_never(profile_meta.auth_retrieval)
+
+def _construct_client_config(
+    profile_meta: ProfileMetadata, job_logger: JobLogger
+) -> ClientConfig | None:
+    if profile_meta.client_options is None:
+        return None
+    return ClientConfig(
+        allow_insecure_tls=profile_meta.client_options.allow_insecure_tls,
+        extra_headers=profile_meta.client_options.extra_headers,
+        logger=job_logger,
+    )
+
+
+def construct_uncountable_client(
+    profile_meta: ProfileMetadata, logger: JobLogger
+) -> Client:
+    return Client(
+        base_url=profile_meta.base_url,
+        auth_details=_construct_auth_details(profile_meta),
+        config=_construct_client_config(profile_meta, logger),
+    )

uncountable/integration/cron.py

@@ -1,10 +1,12 @@
 from dataclasses import dataclass
 
 from pkgs.argument_parser import CachedParser
-from uncountable.integration.construct_client import construct_uncountable_client
-from uncountable.integration.executors.script_executor import resolve_script_executor
-from uncountable.integration.job import CronJobArguments
-from uncountable.integration.types import JobDefinition, ProfileMetadata
+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
 
 
 @dataclass
@@ -18,11 +20,10 @@ cron_args_parser = CachedParser(CronJobArgs)
 
 def cron_job_executor(**kwargs: dict) -> None:
     args_passed = cron_args_parser.parse_storage(kwargs)
-    args = CronJobArguments(
-        job_definition=args_passed.definition,
-        client=construct_uncountable_client(profile_meta=args_passed.profile_metadata),
+    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(),
     )
-
-    job = resolve_script_executor(args_passed.definition.executor, args_passed.profile_metadata)
-
-    job.run(args)

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,41 +1,11 @@
-import os
-from importlib import resources
-
+from uncountable.integration.db.connect import IntegrationDBService, create_db_engine
+from uncountable.integration.scan_profiles import load_profiles
 from uncountable.integration.server import IntegrationServer
-from uncountable.integration.types import ProfileDefinition
-from pkgs.argument_parser import CachedParser
-from uncountable.integration.db.connect import create_db_engine
-
-
-profile_parser = CachedParser(ProfileDefinition)
 
 
 def main() -> None:
-    profiles_module = os.environ["UNC_PROFILES_MODULE"]
-    with IntegrationServer(create_db_engine()) as server:
-        # TODO: Loop through all job spec yaml files and call server.add_job
-        profiles = [
-            entry
-            for entry in resources.files(profiles_module).iterdir()
-            if entry.is_dir()
-        ]
-        for profile in profiles:
-            profile_name = profile.name
-            try:
-                profile = profile_parser.parse_yaml_resource(
-                    package=".".join([profiles_module, profile_name]),
-                    resource="profile.yaml",
-                )
-            except FileNotFoundError as e:
-                print("WARN: profile.yaml not found", e)
-                continue
-            server.register_profile(
-                profile_name=profile_name,
-                base_url=profile.base_url,
-                auth_retrieval=profile.auth_retrieval,
-                jobs=profile.jobs,
-            )
-
+    with IntegrationServer(create_db_engine(IntegrationDBService.CRON)) as server:
+        server.register_jobs(load_profiles())
         server.serve_forever()
 
 

uncountable/integration/executors/executors.py

@@ -0,0 +1,147 @@
+from typing import assert_never
+
+from uncountable.core.client import Client
+from uncountable.integration.executors.generic_upload_executor import GenericUploadJob
+from uncountable.integration.executors.script_executor import resolve_script_executor
+from uncountable.integration.job import Job, JobArguments
+from uncountable.types import (
+    async_jobs_t,
+    entity_t,
+    field_values_t,
+    identifier_t,
+    integration_server_t,
+    job_definition_t,
+    transition_entity_phase_t,
+)
+
+
+def resolve_executor(
+    job_executor: job_definition_t.JobExecutor,
+    profile_metadata: job_definition_t.ProfileMetadata,
+) -> Job:
+    match job_executor:
+        case job_definition_t.JobExecutorScript():
+            return resolve_script_executor(
+                job_executor, profile_metadata=profile_metadata
+            )
+        case job_definition_t.JobExecutorGenericUpload():
+            return GenericUploadJob(
+                remote_directories=job_executor.remote_directories,
+                upload_strategy=job_executor.upload_strategy,
+                data_source=job_executor.data_source,
+            )
+    assert_never(job_executor)
+
+
+def _create_run_entity(
+    *,
+    client: Client,
+    logging_settings: job_definition_t.JobLoggingSettings,
+    job_uuid: str,
+) -> entity_t.Entity:
+    run_entity = client.create_entity(
+        entity_type=entity_t.EntityType.ASYNC_JOB,
+        definition_key=identifier_t.IdentifierKeyRefName(
+            ref_name="unc_integration_server_run_definition"
+        ),
+        field_values=[
+            field_values_t.FieldRefNameValue(
+                field_ref_name=async_jobs_t.ASYNC_JOB_TYPE_FIELD_REF_NAME,
+                value=async_jobs_t.AsyncJobType.INTEGRATION_SERVER_RUN,
+            ),
+            field_values_t.FieldRefNameValue(
+                field_ref_name=async_jobs_t.ASYNC_JOB_STATUS_FIELD_REF_NAME,
+                value=async_jobs_t.AsyncJobStatus.IN_PROGRESS,
+            ),
+            field_values_t.FieldRefNameValue(
+                field_ref_name=integration_server_t.INTEGRATION_SERVER_RUN_UUID_FIELD_REF_NAME,
+                value=job_uuid,
+            ),
+        ],
+    ).entity
+    client.transition_entity_phase(
+        entity=run_entity,
+        transition=transition_entity_phase_t.TransitionIdentifierPhases(
+            phase_from_key=identifier_t.IdentifierKeyRefName(
+                ref_name="unc_integration_server_run__queued"
+            ),
+            phase_to_key=identifier_t.IdentifierKeyRefName(
+                ref_name="unc_integration_server_run__started"
+            ),
+        ),
+    )
+    if logging_settings.share_with_user_groups is not None:
+        client.grant_entity_permissions(
+            entity_type=entity_t.EntityType.ASYNC_JOB,
+            entity_key=identifier_t.IdentifierKeyId(id=run_entity.id),
+            permission_types=[
+                entity_t.EntityPermissionType.READ,
+                entity_t.EntityPermissionType.WRITE,
+            ],
+            user_group_keys=logging_settings.share_with_user_groups,
+        )
+    return run_entity
+
+
+def execute_job(
+    *,
+    job_definition: job_definition_t.JobDefinition,
+    profile_metadata: job_definition_t.ProfileMetadata,
+    args: JobArguments,
+) -> job_definition_t.JobResult:
+    with args.logger.push_scope(job_definition.name) as job_logger:
+        job = resolve_executor(job_definition.executor, profile_metadata)
+
+        job_logger.log_info("running job")
+
+        run_entity: entity_t.Entity | None = None
+        try:
+            if (
+                job_definition.logging_settings is not None
+                and job_definition.logging_settings.enabled
+            ):
+                run_entity = _create_run_entity(
+                    client=args.client,
+                    logging_settings=job_definition.logging_settings,
+                    job_uuid=args.job_uuid,
+                )
+            result = job.run_outer(args=args)
+        except Exception as e:
+            job_logger.log_exception(e)
+            if run_entity is not None:
+                args.client.set_values(
+                    entity=run_entity,
+                    values=[
+                        field_values_t.ArgumentValueRefName(
+                            field_ref_name=async_jobs_t.ASYNC_JOB_STATUS_FIELD_REF_NAME,
+                            value=async_jobs_t.AsyncJobStatus.ERROR,
+                        ),
+                    ],
+                )
+            return job_definition_t.JobResult(success=False)
+
+        if args.batch_processor.current_queue_size() != 0:
+            args.batch_processor.send()
+
+        submitted_batch_job_ids = args.batch_processor.get_submitted_job_ids()
+        job_logger.log_info(
+            "completed job",
+            attributes={
+                "submitted_batch_job_ids": submitted_batch_job_ids,
+                "success": result.success,
+            },
+        )
+        if run_entity is not None:
+            args.client.set_values(
+                entity=run_entity,
+                values=[
+                    field_values_t.ArgumentValueRefName(
+                        field_ref_name=async_jobs_t.ASYNC_JOB_STATUS_FIELD_REF_NAME,
+                        value=async_jobs_t.AsyncJobStatus.COMPLETED
+                        if result.success
+                        else async_jobs_t.AsyncJobStatus.ERROR,
+                    ),
+                ],
+            )
+
+        return result