Refactor: remove lots of repeated rebar

src/cdm_data_loaders/core/fields.py

@@ -69,7 +69,9 @@ class LoaderFileFormatEnum(StrEnum):
     }
 )
 
+
 NonEmptyStr = Annotated[str, StringConstraints(min_length=1)]
+CharStr = Annotated[str, StringConstraints(min_length=1, max_length=1)]
 
 BatchSize = Annotated[
     PositiveInt,

src/cdm_data_loaders/core/settings.py

@@ -1,14 +1,16 @@
 """Settings classes for data loading and wrangling pipelines."""
 
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Final
+from typing import TYPE_CHECKING, Annotated, Any, Final
 
 from frozendict import frozendict
-from pydantic import AliasChoices, AliasPath, computed_field
+from pydantic import AliasChoices, AliasPath, Field, computed_field
 from pydantic.fields import FieldInfo
 from pydantic_settings import BaseSettings, SettingsConfigDict, SettingsError
 
 from cdm_data_loaders.core.fields import (
+    DEFAULT_XML_FILE_GLOB,
+    FILE_GLOB,
     INPUT_DIR,
     LOADER_FILE_FORMAT,
     OUTPUT_DIR,
@@ -17,6 +19,7 @@
     BufferSize,
     DatasetName,
     DltDevMode,
+    FileGlob,
     LoaderFileFormat,
     LogConfigFile,
     LogInterval,
@@ -47,6 +50,18 @@
         USE_OUTPUT_DIR_FOR_PIPELINE_METADATA: "p",
     }
 )
+FILE_FORMAT_CLI_SHORTCUTS: frozendict[str, str] = frozendict(
+    {
+        **CLI_SHORTCUTS,
+        LOADER_FILE_FORMAT: "f",
+    }
+)
+XML_CLI_SHORTCUTS: frozendict[str, str] = frozendict(
+    {
+        **FILE_FORMAT_CLI_SHORTCUTS,
+        FILE_GLOB: "g",
+    }
+)
 
 # argparse registers -h/--help on every parser; claiming either makes the parser fail to build
 RESERVED_CLI_NAMES: Final[frozenset[str]] = frozenset({"h", "help"})
@@ -91,7 +106,7 @@
     model_config = SettingsConfigDict(**DEFAULT_SETTINGS_CONFIG_DICT)
 
     @classmethod
     def __pydantic_init_subclass__(cls, **kwargs: Any) -> None:
         """Validate the subclass's cli shortcuts and aliases."""
         super().__pydantic_init_subclass__(**kwargs)
         cls.check_aliases()
@@ -210,9 +225,21 @@
 class FileFormatIngestSettings(CtsSettings):
     """Configuration for DLT pipelines that import a particular file format."""
 
-    model_config: SettingsConfigDict = default_settings_with_shortcuts(
-        cli_shortcuts={**CLI_SHORTCUTS, LOADER_FILE_FORMAT: "f"}
-    )
+    model_config: SettingsConfigDict = default_settings_with_shortcuts(cli_shortcuts=FILE_FORMAT_CLI_SHORTCUTS)
 
     dataset_name: DatasetName
     loader_file_format: LoaderFileFormat
+
+
+class XmlIngestSettings(FileFormatIngestSettings):
+    """Configuration for DLT pipelines that import XML files."""
+
+    model_config: SettingsConfigDict = default_settings_with_shortcuts(cli_shortcuts=XML_CLI_SHORTCUTS)
+
+    file_glob: Annotated[
+        FileGlob,
+        Field(
+            default=DEFAULT_XML_FILE_GLOB,
+            description="Glob pattern for XML files",
+        ),
+    ]

src/cdm_data_loaders/pipelines/all_the_bacteria.py

@@ -16,16 +16,16 @@
 from functools import cached_property
 from logging import Logger, getLogger
 from pathlib import Path
-from typing import Annotated, Any, Final, Self
+from typing import Any, Final, Self
 
 import dlt
-import dlt.sources.helpers.requests
 from dlt.common.pipeline import LoadInfo
 from dlt.extract.items import DataItemWithMeta
 from dlt.sources.helpers.rest_client.client import RESTClient
-from pydantic import Field, StringConstraints, computed_field, model_validator
+from pydantic import Field, computed_field, model_validator
 from pydantic_settings import SettingsConfigDict
 
