arborist/aborist/sources/wikipedia.py

"""MediaWiki 'cur' table SQL dump source.

Handles 2003-era SQL dumps in bz2 format (e.g. 20030516_cur_tablesql.bz2).
Yields one Document per non-redirect main-namespace article, with
[[wikilinks]] extracted as outbound edges.

Implements a stream parser for MySQL extended INSERT syntax. The cur table
schema for that era starts: cur_id, cur_namespace, cur_title, cur_text, ...
We rely on positional access for the first four columns.
"""

from __future__ import annotations

import bz2
import re
from pathlib import Path
from typing import IO, Iterator

from aborist.document import Document, Edge
from aborist.source import Source


# Match [[Target]], [[Target|display]], [[Target#anchor]] forms.
WIKILINK_RE = re.compile(r"\[\[([^\]\|#\n\r]+)(?:#[^\]\|\n\r]*)?(?:\|[^\]\n\r]*)?\]\]")
NAMESPACE_MAIN = 0


def _decode_mysql_string(s: str) -> str:
    """Decode MySQL-escaped string content (outer quotes already stripped)."""
    out: list[str] = []
    i = 0
    n = len(s)
    while i < n:
        c = s[i]
        if c == "\\" and i + 1 < n:
            nxt = s[i + 1]
            mapped = {
                "n": "\n",
                "r": "\r",
                "t": "\t",
                "0": "\0",
                "\\": "\\",
                "'": "'",
                '"': '"',
                "Z": "\x1a",
            }.get(nxt, nxt)
            out.append(mapped)
            i += 2
        else:
            out.append(c)
            i += 1
    return "".join(out)


def _split_values_tuples(payload: str) -> list[list[str | None]]:
    """Parse `(v1,v2,...),(...),...` into tuples of decoded strings or None."""
    rows: list[list[str | None]] = []
    i = 0
    n = len(payload)
    while i < n:
        while i < n and payload[i] in " \t\n\r,":
            i += 1
        if i >= n or payload[i] != "(":
            break
        i += 1
        values: list[str | None] = []
        while True:
            while i < n and payload[i] in " \t":
                i += 1
            if i >= n:
                break
            if payload[i] == "'":
                # Quoted string
                i += 1
                buf: list[str] = []
                while i < n:
                    c = payload[i]
                    if c == "\\" and i + 1 < n:
                        # Preserve escape pair; decode later in one pass.
                        buf.append(c)
                        buf.append(payload[i + 1])
                        i += 2
                    elif c == "'":
                        i += 1
                        break
                    else:
                        buf.append(c)
                        i += 1
                values.append(_decode_mysql_string("".join(buf)))
            else:
                # NULL or number
                start = i
                while i < n and payload[i] not in ",)":
                    i += 1
                v = payload[start:i].strip()
                values.append(None if v.upper() == "NULL" else v)
            while i < n and payload[i] in " \t":
                i += 1
            if i < n and payload[i] == ",":
                i += 1
                continue
            if i < n and payload[i] == ")":
                i += 1
                break
            break
        rows.append(values)
    return rows


def _iter_insert_statements(file_obj: IO[str]) -> Iterator[str]:
    """Yield complete SQL statements (text up to ; outside a quoted string)."""
    buf: list[str] = []
    in_string = False
    escape = False
    while True:
        chunk = file_obj.read(1 << 17)
        if not chunk:
            break
        for c in chunk:
            buf.append(c)
            if escape:
                escape = False
                continue
            if in_string:
                if c == "\\":
                    escape = True
                elif c == "'":
                    in_string = False
                continue
            if c == "'":
                in_string = True
            elif c == ";":
                yield "".join(buf)
                buf = []
    if buf:
        tail = "".join(buf).strip()
        if tail:
            yield tail


def _extract_wikilinks(text: str, base_uri: str) -> list[Edge]:
    seen: set[str] = set()
    edges: list[Edge] = []
    for m in WIKILINK_RE.finditer(text):
        target = m.group(1).strip()
        if not target or target.startswith(":"):
            continue
        # Skip image / file / category interlinks (they often start "Image:" etc.)
        if ":" in target:
            continue
        uri = base_uri + target.replace(" ", "_")
        if uri in seen:
            continue
        seen.add(uri)
        edges.append(Edge(edge_type="wikilink", dst_uri=uri))
    return edges


class WikipediaCurDump(Source):
    """Iterates a MediaWiki 'cur' table SQL dump (bz2-compressed)."""

    source_type = "wikipedia_cur"

    def __init__(
        self,
        path: str | Path,
        namespace: int = NAMESPACE_MAIN,
        base_uri: str = "https://en.wikipedia.org/wiki/",
    ):
        self.path = Path(path)
        self.namespace = namespace
        self.base_uri = base_uri

    def iter_documents(self) -> Iterator[Document]:
        opener = bz2.open if str(self.path).endswith(".bz2") else open
        with opener(self.path, "rt", encoding="utf-8", errors="replace") as f:
            for stmt in _iter_insert_statements(f):
                head = stmt.lstrip()
                if not head.upper().startswith("INSERT INTO"):
                    continue
                up = head.upper()
                vidx = up.find("VALUES")
                if vidx < 0:
                    continue
                table_clause = head[:vidx].lower()
                if "cur" not in table_clause:
                    continue
                payload = head[vidx + len("VALUES"):]
                payload = payload.rstrip().rstrip(";").rstrip()
                for row in _split_values_tuples(payload):
                    yield from self._row_to_doc(row)

    def _row_to_doc(self, row: list[str | None]) -> Iterator[Document]:
        if len(row) < 4:
            return
        try:
            ns = int(row[1]) if row[1] is not None else None
        except (ValueError, TypeError):
            return
        if ns != self.namespace:
            return
        title = row[2] or ""
        text = row[3] or ""
        if not title or not text:
            return
        # is_redirect lives at position 10 in the standard 2003 schema; be defensive.
        is_redirect = False
        if len(row) > 10 and row[10] is not None:
            try:
                is_redirect = bool(int(row[10]))
            except (ValueError, TypeError):
                is_redirect = False
        if is_redirect:
            return
        edges = _extract_wikilinks(text, self.base_uri)
        uri = self.base_uri + title.replace(" ", "_")
        yield Document(
            uri=uri,
            content=text,
            source_type=self.source_type,
            title=title,
            edges=edges,
        )