sdmxlib 0.8.0__py3-none-any.whl

sdmxlib/storage/lazy.py

@@ -0,0 +1,329 @@
+"""LazyCodelist and LazyItemList — DuckDB-backed lazy code access.
+
+DuckDB adaptation of the SQLite-based LazyCodelist. Uses flat
+``label_{lang}`` columns instead of JSON, and ``ORDER BY position``
+instead of ``ORDER BY rowid``.
+"""
+
+import json
+from collections.abc import Iterator
+from typing import Any, overload
+
+import duckdb
+
+from sdmxlib.model.annotations import Annotation, Annotations
+from sdmxlib.model.codelist import Code, Codelist
+from sdmxlib.model.collections import ItemList
+from sdmxlib.model.istring import InternationalString
+from sdmxlib.model.organisation import Agency
+from sdmxlib.storage.schema import get_languages
+
+
+def _row_to_code(
+    row: tuple[Any, ...],
+    languages: list[str],
+) -> Code:
+    """Convert a DuckDB row to a Code object.
+
+    Expected column order: code_id, parent_code_id, label_{lang}...,
+    description_{lang}..., _annotations
+    """
+    n_langs = len(languages)
+    code_id = row[0]
+    parent_id = row[1]
+    label_vals = row[2 : 2 + n_langs]
+    desc_vals = row[2 + n_langs : 2 + 2 * n_langs]
+    ann_json = row[2 + 2 * n_langs]
+
+    name_dict = {lang: val for lang, val in zip(languages, label_vals, strict=True) if val is not None}
+    name = InternationalString.of(name_dict) if name_dict else InternationalString()
+
+    desc_dict = {lang: val for lang, val in zip(languages, desc_vals, strict=True) if val is not None}
+    desc = InternationalString.of(desc_dict) if desc_dict else InternationalString()
+
+    annotations: list[Annotation] = []
+    if ann_json:
+        for a in json.loads(ann_json):
+            text = InternationalString.of(a["text"]) if a.get("text") else InternationalString()
+            annotations.append(
+                Annotation(
+                    id=a.get("id"),
+                    title=a.get("title"),
+                    type=a.get("type"),
+                    url=a.get("url"),
+                    text=text,
+                )
+            )
+
+    return Code(
+        id=code_id,
+        name=name,
+        description=desc,
+        parent_id=parent_id,
+        annotations=Annotations(annotations),
+    )
+
+
+def _code_select_cols(languages: list[str]) -> str:
+    """Build the SELECT column list for code queries."""
+    label_cols = ", ".join(f"label_{lang}" for lang in languages)
+    desc_cols = ", ".join(f"description_{lang}" for lang in languages)
+    return f"code_id, parent_code_id, {label_cols}, {desc_cols}, _annotations"
+
+
+class LazyItemList:
+    """A read-only sequence of Codes backed by DuckDB queries.
+
+    Supports ``__getitem__`` by int index or string id, ``__len__``,
+    ``__iter__``, and ``__contains__``.
+    """
+
+    __slots__ = ("_conn", "_langs", "_urn")
+
+    def __init__(self, conn: duckdb.DuckDBPyConnection, codelist_urn: str) -> None:
+        self._conn = conn
+        self._urn = codelist_urn
+        self._langs = get_languages(conn)
+
+    @overload
+    def __getitem__(self, index: int) -> Code: ...
+
+    @overload
+    def __getitem__(self, index: str) -> Code: ...
+
+    @overload
+    def __getitem__(self, index: slice) -> ItemList[Code]: ...
+
+    def __getitem__(self, index: int | str | slice) -> "Code | ItemList[Code]":
+        cols = _code_select_cols(self._langs)
+        if isinstance(index, str):
+            row = self._conn.execute(
+                f"""
+                SELECT {cols} FROM codes
+                WHERE codelist_urn = $urn AND code_id = $code_id
+                """,  # noqa: S608
+                {"urn": self._urn, "code_id": index},
+            ).fetchone()
+            if row is None:
+                raise KeyError(index)
+            return _row_to_code(row, self._langs)
+        if isinstance(index, slice):
+            all_codes = list(self)
+            return ItemList(all_codes[index])
+        # int index
+        offset = index if index >= 0 else len(self) + index
+        row = self._conn.execute(
+            f"""
+            SELECT {cols} FROM codes
+            WHERE codelist_urn = $urn
+            ORDER BY position LIMIT 1 OFFSET $offset
+            """,  # noqa: S608
+            {"urn": self._urn, "offset": offset},
+        ).fetchone()
+        if row is None:
+            raise IndexError(index)
+        return _row_to_code(row, self._langs)
+
+    def __len__(self) -> int:
+        row = self._conn.execute(
+            """
+            SELECT COUNT(*) FROM codes WHERE codelist_urn = $urn
+            """,
+            {"urn": self._urn},
+        ).fetchone()
+        return row[0]  # type: ignore[index]
+
+    def __iter__(self) -> Iterator[Code]:
+        cols = _code_select_cols(self._langs)
+        rows = self._conn.execute(
+            f"""
+            SELECT {cols} FROM codes
+            WHERE codelist_urn = $urn
+            ORDER BY position
+            """,  # noqa: S608
+            {"urn": self._urn},
+        ).fetchall()
+        for row in rows:
+            yield _row_to_code(row, self._langs)
+
+    def __contains__(self, key: object) -> bool:
+        if isinstance(key, str):
+            row = self._conn.execute(
+                """
+                SELECT 1 FROM codes
+                WHERE codelist_urn = $urn AND code_id = $code_id
+                """,
+                {"urn": self._urn, "code_id": key},
+            ).fetchone()
+            return row is not None
+        if isinstance(key, Code):
+            return key.id in self  # type: ignore[operator]
+        return False
+
+    def get(self, code_id: str) -> Code | None:
+        """Look up a code by id, returning None if absent."""
+        try:
+            return self[code_id]
+        except KeyError:
+            return None
+
+    def ids(self) -> list[str]:
+        """All code ids in position order."""
+        rows = self._conn.execute(
+            """
+            SELECT code_id FROM codes
+            WHERE codelist_urn = $urn
+            ORDER BY position
+            """,
+            {"urn": self._urn},
+        ).fetchall()
+        return [r[0] for r in rows]
+
+    def filter(self, predicate: "callable") -> ItemList[Code]:  # type: ignore[type-arg]
+        """Filter codes, returning a materialized ItemList."""
+        return ItemList([c for c in self if predicate(c)])
+
+    def first_where(self, predicate: "callable") -> Code | None:  # type: ignore[type-arg]
+        """First code matching predicate, or None."""
+        for c in self:
+            if predicate(c):
+                return c
+        return None
+
+    def to_list(self) -> ItemList[Code]:
+        """Materialize all codes into an ItemList."""
+        return ItemList(list(self))
+
+
+class LazyCodelist:
+    """Codelist facade backed by DuckDB. Codes load on demand.
+
+    Duck-types as ``Codelist`` for read operations. Uses flat
+    ``label_{lang}`` columns for direct SQL access.
+    """
+
+    __slots__ = (
+        "_agency",
+        "_conn",
+        "_description",
+        "_id",
+        "_is_final",
+        "_name",
+        "_urn",
+        "_version",
+    )
+
+    def __init__(
+        self,
+        *,
+        urn: str,
+        conn: duckdb.DuckDBPyConnection,
+        id: str,  # noqa: A002
+        agency: Agency,
+        version: str,
+        name: InternationalString,
+        description: InternationalString,
+        is_final: bool,
+    ) -> None:
+        self._urn = urn
+        self._conn = conn
+        self._id = id
+        self._agency = agency
+        self._version = version
+        self._name = name
+        self._description = description
+        self._is_final = is_final
+
+    # ── MaintainableArtefact-compatible properties ────────────────────────
+
+    @property
+    def id(self) -> str:  # noqa: D102
+        return self._id
+
+    @property
+    def maintainer(self) -> Agency:  # noqa: D102
+        return self._agency
+
+    @property
+    def agency_id(self) -> str:  # noqa: D102
+        return self._agency.id
+
+    @property
+    def version(self) -> str:  # noqa: D102
+        return self._version
+
+    @property
+    def name(self) -> InternationalString:  # noqa: D102
+        return self._name
+
+    @property
+    def description(self) -> InternationalString:  # noqa: D102
+        return self._description
+
+    @property
+    def is_final(self) -> bool:  # noqa: D102
+        return self._is_final
+
+    # ── SDMX class vars ──────────────────────────────────────────────────
+
+    sdmx_package = "codelist"
+    sdmx_class = "Codelist"
+
+    # ── Lazy code access ─────────────────────────────────────────────────
+
+    @property
+    def codes(self) -> LazyItemList:
+        """Lazy sequence of codes backed by DuckDB."""
+        return LazyItemList(self._conn, self._urn)
+
+    def __len__(self) -> int:
+        return len(self.codes)
+
+    def code_map(self, lang: str = "en") -> dict[str, str]:
+        """Direct SQL: {code_id: label} without creating Code objects."""
+        from sdmxlib.storage.readers import code_map  # noqa: PLC0415
+
+        return code_map(self._conn, self._urn, lang)
+
+    def roots(self) -> ItemList[Code]:
+        """Top-level codes (no parent)."""
+        langs = get_languages(self._conn)
+        cols = _code_select_cols(langs)
+        rows = self._conn.execute(
+            f"""
+            SELECT {cols} FROM codes
+            WHERE codelist_urn = $urn AND parent_code_id IS NULL
+            ORDER BY position
+            """,  # noqa: S608
+            {"urn": self._urn},
+        ).fetchall()
+        return ItemList([_row_to_code(r, langs) for r in rows])
+
+    def children_of(self, code_id: str) -> ItemList[Code]:
+        """Direct children of a code."""
+        langs = get_languages(self._conn)
+        cols = _code_select_cols(langs)
+        rows = self._conn.execute(
+            f"""
+            SELECT {cols} FROM codes
+            WHERE codelist_urn = $urn AND parent_code_id = $parent_id
+            ORDER BY position
+            """,  # noqa: S608
+            {"urn": self._urn, "parent_id": code_id},
+        ).fetchall()
+        return ItemList([_row_to_code(r, langs) for r in rows])
+
+    def materialize(self) -> Codelist:
+        """Load all codes and return a real Codelist."""
+        return Codelist(
+            id=self._id,
+            maintainer=self._agency,
+            version=self._version,
+            name=self._name,
+            description=self._description,
+            is_final=self._is_final,
+            codes=list(self.codes),
+        )
+
+    def __repr__(self) -> str:
+        return f"LazyCodelist(id={self._id!r}, agency={self.agency_id!r}, codes=~{len(self)})"