+from cdm_data_loaders.core.fields import NonEmptyStr
 from cdm_data_loaders.core.settings import (
     CLI_SHORTCUTS,
     CtsSettings,
@@ -39,8 +39,6 @@
 
 logger: Logger = getLogger(__name__)
 
-NonEmptyStr = Annotated[str, StringConstraints(min_length=1)]
-
 
 DATASET_NAME = "all_the_bacteria"
 ALL_FILES_TSV_FILE_ID = "R6gcp"

src/cdm_data_loaders/pipelines/core.py

@@ -11,13 +11,20 @@
 import dlt
 from dlt.common.pipeline import LoadInfo
 from dlt.common.runtime.slack import send_slack_message
+from dlt.common.schema.typing import TSchemaContractDict
 from dlt.extract import DltResource
 from dlt.pipeline.exceptions import PipelineStepFailed
+from dlt.sources.filesystem import filesystem
 from pydantic import ValidationError
 from pydantic_settings import CliApp, SettingsError
 
-from cdm_data_loaders.core.fields import JSONL
-from cdm_data_loaders.core.settings import CtsSettings, LoggerSettings
+from cdm_data_loaders.core.fields import DEFAULT_XML_FILE_GLOB, JSONL
+from cdm_data_loaders.core.settings import (
+    CtsSettings,
+    FileFormatIngestSettings,
+    LocalInputOutputSettings,
+    LoggerSettings,
+)
 from cdm_data_loaders.utils.cdm_logger import init_logger
 
 WEBHOOK_NOT_CONFIGURED: Final[str] = "Slack webhook not configured"
@@ -27,6 +34,12 @@
 DESTINATION_NOT_CONFIGURED: Final[str] = (
     "Destination {name!r} is not configured: set destination.{name}.destination_type in the dlt config."
 )
+DEFAULT_SCHEMA_CONTRACT: Final[TSchemaContractDict] = {
+    "tables": "evolve",
+    "columns": "evolve",
+    "data_type": "discard_row",
+}
+DEFAULT_REJECTED_SUFFIX: Final[str] = "_rejected"
 
 logger: Logger = getLogger(__name__)
 
@@ -184,7 +197,19 @@
         try:
             # save the load info as JSONL in the same destination as the ingested data
             load_info_resource: DltResource = dlt.resource(
-                [load_info.asdict()], name=LOAD_INFO_TABLE_NAME, max_table_nesting=0
+                [load_info.asdict()],
+                name=LOAD_INFO_TABLE_NAME,
+                max_table_nesting=0,
+                columns={
+                    k: {"data_type": "text"}
+                    for k in [
+                        "environment",
+                        "staging_type",
+                        "staging_name",
+                        "staging_display_credentials",
+                        "dataset_name",
+                    ]
+                },
             )
             load_info_save_info: LoadInfo = pipeline.run(load_info_resource, loader_file_format=str(JSONL))
             if load_info_save_info.has_failed_jobs:
@@ -300,3 +325,61 @@
         send_slack_message_carefully(slack_hook, msg)
 
     return load_info
+
+
+def build_filesystem_resource(
+    settings: LocalInputOutputSettings,
+    reader_fn: Callable[..., Any] | DltResource,
+    *reader_args: Any,  # noqa: ANN401
+    name: str,
+    file_glob: str | None = None,
+    schema_contract: TSchemaContractDict | None = None,
+    **transformer_kwargs: Any,  # noqa: ANN401
+) -> DltResource:
+    """Build a dlt resource that pipes ``filesystem(settings.input_dir, file_glob)`` into a parallelized transformer.
+
+    :param settings: pipeline settings with ``input_dir`` (and optionally ``file_glob``)
+    :param reader_fn: transformer function or resource accepting ``items`` followed by ``*reader_args``
+    :param reader_args: positional arguments bound to the transformer after ``items``
+    :param name: name of the resulting dlt transformer resource
+    :param file_glob: glob pattern for files; defaults to ``settings.file_glob`` if present, else ``DEFAULT_XML_FILE_GLOB``
+    :param schema_contract: optional dlt schema contract hint applied to the resource
+    :param transformer_kwargs: additional keyword arguments passed to ``dlt.transformer``
+    :return: configured ``DltResource``
+    """
+    resolved_glob = file_glob if file_glob is not None else getattr(settings, "file_glob", DEFAULT_XML_FILE_GLOB)
+    files = filesystem(bucket_url=str(settings.input_dir), file_glob=resolved_glob)
+    reader = (
+        reader_fn
+        if isinstance(reader_fn, DltResource)
+        else dlt.transformer(reader_fn, name=name, parallelized=True, **transformer_kwargs)
+    )
+    resource: DltResource = files | reader(*reader_args)
+    if schema_contract is not None:
+        resource.apply_hints(schema_contract=TSchemaContractDict(**schema_contract))  # pyright: ignore[reportCallIssue]
+    return resource
+
+
+def build_file_ingest_kwargs(
+    settings: FileFormatIngestSettings,
+    pipeline_name: str,
+    *,
+    fallback_dataset_name: bool = False,
+) -> dict[str, dict[str, Any]]:
+    """Build the standard ``pipeline_kwargs`` and ``pipeline_run_kwargs`` for a file-ingest pipeline.
+
+    :param settings: file-format ingest settings with ``dataset_name`` and ``loader_file_format``
+    :param pipeline_name: name of the dlt pipeline
+    :param fallback_dataset_name: whether to fall back to ``pipeline_name`` when ``settings.dataset_name`` is None
+    :return: mapping with ``pipeline_kwargs`` and ``pipeline_run_kwargs`` suitable for ``run_pipeline``
+    """
+    dataset_name = (settings.dataset_name or pipeline_name) if fallback_dataset_name else settings.dataset_name
+    return {
+        "pipeline_kwargs": {
+            "pipeline_name": pipeline_name,
+            "dataset_name": dataset_name,
+        },
+        "pipeline_run_kwargs": {
+            "loader_file_format": str(settings.loader_file_format),
+        },
+    }

src/cdm_data_loaders/pipelines/jsonlines/base.py

@@ -8,6 +8,7 @@
 """
 
 from collections.abc import Callable, Generator
+from importlib import import_module
 from logging import Logger, getLogger
 from typing import TYPE_CHECKING, Any, Final
 
@@ -16,7 +17,11 @@
 from dlt.common.schema.typing import TSchemaContractDict
 from dlt.sources.filesystem import filesystem
 
-from cdm_data_loaders.pipelines.core import run_pipeline
+from cdm_data_loaders.pipelines.core import (
+    DEFAULT_REJECTED_SUFFIX,
+    DEFAULT_SCHEMA_CONTRACT,
+    run_pipeline,
+)
 from cdm_data_loaders.readers.jsonlines import stream_jsonl_lines
 from cdm_data_loaders.utils.buffer import ListBuffer
 
@@ -25,20 +30,35 @@
 
 logger: Logger = getLogger(__name__)
 
-SCHEMA_CONTRACT: Final[TSchemaContractDict] = {
-    "tables": "evolve",
-    "columns": "evolve",
-    "data_type": "discard_row",
-}
-
-DEFAULT_REJECTED_SUFFIX: Final[str] = "_rejected"
+SCHEMA_CONTRACT: Final[TSchemaContractDict] = DEFAULT_SCHEMA_CONTRACT
 
 REJECTED_COLUMNS: Final[tuple[str, ...]] = ("source_file", "line_no", "raw_record")
 
 ParseErrorResolver = Callable[[dict[str, Any]], str | None]
 ResourceFactory = Callable[[str, dict[str, Any], Any], "DltResource"]
 
 
+def load_registry_dict(module_path: str, attribute_name: str, type_hint: str = "") -> dict[str, Any]:
+    """Import ``module_path`` and return its non-empty dictionary attribute ``attribute_name``.
+
+    :param module_path: dotted import path of the registry module
+    :type module_path: str
+    :param attribute_name: name of the dictionary attribute on the module
+    :type attribute_name: str
+    :param type_hint: optional suffix included in the ValueError message when missing or empty
+    :type type_hint: str
+    :return: the registry dictionary
+    :rtype: dict[str, Any]
+    :raises ValueError: if the attribute is missing, empty, or not a dict
+    """
+    module = import_module(module_path)
+    registry = getattr(module, attribute_name, None)
+    if not isinstance(registry, dict) or not registry:
+        err_msg = f"{module_path} does not define a non-empty {attribute_name}{type_hint}"
+        raise ValueError(err_msg)
+    return registry
+
+
 def make_page_router(  # noqa: PLR0913
     table_name: str,
     buffer_size: int,

src/cdm_data_loaders/pipelines/jsonlines/extract_jsonschema_validate_pipeline.py

@@ -8,7 +8,6 @@
 """
 
 import json
-from importlib import import_module
 from logging import Logger, getLogger
 from typing import Any
 
@@ -21,6 +20,7 @@
 from cdm_data_loaders.pipelines.core import PipelineCli
 from cdm_data_loaders.pipelines.jsonlines.base import (
     ParseErrorResolver,
+    load_registry_dict,
     run_registry_ingest,
 )
 from cdm_data_loaders.pipelines.jsonlines.base import (
@@ -44,11 +44,7 @@
     :rtype: dict[str, dict[str, Any]]
     :raises ValueError: if ENTITY_SCHEMAS is missing, empty, or not a dict
     """
-    module = import_module(module_path)
-    registry = getattr(module, "ENTITY_SCHEMAS", None)
-    if not isinstance(registry, dict) or not registry:
-        err_msg = f"{module_path} does not define a non-empty ENTITY_SCHEMAS"
-        raise ValueError(err_msg)
+    registry = load_registry_dict(module_path, "ENTITY_SCHEMAS")
     errors: dict[str, str] = {}
     for table_name, schema in registry.items():
         if not isinstance(table_name, str) or not table_name.strip():

src/cdm_data_loaders/pipelines/jsonlines/extract_pipeline.py

@@ -15,9 +15,13 @@
 from dlt.common.pipeline import LoadInfo
 from dlt.common.storages.fsspec_filesystem import FileItemDict
 from dlt.common.typing import TDataItems
-from dlt.sources.filesystem import filesystem
 
-from cdm_data_loaders.pipelines.core import PipelineCli, run_pipeline
+from cdm_data_loaders.pipelines.core import (
+    PipelineCli,
+    build_file_ingest_kwargs,
+    build_filesystem_resource,
+    run_pipeline,
+)
 from cdm_data_loaders.pipelines.jsonlines.settings import (
     EXTRACT_PIPELINE_NAME,
     JsonlExtractSettings,
@@ -47,20 +51,12 @@
     :return: load information for the pipeline
     :rtype: LoadInfo | None
     """
-    files = filesystem(bucket_url=str(settings.input_dir), file_glob=settings.file_glob)
-
-    jsonl_resource = files | jsonl_reader(settings)
+    jsonl_resource = build_filesystem_resource(settings, jsonl_reader, settings, name="jsonl_reader")
 
     return run_pipeline(
         settings=settings,
         resource=jsonl_resource,
-        pipeline_kwargs={
-            "pipeline_name": EXTRACT_PIPELINE_NAME,
-            "dataset_name": settings.dataset_name or EXTRACT_PIPELINE_NAME,
-        },
-        pipeline_run_kwargs={
-            "loader_file_format": str(settings.loader_file_format),
-        },
+        **build_file_ingest_kwargs(settings, EXTRACT_PIPELINE_NAME, fallback_dataset_name=True),
     )
 
 

src/cdm_data_loaders/pipelines/jsonlines/extract_pydantic_validate_pipeline.py

@@ -31,7 +31,6 @@
 """
 
 import json
-from importlib import import_module
 from typing import Any
 
 from dlt.common.pipeline import LoadInfo
@@ -40,6 +39,7 @@
 from cdm_data_loaders.pipelines.core import PipelineCli
 from cdm_data_loaders.pipelines.jsonlines.base import (
     ParseErrorResolver,
+    load_registry_dict,
     run_registry_ingest,
 )
 from cdm_data_loaders.pipelines.jsonlines.base import (
@@ -60,12 +60,7 @@
     :rtype: dict[str, type[BaseModel]]
     :raises ValueError: if ENTITY_MODELS is missing, empty, or not a dict
     """
-    module = import_module(module_path)
-    registry = getattr(module, "ENTITY_MODELS", None)
-    if not isinstance(registry, dict) or not registry:
-        err_msg = f"{module_path} does not define a non-empty ENTITY_MODELS: dict[str, type[BaseModel]]"
-        raise ValueError(err_msg)
-    return registry
+    return load_registry_dict(module_path, "ENTITY_MODELS", type_hint=": dict[str, type[BaseModel]]")
 
 
 def _make_resolve_error(model: type[BaseModel]) -> ParseErrorResolver: