v0.1.12: new uniprot / uniref XML pipelines

CHANGELOG.md

@@ -1,6 +1,8 @@
 # CDM Data Loaders Changelog
 
 - [CDM Data Loaders Changelog](#cdm-data-loaders-changelog)
+    - [v0.1.12](#v0112)
+    - [v0.1.11](#v0111)
     - [v0.1.10](#v0110)
     - [v0.1.9](#v019)
     - [v0.1.8](#v018)
@@ -14,6 +16,14 @@
     - [v0.1.0](#v010)
 
 
+### v0.1.12
+
+- Dedicated UniProt and UniRef XML import pipelines using custom XSD-aware XML-to-JSON converter.
+
+### v0.1.11
+
+- Preliminary XSV, XML, and JSONLines importers
+
 ### v0.1.10
 
 - Add in the ability to configure loggers using a file

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "cdm-data-loaders"
-version = "0.1.11"
+version = "0.1.12"
 description = "Data loaders and wranglers for the CDM."
 requires-python = ">= 3.13"
 readme = "README.md"
@@ -46,7 +46,9 @@ ncbi_ftp_promote = "cdm_data_loaders.pipelines.ncbi_ftp_promote:cli"
 ncbi_ftp_sync = "cdm_data_loaders.pipelines.ncbi_ftp_download:cli"
 ncbi_rest_api = "cdm_data_loaders.pipelines.ncbi_rest_api:cli"
 pdb_manifest = "cdm_data_loaders.pipelines.pdb_manifest:cli"
+uniprot = "cdm_data_loaders.pipelines.xmltopyarrow.uniprot:cli"
 uniprot_cdm = "cdm_data_loaders.pipelines.uniprot_kb_cdm:cli"
+uniref = "cdm_data_loaders.pipelines.xmltopyarrow.uniref:cli"
 uniref_cdm = "cdm_data_loaders.pipelines.uniref_cdm:cli"
 xmltodict = "cdm_data_loaders.pipelines.xmltodict.pipeline:cli"
 xmltopyarrow = "cdm_data_loaders.pipelines.xmltopyarrow.pipeline:cli"

src/cdm_data_loaders/pipelines/xmltopyarrow/pipeline.py

@@ -42,12 +42,13 @@
     files = filesystem(bucket_url=str(settings.input_dir), file_glob=settings.file_glob)
 
     xmltopyarrow_resource: DltResource = files | reader(settings)
+    pipeline_name = str(settings.model_config.get("cli_prog_name") or PIPELINE_NAME)
 
     return run_pipeline(
         settings=settings,
         resource=xmltopyarrow_resource,
         pipeline_kwargs={
-            "pipeline_name": PIPELINE_NAME,
+            "pipeline_name": pipeline_name,
             "dataset_name": settings.dataset_name,
         },
         pipeline_run_kwargs={

src/cdm_data_loaders/pipelines/xmltopyarrow/settings.py

@@ -1,5 +1,6 @@
 """Settings class for the XML to PyArrow pipeline for the KBase CTS."""
 
+from pathlib import Path
 from typing import Annotated, Final, Self
 
 import pyarrow as pa
@@ -10,6 +11,7 @@
 from cdm_data_loaders.core.fields import (
     DEFAULT_XML_FILE_GLOB,
     FILE_GLOB,
+    DatasetName,
     FileGlob,
     NonEmptyStr,
     TableName,
@@ -23,8 +25,19 @@
     generate_pyarrow_schema,
     load_schema,
 )
+from cdm_data_loaders.schemas.uniprot.schema import SCHEMA as UNIPROT_SCHEMA
+from cdm_data_loaders.schemas.uniref.schema import SCHEMA as UNIREF_SCHEMA
 
 PIPELINE_NAME: Final[str] = "xmltopyarrow"
+UNIPROT_PIPELINE_NAME: Final[str] = "uniprot"
+UNIREF_PIPELINE_NAME: Final[str] = "uniref"
+
+ENTRY_XML_TAG: Final[str] = "entry"
+ENTRY_TABLE_NAME: Final[str] = "entry"
+
+SCHEMAS_DIR: Final[Path] = Path(__file__).resolve().parents[2] / "schemas"
+UNIPROT_XSD_FILE: Final[Path] = SCHEMAS_DIR / "uniprot" / "uniprot.xsd"
+UNIREF_XSD_FILE: Final[Path] = SCHEMAS_DIR / "uniref" / "uniref.xsd"
 
 
 class XmlToPyArrowSettings(FileFormatIngestSettings):
@@ -70,20 +83,24 @@
 
     @model_validator(mode="after")
     def check_xsd_file(self) -> Self:
-        """Check whether the schema file is present and valid; load the XSD and generate the PyArrow schema."""
+        """Check whether the schema file is present and valid; load the XSD and resolve the PyArrow schema."""
         try:
             self._xsd_schema = load_schema(self.xsd_file)
-            self._pyarrow_schema = generate_pyarrow_schema(
-                self._xsd_schema,
-                self.xml_tag,
-                max_recursion_depth=self.max_recursion_depth,
-            )
+            self._pyarrow_schema = self._resolve_pyarrow_schema()
         except Exception as e:
             err_msg = f"Could not generate PyArrow schema from XSD: {e!s}"
             raise ValueError(err_msg) from e
 
         return self
 
+    def _resolve_pyarrow_schema(self) -> pa.Schema:
+        """Return the PyArrow schema for this settings instance."""
+        return generate_pyarrow_schema(
+            self._xsd_schema,
+            self.xml_tag,
+            max_recursion_depth=self.max_recursion_depth,
+        )
+
     @property
     def xsd_schema(self) -> xmlschema.XMLSchema:
         """Return the loaded XMLSchema instance."""
@@ -93,3 +110,93 @@
     def pyarrow_schema(self) -> pa.Schema:
         """Return the generated PyArrow schema for xml_tag."""
         return self._pyarrow_schema
+
+
+class UniProtXmlToPyArrowSettings(XmlToPyArrowSettings):
+    """Settings for the UniProt XML to PyArrow ingestion pipeline."""
+
+    model_config: SettingsConfigDict = default_settings_with_shortcuts(
+        cli_prog_name=UNIPROT_PIPELINE_NAME,
+        cli_shortcuts={
+            **FileFormatIngestSettings.model_config["cli_shortcuts"],
+            FILE_GLOB: "g",
+        },
+    )
+
+    dataset_name: Annotated[
+        DatasetName,
+        Field(
+            default=UNIPROT_PIPELINE_NAME,
+            description="The name of the dataset being produced",
+        ),
+    ]
+    table_name: Annotated[
+        TableName,
+        Field(
+            default=ENTRY_TABLE_NAME,
+            description="The name of the table to export parsed data to",
+        ),
+    ]
+    xml_tag: Annotated[
+        NonEmptyStr,
+        Field(
+            default=ENTRY_XML_TAG,
+            description="XML tag to capture the contents of",
+        ),
+    ]
+    xsd_file: Annotated[
+        FilePath,
+        Field(
+            default=UNIPROT_XSD_FILE,
+            description="Path to the UniProt XSD file describing the XML documents to be imported.",
+        ),
+    ]
+
+    def _resolve_pyarrow_schema(self) -> pa.Schema:
+        """Return the precomputed UniProt PyArrow schema."""
+        return UNIPROT_SCHEMA
+
+
+class UniRefXmlToPyArrowSettings(XmlToPyArrowSettings):
+    """Settings for the UniRef XML to PyArrow ingestion pipeline."""
+
+    model_config: SettingsConfigDict = default_settings_with_shortcuts(
+        cli_prog_name=UNIREF_PIPELINE_NAME,
+        cli_shortcuts={
+            **FileFormatIngestSettings.model_config["cli_shortcuts"],
+            FILE_GLOB: "g",
+        },
+    )
+
+    dataset_name: Annotated[
+        DatasetName,
+        Field(
+            default=UNIREF_PIPELINE_NAME,
+            description="The name of the dataset being produced",
+        ),
+    ]
+    table_name: Annotated[
+        TableName,
+        Field(
+            default=ENTRY_TABLE_NAME,
+            description="The name of the table to export parsed data to",
+        ),
+    ]
+    xml_tag: Annotated[
+        NonEmptyStr,
+        Field(
+            default=ENTRY_XML_TAG,
+            description="XML tag to capture the contents of",
+        ),
+    ]
+    xsd_file: Annotated[
+        FilePath,
+        Field(
+            default=UNIREF_XSD_FILE,
+            description="Path to the UniRef XSD file describing the XML documents to be imported.",
+        ),
+    ]
+
+    def _resolve_pyarrow_schema(self) -> pa.Schema:
+        """Return the precomputed UniRef PyArrow schema."""
+        return UNIREF_SCHEMA

src/cdm_data_loaders/pipelines/xmltopyarrow/uniprot.py

@@ -0,0 +1,32 @@
+"""UniProt XML to PyArrow pipeline for the KBase CTS."""
+
+from dlt.common.pipeline import LoadInfo
+
+from cdm_data_loaders.pipelines.core import PipelineCli
+from cdm_data_loaders.pipelines.xmltopyarrow.pipeline import run_xmltopyarrow_pipeline
+from cdm_data_loaders.pipelines.xmltopyarrow.settings import (
+    UNIPROT_PIPELINE_NAME as PIPELINE_NAME,
+)
+from cdm_data_loaders.pipelines.xmltopyarrow.settings import (
+    UniProtXmlToPyArrowSettings,
+)
+
+__all__ = ["PIPELINE_NAME", "UniProtXmlToPyArrowSettings", "cli", "run_uniprot_pipeline"]
+
+
+def run_uniprot_pipeline(settings: UniProtXmlToPyArrowSettings) -> LoadInfo | None:
+    """Run the UniProt XML to PyArrow pipeline on matching files in settings.input_dir.
+
+    :param settings: pipeline configuration
+    :type  settings: UniProtXmlToPyArrowSettings
+    :return: load information for the pipeline
+    :rtype: LoadInfo | None
+    """
+    return run_xmltopyarrow_pipeline(settings)
+
+
+cli = PipelineCli(UniProtXmlToPyArrowSettings, run_uniprot_pipeline)
+
+
+if __name__ == "__main__":
+    cli()

src/cdm_data_loaders/pipelines/xmltopyarrow/uniref.py

@@ -0,0 +1,32 @@
+"""UniRef XML to PyArrow pipeline for the KBase CTS."""
+
+from dlt.common.pipeline import LoadInfo
+
+from cdm_data_loaders.pipelines.core import PipelineCli
+from cdm_data_loaders.pipelines.xmltopyarrow.pipeline import run_xmltopyarrow_pipeline
+from cdm_data_loaders.pipelines.xmltopyarrow.settings import (
+    UNIREF_PIPELINE_NAME as PIPELINE_NAME,
+)
+from cdm_data_loaders.pipelines.xmltopyarrow.settings import (
+    UniRefXmlToPyArrowSettings,
+)
+
+__all__ = ["PIPELINE_NAME", "UniRefXmlToPyArrowSettings", "cli", "run_uniref_pipeline"]
+
+
+def run_uniref_pipeline(settings: UniRefXmlToPyArrowSettings) -> LoadInfo | None:
+    """Run the UniRef XML to PyArrow pipeline on matching files in settings.input_dir.
+
+    :param settings: pipeline configuration
+    :type  settings: UniRefXmlToPyArrowSettings
+    :return: load information for the pipeline
+    :rtype: LoadInfo | None
+    """
+    return run_xmltopyarrow_pipeline(settings)
+
+
+cli = PipelineCli(UniRefXmlToPyArrowSettings, run_uniref_pipeline)
+
+
+if __name__ == "__main__":
+    cli()