sdmxlib/storage/readers.py

@@ -0,0 +1,310 @@
+"""Read SDMX domain objects from DuckDB normalized tables.
+
+Provides functions to reconstruct domain objects from the normalized
+schema, including lazy codelist access and direct SQL helpers like
+``code_map`` that avoid materializing Python objects.
+
+Example::
+
+    from sdmxlib.storage import schema, readers
+
+    conn = duckdb.connect("store.duckdb")
+    cl = readers.get_artefact(conn, codelist_urn)  # LazyCodelist
+    cm = readers.code_map(conn, codelist_urn, "en")  # {code_id: label}
+"""
+
+import json
+from typing import Any
+
+import duckdb
+
+from sdmxlib.model.convert import from_dict
+from sdmxlib.model.istring import InternationalString
+from sdmxlib.model.organisation import Agency
+from sdmxlib.storage.lazy import LazyCodelist
+from sdmxlib.storage.schema import get_languages
+from sdmxlib.storage.writers import _CLASS_TO_NAME
+
+# Reverse map: type_name → Python class
+_TYPE_REGISTRY: dict[str, type] = {v: k for k, v in _CLASS_TO_NAME.items()}
+
+
+# ── Code loading ────────────────────────────────────────────────────────
+
+
+def load_code_dicts(conn: duckdb.DuckDBPyConnection, codelist_urn: str) -> list[dict[str, Any]]:
+    """Load all codes for a codelist as plain dicts suitable for ``from_dict``.
+
+    Returns dicts with keys: id, parent_id, name, description, annotations.
+    """
+    languages = get_languages(conn)
+
+    label_cols = ", ".join(f"label_{lang}" for lang in languages)
+    desc_cols = ", ".join(f"description_{lang}" for lang in languages)
+
+    rows = conn.execute(
+        f"""
+        SELECT code_id, parent_code_id, {label_cols}, {desc_cols}, _annotations
+        FROM codes
+        WHERE codelist_urn = $codelist_urn
+        ORDER BY position
+        """,  # noqa: S608
+        {"codelist_urn": codelist_urn},
+    ).fetchall()
+
+    result = []
+    n_langs = len(languages)
+    for row in rows:
+        code_id = row[0]
+        parent_id = row[1]
+        label_vals = row[2 : 2 + n_langs]
+        desc_vals = row[2 + n_langs : 2 + 2 * n_langs]
+        ann_json = row[2 + 2 * n_langs]
+
+        d: dict[str, Any] = {"id": code_id}
+        if parent_id is not None:
+            d["parent_id"] = parent_id
+
+        name_dict = {lang: val for lang, val in zip(languages, label_vals, strict=True) if val is not None}
+        if name_dict:
+            d["name"] = name_dict
+
+        desc_dict = {lang: val for lang, val in zip(languages, desc_vals, strict=True) if val is not None}
+        if desc_dict:
+            d["description"] = desc_dict
+
+        if ann_json:
+            d["annotations"] = json.loads(ann_json)
+
+        result.append(d)
+    return result
+
+
+# ── Hierarchy tree loading ──────────────────────────────────────────────
+
+
+def load_hierarchy_tree(conn: duckdb.DuckDBPyConnection, hierarchy_urn: str) -> list[dict[str, Any]]:
+    """Load hierarchy nodes and reconstruct the tree as nested dicts."""
+    languages = get_languages(conn)
+
+    name_cols = ", ".join(f"name_{lang}" for lang in languages)
+    desc_cols = ", ".join(f"description_{lang}" for lang in languages)
+
+    rows = conn.execute(
+        f"""
+        SELECT node_id, parent_node_id, code_urn,
+            {name_cols}, {desc_cols}, _annotations, valid_from, valid_to
+        FROM hierarchy_nodes
+        WHERE hierarchy_urn = $hierarchy_urn
+        """,  # noqa: S608
+        {"hierarchy_urn": hierarchy_urn},
+    ).fetchall()
+
+    n_langs = len(languages)
+    nodes_by_id: dict[str, dict[str, Any]] = {}
+    children_map: dict[str | None, list[str]] = {}
+
+    for row in rows:
+        node_id = row[0]
+        parent_node_id = row[1]
+        code_urn = row[2]
+        name_vals = row[3 : 3 + n_langs]
+        desc_vals = row[3 + n_langs : 3 + 2 * n_langs]
+        ann_json = row[3 + 2 * n_langs]
+        vf = row[3 + 2 * n_langs + 1]
+        vt = row[3 + 2 * n_langs + 2]
+
+        node: dict[str, Any] = {"id": node_id}
+        if code_urn:
+            node["code"] = code_urn
+
+        name_dict = {lang: val for lang, val in zip(languages, name_vals, strict=True) if val is not None}
+        if name_dict:
+            node["name"] = name_dict
+
+        desc_dict = {lang: val for lang, val in zip(languages, desc_vals, strict=True) if val is not None}
+        if desc_dict:
+            node["description"] = desc_dict
+
+        if ann_json:
+            node["annotations"] = json.loads(ann_json)
+        if vf:
+            node["valid_from"] = vf
+        if vt:
+            node["valid_to"] = vt
+
+        node["children"] = []
+        nodes_by_id[node_id] = node
+        children_map.setdefault(parent_node_id, []).append(node_id)
+
+    # Build tree from adjacency list
+    for parent_id, child_ids in children_map.items():
+        if parent_id is not None and parent_id in nodes_by_id:
+            nodes_by_id[parent_id]["children"] = [nodes_by_id[cid] for cid in child_ids]
+
+    root_ids = children_map.get(None, [])
+    return [nodes_by_id[rid] for rid in root_ids]
+
+
+# ── Artefact reconstruction ────────────────────────────────────────────
+
+
+def get_artefact(
+    conn: duckdb.DuckDBPyConnection,
+    urn: str,
+    *,
+    eager: bool = False,
+) -> Any | None:
+    """Look up an artefact by URN string.
+
+    For Codelists, returns a ``LazyCodelist`` by default. Pass
+    ``eager=True`` to get a real ``Codelist`` with all codes loaded.
+
+    For Hierarchies, nodes are loaded from the normalized table and
+    injected into the dict before reconstruction.
+
+    Returns ``None`` if the URN is not in the store.
+    """
+    row = conn.execute(
+        """
+        SELECT artefact_type, _document
+        FROM artefacts
+        WHERE urn = $urn
+        """,
+        {"urn": urn},
+    ).fetchone()
+    if row is None:
+        return None
+
+    artefact_type, doc_str = row
+    data = json.loads(doc_str)
+
+    if artefact_type == "Codelist" and not eager:
+        return make_lazy_codelist(conn, urn, data)
+
+    cls = _TYPE_REGISTRY.get(artefact_type)
+    if cls is None:
+        return None
+
+    if artefact_type == "Codelist":
+        data["codes"] = load_code_dicts(conn, urn)
+
+    if artefact_type == "Hierarchy":
+        data["tree"] = load_hierarchy_tree(conn, urn)
+
+    return from_dict(cls, data)
+
+
+def get_all(
+    conn: duckdb.DuckDBPyConnection,
+    artefact_type: type,
+    *,
+    eager: bool = False,
+) -> list[Any]:
+    """Return all artefacts of a given type.
+
+    For Codelists, returns ``LazyCodelist`` instances by default unless
+    ``eager=True``.
+    """
+    type_name = _CLASS_TO_NAME.get(artefact_type)
+    if type_name is None:
+        return []
+
+    rows = conn.execute(
+        """
+        SELECT urn, _document
+        FROM artefacts
+        WHERE artefact_type = $type_name
+        """,
+        {"type_name": type_name},
+    ).fetchall()
+
+    results: list[Any] = []
+    for urn_str, doc_str in rows:
+        data = json.loads(doc_str)
+        if type_name == "Codelist" and not eager:
+            results.append(make_lazy_codelist(conn, urn_str, data))
+        elif type_name == "Hierarchy":
+            data["tree"] = load_hierarchy_tree(conn, urn_str)
+            results.append(from_dict(artefact_type, data))
+        elif type_name == "Codelist":
+            data["codes"] = load_code_dicts(conn, urn_str)
+            results.append(from_dict(artefact_type, data))
+        else:
+            results.append(from_dict(artefact_type, data))
+    return results
+
+
+# ── LazyCodelist construction ───────────────────────────────────────────
+
+
+def make_lazy_codelist(conn: duckdb.DuckDBPyConnection, urn: str, data: dict[str, Any]) -> LazyCodelist:
+    """Construct a LazyCodelist from metadata dict and DuckDB connection."""
+    name = InternationalString.of(data["name"]) if data.get("name") else InternationalString()
+    desc = InternationalString.of(data["description"]) if data.get("description") else InternationalString()
+    agency = Agency(id=data.get("agency_id", data.get("maintainer", "")))
+    return LazyCodelist(
+        urn=urn,
+        conn=conn,
+        id=data["id"],
+        agency=agency,
+        version=data.get("version", "1.0"),
+        name=name,
+        description=desc,
+        is_final=data.get("is_final", False),
+    )
+
+
+# ── Direct SQL helpers ──────────────────────────────────────────────────
+
+
+def code_map(
+    conn: duckdb.DuckDBPyConnection,
+    codelist_urn: str,
+    lang: str = "en",
+) -> dict[str, str]:
+    """Return ``{code_id: label}`` for a codelist via direct SQL.
+
+    If the requested language column doesn't exist, returns an empty dict.
+    """
+    languages = get_languages(conn)
+    if lang not in languages:
+        return {}
+
+    col = f"label_{lang}"
+    rows = conn.execute(
+        f"""
+        SELECT code_id, {col}
+        FROM codes
+        WHERE codelist_urn = $codelist_urn
+        """,  # noqa: S608
+        {"codelist_urn": codelist_urn},
+    ).fetchall()
+    return {r[0]: r[1] for r in rows if r[1] is not None}
+
+
+def code_map_for_dsd(
+    conn: duckdb.DuckDBPyConnection,
+    dsd_urn: str,
+    lang: str = "en",
+) -> dict[str, dict[str, str]]:
+    """Return ``{component_id: {code_id: label}}`` for all enumerated components.
+
+    Joins ``dsd_components`` with ``codes`` to build the map without
+    materializing Python objects.
+    """
+    rows = conn.execute(
+        """
+        SELECT component_id, codelist_urn
+        FROM dsd_components
+        WHERE dsd_urn = $dsd_urn AND codelist_urn IS NOT NULL
+        """,
+        {"dsd_urn": dsd_urn},
+    ).fetchall()
+
+    result: dict[str, dict[str, str]] = {}
+    for comp_id, cl_urn in rows:
+        cm = code_map(conn, cl_urn, lang)
+        if cm:
+            result[comp_id] = cm
+    return result