diff --git a/libs/data_manipulation/src/data_manipulation/ingestion.py b/libs/data_manipulation/src/data_manipulation/ingestion.py index 8ce35f81..95418dcc 100644 --- a/libs/data_manipulation/src/data_manipulation/ingestion.py +++ b/libs/data_manipulation/src/data_manipulation/ingestion.py @@ -1,3 +1,4 @@ +import codecs import logging import os import re @@ -16,7 +17,7 @@ import pyarrow.parquet as pq import requests from geoalchemy2 import Geometry -from pyarrow.lib import ArrowException +from pyarrow import ArrowException from sqlalchemy import MetaData, Table, func, select, text from sqlalchemy.engine import Engine @@ -30,13 +31,12 @@ logger = logging.getLogger(__name__) DEFAULT_SCHEMA = "public" +CHUNK_SIZE = int(os.getenv("DATAFEEDER_CHUNK_SIZE", 50000)) -# Bytes sampled for encoding detection. chardet's accuracy is unchanged for a sample -# this size, and reading only a sample avoids loading multi-GB files into memory. _ENCODING_DETECT_BYTES = 256 * 1024 -# Number of rows read and written to PostGIS per chunk. Keeps the memory footprint low -# (only one chunk is held in memory / converted to WKB at a time) for large files. -CHUNK_SIZE = int(os.getenv("DATAFEEDER_CHUNK_SIZE", 50000)) +_DEFAULT_ENCODING = "utf-8" +_UTF8_COMPATIBLE_ENCODINGS = (_DEFAULT_ENCODING, "ascii") +_SHAPEFILE_FALLBACK_ENCODING = "cp1252" def _get_table_row_count(table_name: str, engine: Engine, schema: str) -> int: @@ -48,8 +48,63 @@ def _get_table_row_count(table_name: str, engine: Engine, schema: str) -> int: return conn.execute(count_query).scalar() or 0 +def _parse_cpg_encoding(sample: bytes) -> str | None: + """Read the encoding a shapefile .cpg sidecar declares. + + A .cpg *declares* the .dbf encoding as ASCII text ("ISO-8859-1", "UTF-8", "1252", + "ANSI 1252", ...), so it has to be parsed rather than handed to a charset detector: + chardet reports plain ASCII text as "ascii", which decodes nothing. + + A declaration of UTF-8 or ASCII is discarded. The sidecar is only consulted after a + UTF-8 read has already failed, so honouring such a declaration would do no more than + repeat that read. + + Args: + sample: Raw bytes of the .cpg sidecar + + Returns: + Normalized codec name, or None when the sidecar declares nothing usable + """ + try: + declared = sample.decode("ascii").strip() + except UnicodeDecodeError: + return None + + # The codepage is the trailing token: "ANSI 1252" -> "1252", "LDID/87" -> "87". + candidate = re.split(r"[\s/]+", declared)[-1] if declared else "" + + try: + name = codecs.lookup(candidate).name + except LookupError: + logger.warning("Ignoring unusable .cpg encoding declaration: %r", declared) + return None + + return None if name in _UTF8_COMPATIBLE_ENCODINGS else name + + +def _detect_shapefile_encoding(cpg_sample: bytes | None) -> str: + """Pick the encoding to re-read a shapefile with, from its .cpg sidecar. + + Args: + cpg_sample: Raw bytes of the .cpg sidecar, or None when there is none + + Returns: + Encoding string + """ + declared = _parse_cpg_encoding(cpg_sample) if cpg_sample else None + if declared is not None: + return declared + + logger.warning( + "Shapefile declares no usable encoding, falling back to %s. Check the ingested " + "text for mojibake; a .cpg sidecar naming the right encoding fixes it at the source.", + _SHAPEFILE_FALLBACK_ENCODING, + ) + return _SHAPEFILE_FALLBACK_ENCODING + + def _detect_file_encoding(file_path: str) -> str: - """Detect encoding for geospatial files. + """Detect the encoding to retry a failed UTF-8 read with. Args: file_path: Path to the file @@ -57,44 +112,35 @@ def _detect_file_encoding(file_path: str) -> str: Returns: Detected encoding string """ - file_path_to_read = file_path path = Path(file_path) # GeoJSON must be UTF-8 according to RFC 7946 if path.suffix.lower() in (".geojson", ".json"): - return "utf-8" - - # Check for .cpg file (encoding file for shapefiles) - if path.suffix.lower() == ".shp": - cpg_file = path.with_suffix(".cpg") - if cpg_file.exists(): - file_path_to_read = str(cpg_file) + return _DEFAULT_ENCODING try: + if path.suffix.lower() == ".shp": + cpg_file = path.with_suffix(".cpg") + return _detect_shapefile_encoding(cpg_file.read_bytes() if cpg_file.exists() else None) + if path.suffix.lower() == ".zip" and zipfile.is_zipfile(file_path): with zipfile.ZipFile(file_path) as zf: names = zf.namelist() - if any(name.lower().endswith(".shp") for name in names): - member = next((n for n in names if n.lower().endswith(".cpg")), None) or next( - (n for n in names if n.lower().endswith(".dbf")), None - ) - if member is not None: - file_path_to_read = f"{file_path}!{member}" - with zf.open(member) as f: - sample = f.read(_ENCODING_DETECT_BYTES) - else: - sample = None - else: - sample = None - else: - with open(file_path_to_read, "rb") as f: - sample = f.read(_ENCODING_DETECT_BYTES) - encoding = chardet.detect(sample)["encoding"] if sample else None + if not any(name.lower().endswith(".shp") for name in names): + return _DEFAULT_ENCODING + + member = next((n for n in names if n.lower().endswith(".cpg")), None) + if member is None: + return _detect_shapefile_encoding(None) + with zf.open(member) as f: + return _detect_shapefile_encoding(f.read(_ENCODING_DETECT_BYTES)) + + with open(file_path, "rb") as f: + sample = f.read(_ENCODING_DETECT_BYTES) + return chardet.detect(sample)["encoding"] or _DEFAULT_ENCODING except Exception as e: - logger.warning(f"Failed to detect encoding for {file_path_to_read}: {e}") - encoding = None - - return encoding or "utf-8" + logger.warning(f"Failed to detect encoding for {file_path}: {e}") + return _DEFAULT_ENCODING def _read_file_encoded(file_path: str, i: int = 0) -> gpd.GeoDataFrame | pd.DataFrame: @@ -124,19 +170,16 @@ def _read_file_encoded(file_path: str, i: int = 0) -> gpd.GeoDataFrame | pd.Data return pd.read_parquet(ds.fragments[i].path) try: - # Try reading with UTF-8 first (common default) - result = gpd.read_file(file_path, rows=rows) # type: ignore[arg-type] - # pyogrio/pyarrow validate text lazily: a bad encoding doesn't raise here, it - # raises whenever the string columns are first materialized (e.g. much later - # in pandas.to_sql when writing to PostGIS). Force that materialization now - # so we can catch and handle it in this try block instead. - for column in result.columns: - if pd.api.types.is_string_dtype(result[column].dtype): - result[column].to_numpy() - return result - except (UnicodeDecodeError, ArrowException): + # Try reading with UTF-8 first (common default). Undecodable text is reported + # right here: with PYOGRIO_USE_ARROW the conversion to pandas raises an + # ArrowException, and without it GDAL recodes from the encoding the shapefile + # declares (.cpg / .dbf language driver id) or from ISO-8859-1. + return gpd.read_file(file_path, rows=rows) # type: ignore[arg-type] + except (UnicodeDecodeError, ArrowException) as e: logger.warning( - "Failed to read file with UTF-8 encoding, attempting to detect encoding and read again." + "Failed to read file with UTF-8 encoding (%s), detecting the encoding " + "and reading again.", + e, ) # Detect encoding (mainly for shapefiles, others default to UTF-8) diff --git a/libs/data_manipulation/tests/test_ingestion.py b/libs/data_manipulation/tests/test_ingestion.py index fb9bfe02..3a289689 100644 --- a/libs/data_manipulation/tests/test_ingestion.py +++ b/libs/data_manipulation/tests/test_ingestion.py @@ -1,5 +1,8 @@ """Tests for data ingestion utilities in data_manipulation library.""" +import codecs +import zipfile +from pathlib import Path from unittest.mock import MagicMock, Mock, patch from urllib.error import URLError @@ -7,14 +10,18 @@ import pytest import requests from geopandas import GeoDataFrame -from pandas import DataFrame +from pandas import DataFrame, Index from shapely.geometry import Point from sqlalchemy.engine import Engine from data_manipulation import IntegrityTransformation, apply_transformations from data_manipulation.constants import POSTGIS_TABLE_NAME_MAX_LENGTH from data_manipulation.ingestion import ( + _DEFAULT_ENCODING, # pyright: ignore[reportPrivateUsage] + _SHAPEFILE_FALLBACK_ENCODING, # pyright: ignore[reportPrivateUsage] + _UTF8_COMPATIBLE_ENCODINGS, # pyright: ignore[reportPrivateUsage] CHUNK_SIZE, + _detect_file_encoding, # pyright: ignore[reportPrivateUsage] _read_file_encoded, # pyright: ignore[reportPrivateUsage] ingest_data_from_database_into_postgis, ingest_data_from_file_into_postgis, @@ -26,6 +33,141 @@ ) +class TestDetectFileEncoding: + """Encoding picked to retry a read that failed as UTF-8.""" + + @staticmethod + def _loose_shapefile(tmp_path: Path, cpg: bytes | None) -> str: + (tmp_path / "t.shp").write_bytes(b"") + (tmp_path / "t.dbf").write_bytes(b"") + if cpg is not None: + (tmp_path / "t.cpg").write_bytes(cpg) + return str(tmp_path / "t.shp") + + @staticmethod + def _zipped_shapefile(tmp_path: Path, cpg: bytes | None) -> str: + archive = tmp_path / "shape.zip" + with zipfile.ZipFile(archive, "w") as zf: + zf.writestr("t.shp", b"") + zf.writestr("t.dbf", b"") + if cpg is not None: + zf.writestr("t.cpg", cpg) + return str(archive) + + def test_encoding_constants_are_normalized_codec_names(self) -> None: + """The constants are compared against codecs.lookup().name, so they must match it. + + A typo such as "utf8" names a real codec but never equals a normalized name, which + would silently stop _parse_cpg_encoding from discarding UTF-8 declarations. + """ + for name in (_DEFAULT_ENCODING, _SHAPEFILE_FALLBACK_ENCODING, *_UTF8_COMPATIBLE_ENCODINGS): + assert codecs.lookup(name).name == name + + def test_loose_shapefile_uses_cpg_declaration(self, tmp_path: Path) -> None: + """A .cpg is parsed as an encoding name, not sniffed as text.""" + assert _detect_file_encoding(self._loose_shapefile(tmp_path, b"ISO-8859-1")) == "iso8859-1" + + def test_zipped_shapefile_uses_cpg_declaration(self, tmp_path: Path) -> None: + """The .cpg member of a zipped shapefile is parsed the same way.""" + assert _detect_file_encoding(self._zipped_shapefile(tmp_path, b"ISO-8859-1")) == "iso8859-1" + + def test_cpg_codepage_number_is_resolved(self, tmp_path: Path) -> None: + """A bare codepage, with or without a vendor prefix, resolves to its codec.""" + assert _detect_file_encoding(self._loose_shapefile(tmp_path, b"ANSI 1252")) == "cp1252" + + def test_unusable_cpg_is_ignored(self, tmp_path: Path) -> None: + """A .cpg that names no known codec never becomes the encoding.""" + encoding = _detect_file_encoding(self._loose_shapefile(tmp_path, b"not-an-encoding\n")) + + assert encoding != "ascii" + codecs.lookup(encoding) + + def test_shapefile_without_cpg_falls_back_to_cp1252(self, tmp_path: Path) -> None: + """No sidecar at all: guess the codepage shapefiles are usually written in.""" + assert _detect_file_encoding(self._loose_shapefile(tmp_path, None)) == "cp1252" + + def test_zipped_shapefile_without_cpg_falls_back_to_cp1252(self, tmp_path: Path) -> None: + """Same for a zipped shapefile, rather than returning the UTF-8 that just failed.""" + assert _detect_file_encoding(self._zipped_shapefile(tmp_path, None)) == "cp1252" + + def test_shapefile_declaring_utf8_falls_back_to_cp1252(self, tmp_path: Path) -> None: + """A .cpg contradicted by the data is a wrong declaration, not an answer.""" + assert _detect_file_encoding(self._zipped_shapefile(tmp_path, b"UTF-8")) == "cp1252" + + def test_geojson_is_always_utf8(self, tmp_path: Path) -> None: + """GeoJSON is UTF-8 by RFC 7946, whatever the bytes look like.""" + geojson = tmp_path / "t.geojson" + geojson.write_bytes("café".encode("latin-1")) + + assert _detect_file_encoding(str(geojson)) == "utf-8" + + def test_plain_text_file_is_sniffed(self, tmp_path: Path) -> None: + """Non-shapefile inputs keep using chardet on their raw bytes.""" + csv = tmp_path / "t.csv" + csv.write_bytes(("name\n" + "Brévalé çà où\n" * 50).encode("latin-1")) + + assert codecs.lookup(_detect_file_encoding(str(csv))).name != "utf-8" + + def test_zip_without_shapefile_defaults_to_utf8(self, tmp_path: Path) -> None: + """Compressed bytes tell nothing about the encoding of the data inside.""" + archive = tmp_path / "data.zip" + with zipfile.ZipFile(archive, "w") as zf: + zf.writestr("t.csv", "name\ncafé\n") + + assert _detect_file_encoding(str(archive)) == "utf-8" + + def test_unreadable_file_defaults_to_utf8(self, tmp_path: Path) -> None: + """Detection never raises; it falls back to UTF-8.""" + assert _detect_file_encoding(str(tmp_path / "missing.csv")) == "utf-8" + + +class TestReadFileEncodedShapefile: + """Reading real shapefiles whose text is not UTF-8.""" + + NAMES = ["Brévalé çà où ü0", "Ångström", "Müller & Cie", "Zoé"] + + def _latin1_shapefile_zip(self, tmp_path: Path, cpg: bytes | None) -> str: + """Write a latin-1 shapefile with the given .cpg declaration and zip it.""" + source = gpd.GeoDataFrame( + {"name": self.NAMES, "n": range(len(self.NAMES))}, + geometry=[Point(i, i) for i in range(len(self.NAMES))], + crs="EPSG:4326", + ) + source.to_file(tmp_path / "t.shp", encoding="ISO-8859-1") + + cpg_file = tmp_path / "t.cpg" + cpg_file.unlink(missing_ok=True) + if cpg is not None: + cpg_file.write_bytes(cpg) + + archive = tmp_path / "shape.zip" + with zipfile.ZipFile(archive, "w") as zf: + for member in sorted(tmp_path.glob("t.*")): + zf.write(member, member.name) + return str(archive) + + @pytest.mark.parametrize("cpg", [None, b"UTF-8"], ids=["no-cpg", "cpg-lies-utf8"]) + @pytest.mark.parametrize("use_arrow", ["0", "1"]) + def test_zipped_latin1_shapefile_reads_its_text_intact( + self, + tmp_path: Path, + cpg: bytes | None, + use_arrow: str, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + """A shapefile GDAL cannot decode is re-read with the fallback encoding. + + Parametrized over PYOGRIO_USE_ARROW because only the Arrow path raises on + undecodable text, and Airflow enables it while local runs do not. + """ + monkeypatch.setenv("PYOGRIO_USE_ARROW", use_arrow) + archive = self._latin1_shapefile_zip(tmp_path, cpg) + + result = _read_file_encoded(archive) + + assert list(result["name"]) == self.NAMES + + class TestReadFileEncodedParquet: """Parquet/GeoParquet dispatch in _read_file_encoded.""" @@ -129,6 +271,14 @@ def test_read_file_encoded_falls_back_to_detected_encoding( mock_detect_encoding.assert_called_once_with("test.shp") assert result is mock_gdf + @patch("data_manipulation.ingestion.gpd.read_file") + def test_read_file_encoded_accepts_duplicate_column_labels(self, mock_read_file: Mock) -> None: + """GDAL can hand back repeated field names; the frame is returned untouched.""" + mock_df = DataFrame([["a", "b", 1]], columns=Index(["name", "name", "val"])) + mock_read_file.return_value = mock_df + + assert _read_file_encoded("dup.csv") is mock_df + class TestIngestDataFromUrlIntoPostgis: """Test cases for ingest_data_from_url_into_postgis function."""