UncountablePythonSDK 0.0.7__py3-none-any.whl → 0.0.92__py3-none-any.whl

uncountable/integration/construct_client.py

@@ -0,0 +1,51 @@
+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 _construct_auth_details(profile_meta: ProfileMetadata) -> AuthDetailsAll:
+    match profile_meta.auth_retrieval:
+        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,
+            )
+
+            return AuthDetailsApiKey(api_id=api_id, api_secret_key=api_key)
+
+
+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

@@ -0,0 +1,29 @@
+from dataclasses import dataclass
+
+from pkgs.argument_parser import CachedParser
+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
+class CronJobArgs:
+    definition: JobDefinition
+    profile_metadata: ProfileMetadata
+
+
+cron_args_parser = CachedParser(CronJobArgs)
+
+
+def cron_job_executor(**kwargs: dict) -> None:
+    args_passed = cron_args_parser.parse_storage(kwargs)
+    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

@@ -0,0 +1,18 @@
+import os
+from enum import StrEnum
+
+from sqlalchemy import create_engine
+from sqlalchemy.engine.base import Engine
+
+
+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

@@ -0,0 +1,13 @@
+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() -> None:
+    with IntegrationServer(create_db_engine(IntegrationDBService.CRON)) as server:
+        server.register_jobs(load_profiles())
+        server.serve_forever()
+
+
+if __name__ == "__main__":
+    main()

uncountable/integration/executors/executors.py

@@ -0,0 +1,148 @@
+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_uuid: str,
+) -> 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=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

uncountable/integration/executors/generic_upload_executor.py

