224 lines
6.6 KiB
Python
224 lines
6.6 KiB
Python
from __future__ import annotations
|
|
|
|
import csv
|
|
import hashlib
|
|
import io
|
|
import json
|
|
from collections.abc import Mapping, Sequence
|
|
from datetime import date, datetime
|
|
from decimal import Decimal
|
|
from typing import Any
|
|
|
|
from govoplan_core.core.datasources import DatasourceField, DatasourceValidationError
|
|
|
|
|
|
MAX_STAGE_ROWS = 10_000
|
|
MAX_STAGE_BYTES = 5_000_000
|
|
MAX_READ_ROWS = 500
|
|
|
|
|
|
def normalize_rows(
|
|
rows: Sequence[Mapping[str, object]],
|
|
) -> tuple[dict[str, Any], ...]:
|
|
if len(rows) > MAX_STAGE_ROWS:
|
|
raise DatasourceValidationError(
|
|
f"Staging is limited to {MAX_STAGE_ROWS:,} rows."
|
|
)
|
|
normalized = tuple(_json_row(row) for row in rows)
|
|
if encoded_size(normalized) > MAX_STAGE_BYTES:
|
|
raise DatasourceValidationError(
|
|
f"Staging is limited to {MAX_STAGE_BYTES // 1_000_000} MB."
|
|
)
|
|
return normalized
|
|
|
|
|
|
def parse_csv_rows(
|
|
csv_text: str,
|
|
*,
|
|
delimiter: str = ",",
|
|
) -> tuple[dict[str, Any], ...]:
|
|
if delimiter not in {",", ";", "\t", "|"}:
|
|
raise DatasourceValidationError("Unsupported CSV delimiter.")
|
|
if len(csv_text.encode("utf-8")) > MAX_STAGE_BYTES:
|
|
raise DatasourceValidationError(
|
|
f"Staging is limited to {MAX_STAGE_BYTES // 1_000_000} MB."
|
|
)
|
|
try:
|
|
reader = csv.reader(io.StringIO(csv_text), delimiter=delimiter)
|
|
raw_header = next(reader, None)
|
|
if not raw_header:
|
|
raise DatasourceValidationError("CSV input needs a non-empty header row.")
|
|
header = [
|
|
str(name or "").removeprefix("\ufeff").strip()
|
|
for name in raw_header
|
|
]
|
|
if any(not name for name in header):
|
|
raise DatasourceValidationError("CSV input needs a non-empty header row.")
|
|
if len(set(header)) != len(header):
|
|
raise DatasourceValidationError("CSV headers must be unique.")
|
|
rows: list[dict[str, object]] = []
|
|
for line_number, values in enumerate(reader, start=2):
|
|
if len(values) != len(header):
|
|
raise DatasourceValidationError(
|
|
f"CSV row {line_number} has {len(values)} values; expected {len(header)}."
|
|
)
|
|
rows.append(
|
|
{
|
|
name: _csv_value(value)
|
|
for name, value in zip(header, values, strict=True)
|
|
}
|
|
)
|
|
except csv.Error as exc:
|
|
raise DatasourceValidationError(f"CSV input is invalid: {exc}") from exc
|
|
return normalize_rows(rows)
|
|
|
|
|
|
def infer_schema(
|
|
rows: Sequence[Mapping[str, object]],
|
|
) -> tuple[DatasourceField, ...]:
|
|
names: list[str] = []
|
|
for row in rows:
|
|
for name in row:
|
|
if name not in names:
|
|
names.append(name)
|
|
result: list[DatasourceField] = []
|
|
for name in names:
|
|
values = [row.get(name) for row in rows]
|
|
concrete = [value for value in values if value is not None]
|
|
data_type = _type_name(concrete[0]) if concrete else "unknown"
|
|
if any(_type_name(value) != data_type for value in concrete[1:]):
|
|
data_type = "mixed"
|
|
result.append(
|
|
DatasourceField(
|
|
name=name,
|
|
data_type=data_type,
|
|
nullable=len(concrete) != len(values),
|
|
)
|
|
)
|
|
return tuple(result)
|
|
|
|
|
|
def fingerprint_rows(
|
|
rows: Sequence[Mapping[str, object]],
|
|
schema: Sequence[DatasourceField],
|
|
) -> str:
|
|
payload = {
|
|
"schema": [field_payload(field) for field in schema],
|
|
"rows": [dict(row) for row in rows],
|
|
}
|
|
encoded = json.dumps(payload, sort_keys=True, separators=(",", ":"), default=str)
|
|
return hashlib.sha256(encoded.encode("utf-8")).hexdigest()
|
|
|
|
|
|
def encoded_size(rows: Sequence[Mapping[str, object]]) -> int:
|
|
return len(_encoded_rows(rows))
|
|
|
|
|
|
def payload_checksum(rows: Sequence[Mapping[str, object]]) -> str:
|
|
return hashlib.sha256(_encoded_rows(rows)).hexdigest()
|
|
|
|
|
|
def payload_row_checksum(row: Mapping[str, object]) -> str:
|
|
return hashlib.sha256(_encoded_row(row)).hexdigest()
|
|
|
|
|
|
def _encoded_rows(rows: Sequence[Mapping[str, object]]) -> bytes:
|
|
return json.dumps(
|
|
rows,
|
|
sort_keys=True,
|
|
separators=(",", ":"),
|
|
default=str,
|
|
).encode("utf-8")
|
|
|
|
|
|
def _encoded_row(row: Mapping[str, object]) -> bytes:
|
|
return json.dumps(
|
|
row,
|
|
sort_keys=True,
|
|
separators=(",", ":"),
|
|
default=str,
|
|
).encode("utf-8")
|
|
|
|
|
|
def field_payload(field: DatasourceField) -> dict[str, object]:
|
|
return {
|
|
"name": field.name,
|
|
"data_type": field.data_type,
|
|
"nullable": field.nullable,
|
|
"classification": field.classification,
|
|
}
|
|
|
|
|
|
def _json_row(row: Mapping[str, object]) -> dict[str, Any]:
|
|
normalized = {str(key).strip(): value for key, value in row.items()}
|
|
if not normalized or any(not key for key in normalized):
|
|
raise DatasourceValidationError("Every row needs named columns.")
|
|
try:
|
|
encoded = json.dumps(normalized, default=_json_value)
|
|
except (TypeError, ValueError) as exc:
|
|
raise DatasourceValidationError(
|
|
f"Datasource values must be JSON compatible: {exc}"
|
|
) from exc
|
|
return json.loads(encoded)
|
|
|
|
|
|
def _json_value(value: object) -> object:
|
|
if isinstance(value, (date, datetime, Decimal)):
|
|
return str(value)
|
|
raise TypeError(f"{type(value).__name__} is not JSON serializable")
|
|
|
|
|
|
def _type_name(value: object) -> str:
|
|
if isinstance(value, bool):
|
|
return "boolean"
|
|
if isinstance(value, int):
|
|
return "integer"
|
|
if isinstance(value, (float, Decimal)):
|
|
return "number"
|
|
if isinstance(value, str):
|
|
return "string"
|
|
if isinstance(value, list):
|
|
return "array"
|
|
if isinstance(value, dict):
|
|
return "object"
|
|
return type(value).__name__.lower()
|
|
|
|
|
|
def _csv_value(value: str) -> object:
|
|
cleaned = value.strip()
|
|
if not cleaned:
|
|
return None
|
|
lowered = cleaned.lower()
|
|
if lowered == "true":
|
|
return True
|
|
if lowered == "false":
|
|
return False
|
|
if (
|
|
cleaned.isdigit()
|
|
and (cleaned == "0" or not cleaned.startswith("0"))
|
|
) or (
|
|
cleaned.startswith("-")
|
|
and cleaned[1:].isdigit()
|
|
and (cleaned[1:] == "0" or not cleaned[1:].startswith("0"))
|
|
):
|
|
return int(cleaned)
|
|
try:
|
|
return float(cleaned)
|
|
except ValueError:
|
|
return cleaned
|
|
|
|
|
|
__all__ = [
|
|
"MAX_READ_ROWS",
|
|
"MAX_STAGE_BYTES",
|
|
"MAX_STAGE_ROWS",
|
|
"encoded_size",
|
|
"field_payload",
|
|
"fingerprint_rows",
|
|
"infer_schema",
|
|
"normalize_rows",
|
|
"payload_checksum",
|
|
"payload_row_checksum",
|
|
"parse_csv_rows",
|
|
]
|