UncountablePythonSDK 0.0.40__py3-none-any.whl → 0.0.42__py3-none-any.whl

pkgs/type_spec/config.py

@@ -1,10 +1,9 @@
 import os
 from collections.abc import Callable, Mapping
 from dataclasses import dataclass
-from decimal import Decimal
 from typing import Self, TypeVar
 
-import yaml
+from pkgs.serialization import yaml
 
 ConfigValueType = str | None | Mapping[str, str | None] | list[str]
 
@@ -103,17 +102,6 @@ def _parse_language(config_class: type[_T], raw_value: ConfigValueType) -> _T:
     return config_class(**raw_value)
 
 
-def _decimal_constructor(loader, node):  # type:ignore
-    value = loader.construct_scalar(node)
-    return Decimal(value)
-
-
-# A semi-acceptable patch to force a number to be parsed as a decimal, since pyyaml
-# parses them as lossy floats otherwise. Though a bit ugly, at least this way we have
-# support for decimal constants
-yaml.SafeLoader.add_constructor("!decimal", _decimal_constructor)
-
-
 def parse_yaml_config(config_file: str) -> Config:
     with open(config_file, encoding="utf-8") as input:
         raw_config: dict[str, ConfigValueType] = yaml.safe_load(input)

pkgs/type_spec/emit_open_api.py

@@ -9,7 +9,8 @@ import json
 import re
 from typing import Collection, cast
 
-import yaml
+from pkgs.serialization import yaml
+from pkgs.serialization_util.serialization_helpers import serialize_for_api
 
 from . import builder, util
 from .builder import EndpointGuideKey, RootGuideKey
@@ -169,7 +170,7 @@ def _serialize_global_context(ctx: EmitOpenAPIGlobalContext) -> str:
         oa_paths[path.path] = {"$ref": path.ref}
     oa_root["paths"] = oa_paths
 
-    return yaml.dump(oa_root, sort_keys=False)
+    return yaml.dumps(oa_root, sort_keys=False)
 
 
 def _is_empty_object_type(typ: OpenAPIType) -> bool:
@@ -413,7 +414,7 @@ def _emit_namespace(
 
     path = f"{config.types_output}/common/{'/'.join(namespace.path)}.yaml"
     oa_namespace = {"components": oa_components}
-    _rewrite_with_notice(path, yaml.dump(oa_namespace, sort_keys=False))
+    _rewrite_with_notice(path, yaml.dumps(oa_namespace, sort_keys=False))
 
 
 def _emit_type(
@@ -568,6 +569,7 @@ def _emit_endpoint(
         description = f"**[External API-Endpoint]** <br/> {description}"
 
     path_cutoff = min(3, len(namespace.path) - 1)
+
     ctx.endpoint = EmitOpenAPIEndpoint(
         method=namespace.endpoint.method.lower(),
         tags=[tag_name],
@@ -580,8 +582,8 @@ def _emit_endpoint(
                 ref_name=f"ex_{i}",
                 summary=example.summary,
                 description=example.description,
-                arguments=example.arguments,
-                data=example.data,
+                arguments=serialize_for_api(example.arguments),
+                data=serialize_for_api(example.data),
             )
             for i, example in enumerate(endpoint_examples)
         ],

pkgs/type_spec/emit_open_api_util.py

@@ -8,6 +8,8 @@ from collections import defaultdict
 from dataclasses import dataclass, field
 from typing import TypeAlias
 
+from pkgs.serialization_util.serialization_helpers import JsonValue
+
 from . import builder
 from .open_api_util import OpenAPIType
 
@@ -68,8 +70,8 @@ class EmitOpenAPIEndpointExample:
     ref_name: str
     summary: str
     description: str
-    arguments: dict[str, object]
-    data: dict[str, object]
+    arguments: dict[str, JsonValue]
+    data: dict[str, JsonValue]
 
 
 EmitOpenAPIStabilityLevel = builder.StabilityLevel

pkgs/type_spec/emit_python.py

@@ -1,6 +1,6 @@
+import dataclasses
 import io
 import os
-from dataclasses import dataclass, field
 from decimal import Decimal
 from typing import Any, Optional
 
@@ -35,11 +35,11 @@ QUEUED_BATCH_REQUEST_STYPE = builder.SpecTypeDefnObject(
 )
 
 
-@dataclass(kw_only=True)
+@dataclasses.dataclass(kw_only=True)
 class TrackingContext:
     namespace: Optional[builder.SpecNamespace] = None
-    namespaces: set[builder.SpecNamespace] = field(default_factory=set)
-    names: set[str] = field(default_factory=set)
+    namespaces: set[builder.SpecNamespace] = dataclasses.field(default_factory=set)
+    names: set[str] = dataclasses.field(default_factory=set)
 
     use_enum: bool = False
     use_serial_string_enum: bool = False
@@ -50,7 +50,7 @@ class TrackingContext:
     use_opaque_key: bool = False
 
 
-@dataclass(kw_only=True)
+@dataclasses.dataclass(kw_only=True)
 class Context(TrackingContext):
     out: io.StringIO
     namespace: builder.SpecNamespace
@@ -221,7 +221,7 @@ def _emit_types_imports(*, out: io.StringIO, ctx: Context) -> None:
     if ctx.use_enum:
         out.write("from pkgs.strenum_compat import StrEnum\n")
     if ctx.use_dataclass:
-        out.write("from dataclasses import dataclass\n")
+        out.write("import dataclasses\n")
     if ctx.use_serial_class:
         out.write("from pkgs.serialization import serial_class\n")
     if ctx.use_serial_union:
@@ -652,7 +652,7 @@ def _emit_string_enum(ctx: Context, stype: builder.SpecTypeDefnStringEnum) -> No
         )
 
 
-@dataclass
+@dataclasses.dataclass
 class EmittedPropertiesMetadata:
     unconverted_keys: set[str]
     unconverted_values: set[str]
@@ -719,7 +719,6 @@ def _emit_properties(
                 default = "None"
             elif prop.has_default:
                 default = _emit_value(ctx, prop.spec_type, prop.default)
-
             class_out.write(f"{INDENT * num_indent}{py_name}: {ref_type}")
             if default:
                 class_out.write(f" = {default}")
@@ -833,7 +832,7 @@ def _emit_type(ctx: Context, stype: builder.SpecType) -> None:
 
         ctx.out.write(")\n")
 
-    dataclass = "@dataclass"
+    dataclass = "@dataclasses.dataclass"
     dc_args = []
     if stype.is_kw_only():
         dc_args.append("kw_only=True")
@@ -1159,7 +1158,7 @@ def _emit_api_argument_lookup(
     for endpoint_root in builder.api_endpoints:
         routes_output = config.routes_output[endpoint_root]
 
-        imports = ["import typing", "from dataclasses import dataclass"]
+        imports = ["import typing", "import dataclasses"]
         mappings = []
         for namespace in sorted(
             builder.namespaces.values(),
@@ -1212,13 +1211,13 @@ AT = typing.TypeVar("AT")
 DT = typing.TypeVar("DT")
 
 
-@dataclass(kw_only=True, frozen=True)
+@dataclasses.dataclass(kw_only=True, frozen=True)
 class ApiEndpointKey:
     method: str
     route: str
 
 
-@dataclass(kw_only=True)
+@dataclasses.dataclass(kw_only=True)
 class ApiEndpointSpec(typing.Generic[AT, DT]):
     route: str
     arguments_type: type[AT]
@@ -1246,7 +1245,7 @@ __all__ = ["{API_ARGUMENTS_NAME}"]
 CLIENT_CLASS_FILENAME = "client_base"
 CLIENT_CLASS_IMPORTS = [
     "from abc import ABC, abstractmethod",
-    "from dataclasses import dataclass",
+    "import dataclasses",
 ]
 ASYNC_BATCH_PROCESSOR_FILENAME = "async_batch_processor"
 ASYNC_BATCH_PROCESSOR_IMPORTS = [
@@ -1339,7 +1338,7 @@ def _emit_client_class(
 DT = typing.TypeVar("DT")
 
 
-@dataclass(kw_only=True)
+@dataclasses.dataclass(kw_only=True)
 class APIRequest:
     method: str
     endpoint: str

pkgs/type_spec/load_types.py

@@ -3,9 +3,10 @@ from collections.abc import Callable
 from io import StringIO
 from typing import Optional
 
-import yaml
 from shelljob import fs
 
+from pkgs.serialization import yaml
+
 from .builder import SpecBuilder
 from .config import Config
 

pkgs/type_spec/value_spec/__main__.py

@@ -20,9 +20,9 @@ import sys
 from typing import TypeVar, cast
 
 import regex as re
-import yaml
 
 from main.base.types import base_t, value_spec_t
+from pkgs.serialization import yaml
 
 from ..util import parse_type_str, rewrite_file
 from .convert_type import convert_to_value_spec_type

uncountable/core/file_upload.py

@@ -16,6 +16,7 @@ _CHUNK_SIZE = 5 * 1024 * 1024  # s3 requires 5MiB minimum
 
 class FileUploadType(StrEnum):
     MEDIA_FILE_UPLOAD = "MEDIA_FILE_UPLOAD"
+    DATA_FILE_UPLOAD = "DATA_FILE_UPLOAD"
 
 
 @dataclass(kw_only=True)
@@ -26,7 +27,14 @@ class MediaFileUpload:
     type: Literal[FileUploadType.MEDIA_FILE_UPLOAD] = FileUploadType.MEDIA_FILE_UPLOAD
 
 
-FileUpload = MediaFileUpload
+@dataclass(kw_only=True)
+class DataFileUpload:
+    data: BytesIO
+    name: str
+    type: Literal[FileUploadType.DATA_FILE_UPLOAD] = FileUploadType.DATA_FILE_UPLOAD
+
+
+FileUpload = MediaFileUpload | DataFileUpload
 
 
 @dataclass(kw_only=True)
@@ -37,12 +45,14 @@ class FileBytes:
 
 @contextmanager
 def file_upload_data(file_upload: FileUpload) -> Generator[FileBytes, None, None]:
-    match file_upload.type:
-        case FileUploadType.MEDIA_FILE_UPLOAD:
+    match file_upload:
+        case MediaFileUpload():
             with open(file_upload.path, "rb") as f:
                 yield FileBytes(
                     name=Path(file_upload.path).name, bytes_data=BytesIO(f.read())
                 )
+        case DataFileUpload():
+            yield FileBytes(name=file_upload.name, bytes_data=file_upload.data)
 
 
 @dataclass(kw_only=True)

uncountable/integration/construct_client.py

@@ -2,7 +2,7 @@ import os
 from typing import assert_never
 
 from uncountable.core import AuthDetailsApiKey, Client
-from uncountable.integration.types import (
+from uncountable.types.job_definition_t import (
     AuthRetrievalEnv,
     ProfileMetadata,
 )

uncountable/integration/cron.py

@@ -2,9 +2,9 @@ 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.integration.executors.executors import resolve_executor
+from uncountable.integration.job import CronJobArguments, JobLogger
+from uncountable.types.job_definition_t import JobDefinition, ProfileMetadata
 
 
 @dataclass
@@ -21,10 +21,15 @@ def cron_job_executor(**kwargs: dict) -> None:
     args = CronJobArguments(
         job_definition=args_passed.definition,
         client=construct_uncountable_client(profile_meta=args_passed.profile_metadata),
+        profile_metadata=args_passed.profile_metadata,
+        logger=JobLogger(
+            profile_metadata=args_passed.profile_metadata,
+            job_definition=args_passed.definition,
+        ),
     )
 
-    job = resolve_script_executor(
-        args_passed.definition.executor, args_passed.profile_metadata
-    )
+    job = resolve_executor(args_passed.definition.executor, args_passed.profile_metadata)
+
+    print(f"running job {args_passed.definition.name}")
 
     job.run(args=args)

uncountable/integration/entrypoint.py

@@ -4,15 +4,14 @@ from importlib import resources
 from pkgs.argument_parser import CachedParser
 from uncountable.integration.db.connect import create_db_engine
 from uncountable.integration.server import IntegrationServer
-from uncountable.integration.types import ProfileDefinition
+from uncountable.types.job_definition_t import ProfileDefinition
 
 profile_parser = CachedParser(ProfileDefinition)
 
 
-def main() -> None:
+def main(blocking: bool) -> 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()
@@ -26,7 +25,7 @@ def main() -> None:
                     resource="profile.yaml",
                 )
             except FileNotFoundError as e:
-                print("WARN: profile.yaml not found", e)
+                print(f"WARN: profile.yaml not found for {profile_name}", e)
                 continue
             server.register_profile(
                 profile_name=profile_name,
@@ -35,8 +34,8 @@ def main() -> None:
                 jobs=profile.jobs,
             )
 
-        server.serve_forever()
+        if blocking:
+            server.serve_forever()
 
 
-if __name__ == "__main__":
-    main()
+main(__name__ == "__main__")

uncountable/integration/executors/executors.py

@@ -0,0 +1,24 @@
+from typing import assert_never
+
+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
+from uncountable.types import job_definition_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)

uncountable/integration/executors/generic_upload_executor.py

@@ -0,0 +1,245 @@
+import io
+import os
+import re
+from datetime import datetime, timezone
+
+import paramiko
+
+from pkgs.filesystem_utils import (
+    FileObjectData,
+    FileSystemFileReference,
+    FileSystemObject,
+    FileSystemSFTPConfig,
+    FileTransfer,
+    SFTPSession,
+)
+from pkgs.filesystem_utils.filesystem_session import FileSystemSession
+from uncountable.core.async_batch import AsyncBatchProcessor
+from uncountable.core.file_upload import DataFileUpload, FileUpload
+from uncountable.integration.job import Job, JobArguments, JobLogger, JobResult
+from uncountable.integration.secret_retrieval import retrieve_secret
+from uncountable.types.generic_upload_t import (
+    GenericRemoteDirectoryScope,
+    GenericUploadStrategy,
+)
+from uncountable.types.job_definition_t import (
+    GenericUploadDataSource,
+    GenericUploadDataSourceSFTP,
+)
+
+
+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):
+    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
+
+    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
+                )
+                private_key = paramiko.AgentKey.from_private_key(io.StringIO(pem_secret))
+                sftp_config = FileSystemSFTPConfig(
+                    ip=self.data_source.host,
+                    username=self.data_source.username,
+                    pem_key=private_key,
+                    pem_path=None,
+                )
+                return SFTPSession(sftp_config=sftp_config)
+
+    def run(self, args: JobArguments) -> JobResult:
+        client = args.client
+        logger = args.logger
+
+        batch_executor = AsyncBatchProcessor(client=client)
+        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]
+
+        if self.upload_strategy.parse_files_individually:
+            for file_id in file_ids:
+                batch_executor.invoke_uploader(
+                    file_ids=[file_id],
+                    uploader_key=self.upload_strategy.uploader_key,
+                    material_family_keys=self.upload_strategy.material_family_keys,
+                )
+        else:
+            batch_executor.invoke_uploader(
+                file_ids=file_ids,
+                uploader_key=self.upload_strategy.uploader_key,
+                material_family_keys=self.upload_strategy.material_family_keys,
+            )
+
+        batch_executor.send()
+
+        return JobResult(success=True)

uncountable/integration/executors/script_executor.py

@@ -3,7 +3,7 @@ import inspect
 import os
 
 from uncountable.integration.job import Job
-from uncountable.integration.types import JobExecutorScript, ProfileMetadata
+from uncountable.types.job_definition_t import JobExecutorScript, ProfileMetadata
 
 
 def resolve_script_executor(

uncountable/integration/job.py

@@ -2,13 +2,30 @@ from abc import ABC, abstractmethod
 from dataclasses import dataclass
 
 from uncountable.core.client import Client
-from uncountable.integration.types import JobDefinition
+from uncountable.types.job_definition_t import JobDefinition, ProfileMetadata
+
+
+class JobLogger:
+    def __init__(
+        self, *, profile_metadata: ProfileMetadata, job_definition: JobDefinition
+    ) -> None:
+        self.profile_metadata = profile_metadata
+        self.job_definition = job_definition
+
+    def log_info(self, *log_objects: object) -> None:
+        # IMPROVE: log a json message with context that can be parsed by OT
+        print(
+            f"[{self.job_definition.id}] in profile ({self.profile_metadata.name}): ",
+            *log_objects,
+        )
 
 
 @dataclass
 class JobArgumentsBase:
     job_definition: JobDefinition
+    profile_metadata: ProfileMetadata
     client: Client
+    logger: JobLogger
 
 
 @dataclass
@@ -37,6 +54,6 @@ class CronJob(Job):
     def run(self, args: CronJobArguments) -> JobResult: ...
 
 
-def register_job(cls: Job) -> Job:
+def register_job(cls: type[Job]) -> type[Job]:
     cls._unc_job_registered = True
     return cls

uncountable/integration/secret_retrieval/__init__.py

@@ -0,0 +1,3 @@
+from .retrieve_secret import retrieve_secret
+
+__all__: list[str] = ["retrieve_secret"]