@@ -0,0 +1,284 @@
+import io
+import os
+import re
+from datetime import datetime, timezone
+
+import paramiko
+
+from pkgs.filesystem_utils import (
+    FileObjectData,
+    FileSystemFileReference,
+    FileSystemObject,
+    FileSystemS3Config,
+    FileSystemSession,
+    FileSystemSFTPConfig,
+    FileTransfer,
+    S3Session,
+    SFTPSession,
+)
+from uncountable.core.file_upload import DataFileUpload, FileUpload
+from uncountable.integration.job import Job, JobArguments
+from uncountable.integration.secret_retrieval import retrieve_secret
+from uncountable.integration.telemetry import JobLogger
+from uncountable.types.generic_upload_t import (
+    GenericRemoteDirectoryScope,
+    GenericUploadStrategy,
+)
+from uncountable.types.job_definition_t import (
+    GenericUploadDataSource,
+    GenericUploadDataSourceS3,
+    GenericUploadDataSourceSFTP,
+    JobResult,
+    S3CloudProvider,
+)
+
+
+def _filter_files_by_keyword(
+    remote_directory: GenericRemoteDirectoryScope,
+    files: list[FileObjectData],
+    logger: JobLogger,
+) -> list[FileObjectData]:
+    if remote_directory.detection_keyword is None:
+        return files
+
+    raise NotImplementedError("keyword detection not implemented yet")
+
+
+def _filter_by_filename(
+    remote_directory: GenericRemoteDirectoryScope, files: list[FileSystemObject]
+) -> list[FileSystemObject]:
+    if remote_directory.filename_regex is None:
+        return files
+
+    return [
+        file
+        for file in files
+        if file.filename is not None
+        and re.search(remote_directory.filename_regex, file.filename)
+    ]
+
+
+def _filter_by_file_extension(
+    remote_directory: GenericRemoteDirectoryScope, files: list[FileSystemObject]
+) -> list[FileSystemObject]:
+    if remote_directory.valid_file_extensions is None:
+        return files
+
+    return [
+        file
+        for file in files
+        if file.filename is not None
+        and os.path.splitext(file.filename)[-1]
+        in remote_directory.valid_file_extensions
+    ]
+
+
+def _filter_by_max_files(
+    remote_directory: GenericRemoteDirectoryScope, files: list[FileSystemObject]
+) -> list[FileSystemObject]:
+    if remote_directory.max_files is None:
+        return files
+
+    return files[: remote_directory.max_files]
+
+
+def _pull_remote_directory_data(
+    *,
+    filesystem_session: FileSystemSession,
+    remote_directory: GenericRemoteDirectoryScope,
+    logger: JobLogger,
+) -> list[FileObjectData]:
+    files_to_pull = filesystem_session.list_files(
+        dir_path=FileSystemFileReference(
+            filepath=remote_directory.src_path,
+        ),
+        recursive=remote_directory.recursive,
+    )
+    logger.log_info(
+        f"Pulled the following files {files_to_pull} from the remote directory {remote_directory}.",
+    )
+
+    files_to_pull = _filter_by_file_extension(remote_directory, files_to_pull)
+    files_to_pull = _filter_by_filename(remote_directory, files_to_pull)
+    files_to_pull = _filter_by_max_files(remote_directory, files_to_pull)
+
+    logger.log_info(
+        f"Accessing SFTP directory: {remote_directory.src_path} and pulling files: {', '.join([f.filename for f in files_to_pull if f.filename is not None])}",
+    )
+    return filesystem_session.download_files(files_to_pull)
+
+
+def _filter_downloaded_file_data(
+    remote_directory: GenericRemoteDirectoryScope,
+    pulled_file_data: list[FileObjectData],
+    logger: JobLogger,
+) -> list[FileObjectData]:
+    filtered_file_data = _filter_files_by_keyword(
+        remote_directory=remote_directory, files=pulled_file_data, logger=logger
+    )
+    return filtered_file_data
+
+
+def _move_files_post_upload(
+    *,
+    filesystem_session: FileSystemSession,
+    remote_directory_scope: GenericRemoteDirectoryScope,
+    success_file_paths: list[str],
+    failed_file_paths: list[str],
+) -> None:
+    success_file_transfers: list[FileTransfer] = []
+    appended_text = ""
+
+    if remote_directory_scope.prepend_date_on_archive:
+        appended_text = f"-{datetime.now(timezone.utc).timestamp()}"
+
+    for file_path in success_file_paths:
+        filename = os.path.split(file_path)[-1]
+        root, extension = os.path.splitext(filename)
+        new_filename = f"{root}{appended_text}{extension}"
+        # format is source, dest in the tuple
+        success_file_transfers.append((
+            FileSystemFileReference(file_path),
+            FileSystemFileReference(
+                os.path.join(
+                    remote_directory_scope.success_archive_path,
+                    new_filename,
+                )
+            ),
+        ))
+
+    failed_file_transfers: list[FileTransfer] = []
+    for file_path in failed_file_paths:
+        filename = os.path.split(file_path)[-1]
+        root, extension = os.path.splitext(filename)
+        new_filename = f"{root}{appended_text}{extension}"
+        failed_file_transfers.append((
+            FileSystemFileReference(file_path),
+            FileSystemFileReference(
+                os.path.join(
+                    remote_directory_scope.failure_archive_path,
+                    new_filename,
+                )
+            ),
+        ))
+
+    filesystem_session.move_files([*success_file_transfers, *failed_file_transfers])
+
+
+class GenericUploadJob(Job[None]):
+    def __init__(
+        self,
+        data_source: GenericUploadDataSource,
+        remote_directories: list[GenericRemoteDirectoryScope],
+        upload_strategy: GenericUploadStrategy,
+    ) -> None:
+        super().__init__()
+        self.remote_directories = remote_directories
+        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():
+                pem_secret = retrieve_secret(
+                    self.data_source.pem_secret, profile_metadata=args.profile_metadata
+                )
+                pem_key = paramiko.RSAKey.from_private_key(io.StringIO(pem_secret))
+                sftp_config = FileSystemSFTPConfig(
+                    ip=self.data_source.host,
+                    username=self.data_source.username,
+                    pem_path=None,
+                    pem_key=pem_key,
+                )
+                return SFTPSession(sftp_config=sftp_config)
+            case GenericUploadDataSourceS3():
+                if self.data_source.access_key_secret is not None:
+                    secret_access_key = retrieve_secret(
+                        self.data_source.access_key_secret,
+                        profile_metadata=args.profile_metadata,
+                    )
+                else:
+                    secret_access_key = None
+
+                if self.data_source.endpoint_url is None:
+                    assert self.data_source.cloud_provider is not None, (
+                        "either cloud_provider or endpoint_url must be specified"
+                    )
+                    match self.data_source.cloud_provider:
+                        case S3CloudProvider.AWS:
+                            endpoint_url = "https://s3.amazonaws.com"
+                        case S3CloudProvider.OVH:
+                            assert self.data_source.region_name is not None, (
+                                "region_name must be specified for cloud_provider OVH"
+                            )
+                            endpoint_url = f"https://s3.{self.data_source.region_name}.cloud.ovh.net"
+                else:
+                    endpoint_url = self.data_source.endpoint_url
+
+                s3_config = FileSystemS3Config(
+                    endpoint_url=endpoint_url,
+                    bucket_name=self.data_source.bucket_name,
+                    region_name=self.data_source.region_name,
+                    access_key_id=self.data_source.access_key_id,
+                    secret_access_key=secret_access_key,
+                    session_token=None,
+                )
+
+                return S3Session(s3_config=s3_config)
+
+    def run_outer(self, args: JobArguments) -> JobResult:
+        client = args.client
+        batch_processor = args.batch_processor
+        logger = args.logger
+
+        with self._construct_filesystem_session(args) as filesystem_session:
+            files_to_upload: list[FileUpload] = []
+            for remote_directory in self.remote_directories:
+                pulled_file_data = _pull_remote_directory_data(
+                    filesystem_session=filesystem_session,
+                    remote_directory=remote_directory,
+                    logger=logger,
+                )
+                filtered_file_data = _filter_downloaded_file_data(
+                    remote_directory=remote_directory,
+                    pulled_file_data=pulled_file_data,
+                    logger=args.logger,
+                )
+                for file_data in filtered_file_data:
+                    files_to_upload.append(
+                        DataFileUpload(
+                            data=io.BytesIO(file_data.file_data),
+                            name=file_data.filename,
+                        )
+                    )
+                if not self.upload_strategy.skip_moving_files:
+                    _move_files_post_upload(
+                        filesystem_session=filesystem_session,
+                        remote_directory_scope=remote_directory,
+                        success_file_paths=[
+                            file.filepath
+                            if file.filepath is not None
+                            else file.filename
+                            for file in filtered_file_data
+                        ],
+                        # IMPROVE: use triggers/webhooks to mark failed files as failed
+                        failed_file_paths=[],
+                    )
+
+            uploaded_files = client.upload_files(file_uploads=files_to_upload)
+
+        file_ids = [file.file_id for file in uploaded_files]
+
+        for destination in self.upload_strategy.destinations:
+            for file_id in file_ids:
+                batch_processor.invoke_uploader(
+                    file_id=file_id,
+                    uploader_key=self.upload_strategy.uploader_key,
+                    destination=destination,
+                )
+
+        return JobResult(success=True)

uncountable/integration/executors/script_executor.py

@@ -0,0 +1,25 @@
+import importlib
+import inspect
+import os
+
+from uncountable.integration.job import Job
+from uncountable.types.job_definition_t import JobExecutorScript, ProfileMetadata
+
+
+def resolve_script_executor(
+    executor: JobExecutorScript, profile_metadata: ProfileMetadata
+) -> Job:
+    job_module_path = ".".join([
+        os.environ["UNC_PROFILES_MODULE"],
+        profile_metadata.name,
+        executor.import_path,
+    ])
+    job_module = importlib.import_module(job_module_path)
+    found_jobs: list[Job] = []
+    for _, job_class in inspect.getmembers(job_module, inspect.isclass):
+        if getattr(job_class, "_unc_job_registered", False):
+            found_jobs.append(job_class())
+    assert len(found_jobs) == 1, (
+        f"expected exactly one job class in {executor.import_path}, found {len(found_jobs)}"
+    )
+    return found_jobs[0]

uncountable/integration/job.py

@@ -0,0 +1,87 @@
+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 base_t, webhook_job_t
+from uncountable.types.job_definition_t import JobDefinition, JobResult, ProfileMetadata
+
+
+@dataclass(kw_only=True)
+class JobArguments:
+    job_definition: JobDefinition
+    profile_metadata: ProfileMetadata
+    client: Client
+    batch_processor: AsyncBatchProcessor
+    logger: JobLogger
+    payload: base_t.JsonValue
+
+
+# only for compatibility:
+CronJobArguments = JobArguments
+
+
+PT = typing.TypeVar("PT")
+
+
+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: 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]: ...
+
+    def run_outer(self, args: JobArguments) -> JobResult:
+        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: JobArguments, payload: WPT) -> JobResult: ...
+
+
+def register_job(cls: type[Job]) -> type[Job]:
+    cls._unc_job_registered = True
+    return cls