diff --git a/docs/INGEST.md b/docs/INGEST.md new file mode 100644 index 0000000..7077868 --- /dev/null +++ b/docs/INGEST.md @@ -0,0 +1,166 @@ +# S3 ingest: the tag contract and provenance columns + +The helpers in `src/g3dt/ingest/ingest.py` are **optional, +bring-your-own-ingestion utilities**. Use them when you already have flat +files (CSV / JSON / XLSX) landing in S3 and want them queryable in Athena +with full provenance. They are *not* the platform's primary submission +path: the supported no-code route is filling in a +gen3-metadata-templates workbook, which flows into the bronze layer +automatically — see +[DATA_LAYERS.md](https://github.com/AustralianBioCommons/aws-gen3-pipeline/blob/main/docs/DATA_LAYERS.md). + +## The tag contract + +Files opt in to ingestion via **S3 object tags** — not by their location, +name, or upload time. + +| Tag | Required | Purpose | +| --- | --- | --- | +| `ingest` | for discovery | `get_ingest_true_files` keeps only objects tagged `ingest=true` | +| `node` | yes | names the target table: `{table_prefix}_{node}` | +| `study_id` | yes | becomes the `study_id` column (and a partition for parquet) | +| `submission_date` | yes | parsed to `YYYY-MM-DD` and stamped on every row | + +`submission_date` accepts `YYYY-MM-DD` or `DD-MM-YYYY`, with `_` or spaces +as separators (`sanitize_submission_date`); anything else raises +`ValueError`. A missing `node` or `submission_date` tag fails the file +with `ValueError`; a missing `study_id` fails with a `KeyError`. + +Every tag on the object — required or not — is also carried into the +table as a `tag_` column. + +> **Untagged objects are silently skipped — by design.** This is the #1 +> surprise. `get_ingest_true_files` treats the `ingest=true` tag as the +> "this file is ready" signal: an object with no tags, unreadable tags, +> or any value other than the literal string `true` is dropped from the +> scan with only a debug-level log line. If a file you expected is +> missing from the bronze table, check its tags first: +> `aws s3api get-object-tagging --bucket --key `. + +## Operational sharp edges + +- **IAM**: the scanner calls `GetObjectTagging` on every listed object, + so the ingesting role needs `s3:GetObjectTagging` (plus + `s3:ListBucket` / `s3:GetObject`). Whoever uploads and tags files + needs `s3:PutObjectTagging`. +- **Copies lose tags.** The high-level `aws s3 cp` / `aws s3 sync` + commands do not preserve object tags, so a copied file silently drops + out of the ingest scan. Use + `aws s3api copy-object ... --tagging-directive COPY` to propagate + tags, or re-tag after copying. +- **Console uploads aren't tagged** unless you expand the upload + wizard's properties step and add tags by hand. Uploads from scripts + are untagged unless the script tags them. +- Tagging an existing object is one call: + + ```bash + aws s3api put-object-tagging \ + --bucket my-bucket \ + --key raw/mystudy/samples.csv \ + --tagging 'TagSet=[{Key=ingest,Value=true},{Key=study_id,Value=mystudy},{Key=node,Value=sample},{Key=submission_date,Value=2026-08-05}]' + ``` + + (Note: `put-object-tagging` *replaces* the whole tag set — include + every tag, not just the one you are adding.) + +## Provenance columns + +`prepare_ingest_metadata` stamps these columns on every row, alongside +the file's raw columns: + +| Column | Meaning | +| --- | --- | +| `study_id` | from the `study_id` tag | +| `submission_date` | the `submission_date` tag, normalised to `YYYY-MM-DD` | +| `ingest_run_id` | one UUID per `ingest_files_to_dataset` call, shared by every file in that call | +| `ingest_received_at` | timestamp the run started (`Australia/Melbourne`) | +| `ingest_timezone` | timezone of the timestamps (`Australia/Melbourne`) | +| `ingest_original_file_path` | full `s3://` URI of the source object | +| `ingest_file_name` | basename of the source object | +| `ingest_submission_id` | caller-supplied submission ID (empty string if none) | +| `ingest_file_etag` | S3 ETag of the object | +| `ingest_file_size_bytes` | object size | +| `ingest_file_last_modified` | S3 LastModified (`Australia/Melbourne`) | +| `tag_` | one column per S3 object tag | +| `ingest_row_hash` | SHA-256 content hash of the row (see below) | + +### `ingest_row_hash`: re-ingest as a no-op + +`compute_row_hash` hashes **raw columns only** — every `ingest_*` and +`tag_*` column plus `study_id` and `submission_date` are excluded. The +remaining columns are sorted by name, serialised as `col=value` pairs +joined with `||`, and SHA-256 hashed. + +Because run-varying metadata (run ID, timestamps, ETag, path) never +enters the hash, re-ingesting the same file — or the same file under a +new name or date — produces identical hashes. Deduplicate on it, e.g. +by passing `merge_cols=["ingest_row_hash"]` to `write_iceberg_to_db`, +or by using it as the unique key in a downstream dbt incremental model, +and a re-ingest becomes a MERGE no-op instead of duplicated rows. (The +default ingest write is an append; deduplication on the hash is up to +the consumer.) + +For XLSX files, every sheet is read (via openpyxl), a `sheet_name` +column is added, and the sheets are concatenated; `sheet_name` counts +as a raw column and participates in the hash. + +## Table format, parallelism, and knobs + +`ingest_table_to_dataset` / `ingest_files_to_dataset` accept: + +- `table_format` — `"iceberg"` (default) or `"parquet"`. + - **iceberg**: written via Athena `MERGE`-capable Iceberg tables + (`write_iceberg_to_db`). The table must already exist, or the + writer needs a table location configured. + - **parquet**: written as a Hive-style dataset + (`write_parquet_to_db`) partitioned by + `study_id / submission_date / ingest_file_name` with snappy + compression and schema evolution enabled. Requires + `dataset_root` (e.g. `s3://bucket/prefix/`) — the dataset's S3 + root; `mode` chooses `"append"` (default), `"overwrite"`, or + `"overwrite_partitions"`. Both are ignored for iceberg. +- `exclude_fn` (`ingest_files_to_dataset`) — file names to skip; + defaults to `['program.json', 'project.json']`. +- `get_ingest_true_files(s3_uri, exclude_directories=None, + max_workers=32)` — the tag scan issues one `GetObjectTagging` call + per object, so it runs them in a thread pool; `max_workers` controls + the concurrency (32 is comfortable for S3). + +Supported formats, chosen by file extension: `csv` (delimiter-sniffed, +UTF-8/cp1252 fallback, headers normalised to snake_case), `json` +(records orient), `xlsx` (all sheets, flattened). Everything is read +and stored as strings. + +## Worked example + +1. Land the file and mark it ready: + + ```bash + aws s3 cp samples.csv s3://my-bucket/raw/mystudy/samples.csv + aws s3api put-object-tagging \ + --bucket my-bucket --key raw/mystudy/samples.csv \ + --tagging 'TagSet=[{Key=ingest,Value=true},{Key=study_id,Value=mystudy},{Key=node,Value=sample},{Key=submission_date,Value=2026-08-05}]' + ``` + +2. Scan and ingest: + + ```python + from g3dt.ingest.ingest import get_ingest_true_files, ingest_files_to_dataset + + files = get_ingest_true_files("s3://my-bucket/raw/mystudy/") + results = ingest_files_to_dataset( + s3_uris=files, + database="my_bronze_db", + table_prefix="raw", + athena_s3_output="s3://my-bucket/athena-output/", + ) + ``` + +3. Query in Athena: + + ```sql + SELECT sample_id, ingest_file_name, ingest_received_at + FROM my_bronze_db.raw_sample + WHERE study_id = 'mystudy' + AND submission_date = '2026-08-05'; + ``` diff --git a/poetry.lock b/poetry.lock index f1fd1cc..b2ff7f6 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1185,6 +1185,18 @@ backoff = ">=1.10.0,<2.0.0" httpx = ">=0.28.1,<1.0.0" requests = ">=2.23.0,<3.0.0" +[[package]] +name = "et-xmlfile" +version = "2.0.0" +description = "An implementation of lxml.xmlfile for the standard library" +optional = false +python-versions = ">=3.8" +groups = ["main"] +files = [ + {file = "et_xmlfile-2.0.0-py3-none-any.whl", hash = "sha256:7a91720bc756843502c3b7504c77b8fe44217c85c537d85037f0f536151b2caa"}, + {file = "et_xmlfile-2.0.0.tar.gz", hash = "sha256:dab3f4764309081ce75662649be815c4c9081e88f0837825f90fd28317d4da54"}, +] + [[package]] name = "exceptiongroup" version = "1.3.1" @@ -2471,6 +2483,21 @@ datalib = ["numpy (>=1)", "pandas (>=1.2.3)", "pandas-stubs (>=1.1.0.11)"] realtime = ["websockets (>=13,<16)"] voice-helpers = ["numpy (>=2.0.2)", "sounddevice (>=0.5.1)"] +[[package]] +name = "openpyxl" +version = "3.1.5" +description = "A Python library to read/write Excel 2010 xlsx/xlsm files" +optional = false +python-versions = ">=3.8" +groups = ["main"] +files = [ + {file = "openpyxl-3.1.5-py2.py3-none-any.whl", hash = "sha256:5282c12b107bffeef825f4617dc029afaf41d0ea60823bbb665ef3079dc79de2"}, + {file = "openpyxl-3.1.5.tar.gz", hash = "sha256:cf0e3cf56142039133628b5acffe8ef0c12bc902d2aadd3e0fe5878dc08d1050"}, +] + +[package.dependencies] +et-xmlfile = "*" + [[package]] name = "packaging" version = "25.0" @@ -4776,4 +4803,4 @@ type = ["pytest-mypy (>=1.0.1) ; platform_python_implementation != \"PyPy\""] [metadata] lock-version = "2.1" python-versions = ">=3.9.5,<4.0.0" -content-hash = "0005c533a5df06cbb6235a4a732f4848fcf553a75aed492355eac8c5acb880fb" +content-hash = "157a8aeec1b13d5881be4f95cf501a0beaca6c5b8875245df8027698a22cb9ed" diff --git a/pyproject.toml b/pyproject.toml index 83e4f21..2b88caf 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -28,6 +28,7 @@ awswrangler = ">=3.14.0,<4.0.0" pyarrow = ">=14.0.0,<19.0.0" gen3-metadata = ">=1.4.0,<2.0.0" typer = ">=0.12" +openpyxl = ">=3.1.0" [tool.poetry.scripts] g3dt = "g3dt.cli.main:main" diff --git a/src/g3dt/ingest/ingest.py b/src/g3dt/ingest/ingest.py index bb13660..f6ab4f9 100644 --- a/src/g3dt/ingest/ingest.py +++ b/src/g3dt/ingest/ingest.py @@ -1,12 +1,15 @@ import awswrangler as wr import pandas as pd import boto3 -from g3dt.utils.athena_utils import write_iceberg_to_db +from g3dt.utils.athena_utils import write_iceberg_to_db, write_parquet_to_db import urllib.parse import os import uuid +import warnings import hashlib +from concurrent.futures import ThreadPoolExecutor from datetime import datetime +from botocore.config import Config from botocore.exceptions import ClientError import logging import pytz # Replaced tzlocal with pytz @@ -17,7 +20,9 @@ logger = logging.getLogger(__name__) logger.setLevel(logging.INFO) -s3 = boto3.client("s3") +# Pool sized above get_ingest_true_files' max_workers so threaded +# GetObjectTagging calls aren't throttled by urllib3's default pool of 10. +s3 = boto3.client("s3", config=Config(max_pool_connections=64)) # ---------- helpers ---------- def parse_s3(uri: str): @@ -352,21 +357,28 @@ def align_and_combine_frames(frames: list[pd.DataFrame]) -> pd.DataFrame: logger.error(f"Failed to align and combine DataFrames: {e}") raise RuntimeError(f"Failed to align and combine DataFrames: {e}") -def get_ingest_true_files(s3_uri: str, exclude_directories: list[str] = None) -> list[str]: +def get_ingest_true_files(s3_uri: str, exclude_directories: list[str] = None, max_workers: int = 32) -> list[str]: """ Get a list of files from an S3 URI, excluding directories. + Tag lookups (GetObjectTagging) are made concurrently with a thread + pool, since checking one object per round-trip is the dominant cost + when scanning buckets with many files. + Parameters ---------- s3_uri : str S3 URI to list files from. exclude_directories : list[str], optional List of directories to exclude. Defaults to None. + max_workers : int, optional + Number of concurrent tag lookups. Defaults to 32, which S3 + handles comfortably for GetObjectTagging. Returns ------- list[str] - List of file paths. + List of file paths, in listing order. """ logger.debug(f"Getting files from {s3_uri}") try: @@ -380,22 +392,25 @@ def get_ingest_true_files(s3_uri: str, exclude_directories: list[str] = None) -> logger.debug(f"Found {len(file_paths)} files after excluding directories: {exclude_directories}") else: logger.debug(f"No directories excluded from {s3_uri}") - - ingest_files = [] - for path in file_paths: + + def check(path): try: tags = get_tags(path) except Exception as e: logger.warning(f"Could not get tags for {path}: {e}") - continue - if not tags or "ingest" not in tags or tags["ingest"] != "true": + return None + if not tags or tags.get("ingest") != "true": logger.debug(f"Skipping {path}: missing or non-true 'ingest' tag") - continue - ingest_files.append(path) + return None + return path + + with ThreadPoolExecutor(max_workers=max_workers) as executor: + results = executor.map(check, file_paths) + ingest_files = [path for path in results if path] logger.debug(f"Found {len(ingest_files)} files with 'ingest' tag set to 'true'") return ingest_files -def ingest_table_to_parquet_dataset( +def ingest_table_to_dataset( s3_uri: str, database: str, table_prefix: str, @@ -405,16 +420,19 @@ def ingest_table_to_parquet_dataset( ingest_submission_id: str = None, ingest_received_at: str = None, ingest_run_id: str = None, + table_format: str = "iceberg", + dataset_root: str = None, + mode: str = "append", ) -> dict: """ - Ingest a single CSV or JSON file from S3, annotate with metadata, and write to an Iceberg table. + Ingest a single CSV or JSON file from S3, annotate with metadata, and write to an Iceberg or Parquet dataset. The file is written to a Glue table named "{table_prefix}_{node}", where node is taken from the S3 tags. Steps: - Reads the file and its S3 tags/head metadata. - Annotates with ingest metadata and S3 tags. - Computes a row hash for deduplication/auditing. - - Writes to an Iceberg table in Glue/Athena under table "{table_prefix}_{node}". + - Writes to an Iceberg or Parquet table in Glue/Athena under table "{table_prefix}_{node}". Parameters ---------- @@ -436,12 +454,29 @@ def ingest_table_to_parquet_dataset( Timestamp for when the ingest was received. If None, will use current UTC time. ingest_run_id : str, optional Unique ID for this ingest run. If None, will generate a new one. + table_format : str, optional + Table format to write: "iceberg" or "parquet". Defaults to "iceberg". + dataset_root : str, optional + S3 URI for the root of the Parquet dataset (e.g., 's3://bucket/prefix/'). + Required when table_format is "parquet"; ignored for "iceberg". + mode : str, optional + Write mode for the Parquet dataset: "append", "overwrite", or + "overwrite_partitions". Defaults to "append". Ignored for "iceberg". Returns ------- dict Summary of the ingest run, including run ID, file count, and tables/partitions written. """ + if table_format not in ("iceberg", "parquet"): + logger.error(f"Invalid table_format: {table_format!r}") + raise ValueError( + f"table_format must be 'iceberg' or 'parquet', got {table_format!r}" + ) + if table_format == "parquet" and not dataset_root: + logger.error("dataset_root is required when table_format='parquet'") + raise ValueError("dataset_root is required when table_format='parquet'") + if ingest_received_at is None: aest_tz = pytz.timezone("Australia/Melbourne") ingest_received_at = datetime.now(aest_tz).strftime("%Y-%m-%dT%H:%M:%S%z") @@ -518,19 +553,31 @@ def ingest_table_to_parquet_dataset( logger.error(f"Failed to prepare ingest metadata for file {uri}: {e}") raise RuntimeError(f"Failed to prepare ingest metadata for file {uri}: {e}") - # Write this file's DataFrame to its Iceberg table + # Write this file's DataFrame to its Iceberg or Parquet table try: - write_iceberg_to_db( - df=annotated_df, - database=database, - table=table_name, - athena_s3_output=athena_s3_output, - workgroup=workgroup, - ) - logger.debug(f"Successfully wrote {len(annotated_df)} rows to Iceberg table {database}.{table_name}") + if table_format == "iceberg": + write_iceberg_to_db( + df=annotated_df, + database=database, + table=table_name, + athena_s3_output=athena_s3_output, + workgroup=workgroup, + ) + else: + write_parquet_to_db( + df=annotated_df, + dataset_root=dataset_root, + database=database, + table=table_name, + partition_cols=["study_id", "submission_date", "ingest_file_name"], + compression="snappy", + mode=mode, + schema_evolution=True, + ) + logger.debug(f"Successfully wrote {len(annotated_df)} rows to {table_format} table {database}.{table_name}") except Exception as e: - logger.error(f"Failed to write DataFrame to Iceberg table {table_name}: {e}") - raise RuntimeError(f"Failed to write DataFrame to Iceberg table {table_name}: {e}") + logger.error(f"Failed to write DataFrame to {table_format} table {table_name}: {e}") + raise RuntimeError(f"Failed to write DataFrame to {table_format} table {table_name}: {e}") # Track partitions written for this table try: @@ -567,17 +614,20 @@ def ingest_table_to_parquet_dataset( "results": results } -def ingest_files_to_parquet_dataset( +def ingest_files_to_dataset( s3_uris: list, database: str, table_prefix: str, athena_s3_output: str, workgroup: str = "primary", ingest_submission_id: str = None, - exclude_fn: list = ['program.json', 'project.json'], + exclude_fn: list = None, + table_format: str = "iceberg", + dataset_root: str = None, + mode: str = "append", ): """ - Ingest multiple files from S3, annotate with metadata, and write to Iceberg tables in Glue/Athena. + Ingest multiple files from S3, annotate with metadata, and write to Iceberg or Parquet tables in Glue/Athena. Each time this function is called, a single ingestion ID is created which will be attached to all the files in the list of s3_uris. @@ -597,15 +647,26 @@ def ingest_files_to_parquet_dataset( Submission ID for ingest metadata. Defaults to None. exclude_fn : list, optional List of file names from the s3_uris to exclude from ingestion. + Defaults to ['program.json', 'project.json']. Example: ['program.json', 'project.json'] or ['randomDatafile.csv'] + table_format : str, optional + Table format to write: "iceberg" or "parquet". Defaults to "iceberg". + dataset_root : str, optional + S3 URI for the root of the Parquet dataset (e.g., 's3://bucket/prefix/'). + Required when table_format is "parquet"; ignored for "iceberg". + mode : str, optional + Write mode for the Parquet dataset: "append", "overwrite", or + "overwrite_partitions". Defaults to "append". Ignored for "iceberg". """ - + if exclude_fn is None: + exclude_fn = ['program.json', 'project.json'] + ingest_run_id = str(uuid.uuid4()) - + # Generate timestamp in Australia/Melbourne timezone aest_tz = pytz.timezone("Australia/Melbourne") ingest_received_at = datetime.now(aest_tz).strftime("%Y-%m-%dT%H:%M:%S%z") - + # Hardcode the timezone for metadata ingest_timezone = "Australia/Melbourne" @@ -613,7 +674,7 @@ def ingest_files_to_parquet_dataset( for uri in s3_uris: if any(uri.endswith(exclude) for exclude in exclude_fn): continue - resp = ingest_table_to_parquet_dataset( + resp = ingest_table_to_dataset( s3_uri=uri, database=database, table_prefix=table_prefix, @@ -622,8 +683,31 @@ def ingest_files_to_parquet_dataset( ingest_timezone=ingest_timezone, ingest_submission_id=ingest_submission_id, ingest_run_id=ingest_run_id, - ingest_received_at=ingest_received_at + ingest_received_at=ingest_received_at, + table_format=table_format, + dataset_root=dataset_root, + mode=mode, ) results.append(resp) return results + + +def ingest_table_to_parquet_dataset(*args, **kwargs) -> dict: + """Deprecated: use ingest_table_to_dataset instead.""" + warnings.warn( + "ingest_table_to_parquet_dataset is deprecated; " + "use ingest_table_to_dataset", + DeprecationWarning, stacklevel=2, + ) + return ingest_table_to_dataset(*args, **kwargs) + + +def ingest_files_to_parquet_dataset(*args, **kwargs): + """Deprecated: use ingest_files_to_dataset instead.""" + warnings.warn( + "ingest_files_to_parquet_dataset is deprecated; " + "use ingest_files_to_dataset", + DeprecationWarning, stacklevel=2, + ) + return ingest_files_to_dataset(*args, **kwargs) diff --git a/src/g3dt/utils/athena_utils.py b/src/g3dt/utils/athena_utils.py index a8e801e..4810ea2 100644 --- a/src/g3dt/utils/athena_utils.py +++ b/src/g3dt/utils/athena_utils.py @@ -336,6 +336,85 @@ def write_iceberg_to_db( ) +def write_parquet_to_db( + df: pd.DataFrame, + dataset_root: str, + database: str, + table: str, + partition_cols: list = None, + compression: str = "snappy", + mode: str = "append", + schema_evolution: bool = True, + boto3_session=None, +) -> None: + """ + Write a DataFrame to S3 as a Parquet dataset and register/update + the Glue table. + + Parameters + ---------- + df : pd.DataFrame + The DataFrame to write. + dataset_root : str + S3 URI for the root of the Parquet dataset + (e.g., 's3://bucket/prefix/'). + database : str + Glue database name. + table : str + Glue table name. + partition_cols : list, optional + List of columns to partition by. Defaults to + ["study_id", "submission_date"]. + compression : str, optional + Parquet compression codec. Defaults to "snappy". + mode : str, optional + Write mode for Parquet dataset. Defaults to "append". + You can also use "overwrite" to overwrite existing data, + or "overwrite_partitions" to overwrite partitions only. + schema_evolution : bool, optional + Whether to allow schema evolution. Defaults to True. + boto3_session : boto3.Session, optional + A boto3 session. If None, the default session is used. + """ + if partition_cols is None: + partition_cols = ["study_id", "submission_date"] + + try: + logger.debug(f"Creating Glue database '{database}' if not exists.") + wr.catalog.create_database(name=database, exist_ok=True) + logger.debug( + f"Writing DataFrame to Parquet at {dataset_root.rstrip('/')}/ " + f"(table: {table}, database: {database}, " + f"partitions: {partition_cols}, compression: {compression}, " + f"mode: {mode}, schema_evolution: {schema_evolution})" + ) + wr.s3.to_parquet( + df=df.astype("string"), + path=dataset_root.rstrip("/") + "/", + dataset=True, + database=database, + table=table, + partition_cols=partition_cols, + compression=compression, + mode=mode, + schema_evolution=schema_evolution, + boto3_session=boto3_session, + ) + logger.debug( + f"Successfully wrote Parquet dataset to S3 at {dataset_root} " + f"(table: {table}, database: {database})" + ) + except Exception as e: + logger.error( + f"Failed to write Parquet dataset to S3 at {dataset_root} " + f"(table: {table}, database: {database}): {e}" + ) + raise RuntimeError( + f"Failed to write Parquet dataset to S3 at {dataset_root} " + f"(table: {table}, database: {database}): {e}" + ) + + def convert_dataframe_types_for_json(df: pd.DataFrame) -> pd.DataFrame: """ Converts DataFrame column types to JSON-serialisable formats. diff --git a/tests/test_ingest.py b/tests/test_ingest.py index 3bf9e27..d317b06 100644 --- a/tests/test_ingest.py +++ b/tests/test_ingest.py @@ -142,7 +142,19 @@ def test_read_csv_robust_fallback_encoding(self, mock_wrangler): assert mock_wrangler.s3.read_csv.call_count == 2 def test_get_ingest_true_files(self, mock_wrangler): - """Tests filtering of files based on the 'ingest=true' tag.""" + """ + Tests filtering of files based on the 'ingest=true' tag. + + Background: get_ingest_true_files now checks tags concurrently with + a thread pool — the serial one-GetObjectTagging-per-file scan was + the bottleneck on large corpora. Because worker threads can run in + any order, this test asserts on the SET of returned files rather + than the order of get_tags calls. + + Input: 3 listed files — one tagged ingest=true, one ingest=false, + one untagged. + Expected: only the ingest=true file is returned. + """ file_list = ["s3://b/f1.csv", "s3://b/f2.csv", "s3://b/f3.csv"] mock_wrangler.s3.list_objects.return_value = file_list @@ -153,11 +165,65 @@ def get_tags_side_effect(uri): if uri == "s3://b/f2.csv": return {'ingest': 'false'} return {} # f3 has no tags - + with patch('g3dt.ingest.ingest.get_tags', side_effect=get_tags_side_effect): ingest_files = ingest_module.get_ingest_true_files("s3://b/") - - assert ingest_files == ["s3://b/f1.csv"] + + assert set(ingest_files) == {"s3://b/f1.csv"} + + def test_get_ingest_true_files_preserves_listing_order(self, mock_wrangler): + """ + Tests that concurrent tag lookups don't reorder the result. + + get_ingest_true_files checks the 'ingest' tag of each file with a + thread pool (one GetObjectTagging round-trip per file was the + bottleneck on large buckets). Threads can finish in any order, so + this test feeds in 20 files where every second one is tagged + ingest=true and asserts the selected files come back in the exact + order they were listed. Downstream, partition tracking and ingest + summaries assume a deterministic file order. + + Input: 20 listed files, even-numbered ones tagged ingest=true. + Expected: exactly the even-numbered files, in listing order. + """ + file_list = [f"s3://b/f{i}.csv" for i in range(20)] + mock_wrangler.s3.list_objects.return_value = file_list + + def get_tags_side_effect(uri): + index = int(uri.removeprefix("s3://b/f").removesuffix(".csv")) + return {'ingest': 'true'} if index % 2 == 0 else {'ingest': 'false'} + + with patch('g3dt.ingest.ingest.get_tags', side_effect=get_tags_side_effect): + ingest_files = ingest_module.get_ingest_true_files("s3://b/") + + assert ingest_files == [f"s3://b/f{i}.csv" for i in range(0, 20, 2)] + + def test_get_ingest_true_files_skips_file_when_tag_lookup_fails(self, mock_wrangler): + """ + Tests that one failing tag lookup doesn't lose the other files. + + Tag lookups run in worker threads, so an exception for one file + must be caught inside the worker: the file is skipped with a + warning while every other file is still checked and returned. + Without this isolation, a single deleted or permission-restricted + object would abort the whole bucket scan. + + Input: 3 listed files; the tag lookup for the middle one raises. + Expected: the two healthy ingest=true files are returned in order, + the failing file is simply absent. + """ + file_list = ["s3://b/ok1.csv", "s3://b/broken.csv", "s3://b/ok2.csv"] + mock_wrangler.s3.list_objects.return_value = file_list + + def get_tags_side_effect(uri): + if uri == "s3://b/broken.csv": + raise RuntimeError("access denied") + return {'ingest': 'true'} + + with patch('g3dt.ingest.ingest.get_tags', side_effect=get_tags_side_effect): + ingest_files = ingest_module.get_ingest_true_files("s3://b/") + + assert ingest_files == ["s3://b/ok1.csv", "s3://b/ok2.csv"] # --- Tests for Main Ingest Logic --- @@ -208,23 +274,36 @@ def test_align_and_combine_frames(self): assert combined.loc[0, 'c'] == "" assert combined.loc[1, 'a'] == "" + @patch('g3dt.ingest.ingest.write_parquet_to_db') @patch('g3dt.ingest.ingest.write_iceberg_to_db') @patch('g3dt.ingest.ingest.prepare_ingest_metadata') @patch('g3dt.ingest.ingest.get_head_meta') @patch('g3dt.ingest.ingest.read_csv_robust') @patch('g3dt.ingest.ingest.get_tags') - def test_ingest_table_to_parquet_dataset( - self, mock_get_tags, mock_read_csv, mock_get_head, mock_prepare, mock_write_iceberg + def test_ingest_table_to_dataset_defaults_to_iceberg( + self, mock_get_tags, mock_read_csv, mock_get_head, mock_prepare, + mock_write_iceberg, mock_write_parquet ): - """End-to-end test of the single-file ingest pipeline.""" + """ + End-to-end test of the single-file ingest pipeline with default settings. + + The `table_format` parameter defaults to "iceberg", so when it is not + supplied the file must be written via write_iceberg_to_db and the + Parquet writer must never be touched. This protects the default + behavior of every existing caller that doesn't pass table_format. + + Input: one CSV tagged with node 'diagnosis' and table_prefix 'test_prefix'. + Expected: write_iceberg_to_db called once for table 'test_prefix_diagnosis', + write_parquet_to_db not called, and a summary dict for 1 file. + """ mock_get_tags.return_value = {'study_id': 's1', 'node': 'diagnosis', 'submission_date': '2025-10-30'} mock_read_csv.return_value = pd.DataFrame({'snomed': ['123']}) mock_get_head.return_value = {} mock_prepare.return_value = pd.DataFrame({ 'snomed': ['123'], 'study_id': ['s1'], 'submission_date': ['2025-10-30'], 'ingest_file_name': ['diag.csv'] }) - - result = ingest_module.ingest_table_to_parquet_dataset( + + result = ingest_module.ingest_table_to_dataset( s3_uri="s3://b/diag.csv", database="test_db", table_prefix="test_prefix", @@ -236,11 +315,174 @@ def test_ingest_table_to_parquet_dataset( mock_write_iceberg.assert_called_once() call_args, call_kwargs = mock_write_iceberg.call_args assert call_kwargs['table'] == expected_table_name + mock_write_parquet.assert_not_called() # Assert the result dictionary has the expected structure assert result['files_processed'] == 1 assert result['tables_written'] == [expected_table_name] + @patch('g3dt.ingest.ingest.write_parquet_to_db') + @patch('g3dt.ingest.ingest.write_iceberg_to_db') + @patch('g3dt.ingest.ingest.prepare_ingest_metadata') + @patch('g3dt.ingest.ingest.get_head_meta') + @patch('g3dt.ingest.ingest.read_csv_robust') + @patch('g3dt.ingest.ingest.get_tags') + def test_ingest_table_to_dataset_parquet_format( + self, mock_get_tags, mock_read_csv, mock_get_head, mock_prepare, + mock_write_iceberg, mock_write_parquet + ): + """ + Tests the Parquet write path selected via table_format="parquet". + + Legacy bronze tables (created before the Iceberg migration) are plain + Parquet Glue tables; wr.athena.to_iceberg cannot write into them and + fails with a confusing "Schema change detected" error. This flag lets + an ingest keep appending to those tables, so the test pins down the + exact write call the legacy path relied on: the caller-supplied + dataset_root, mode, and the three-level partitioning + (study_id / submission_date / ingest_file_name). + + Input: one CSV ingested with table_format="parquet", + dataset_root="s3://legacy-bronze/", and mode="append". + Expected: write_parquet_to_db called once with those values and the + legacy partition columns; write_iceberg_to_db not called. + """ + mock_get_tags.return_value = {'study_id': 's1', 'node': 'diagnosis', 'submission_date': '2025-10-30'} + mock_read_csv.return_value = pd.DataFrame({'snomed': ['123']}) + mock_get_head.return_value = {} + mock_prepare.return_value = pd.DataFrame({ + 'snomed': ['123'], 'study_id': ['s1'], 'submission_date': ['2025-10-30'], 'ingest_file_name': ['diag.csv'] + }) + + result = ingest_module.ingest_table_to_dataset( + s3_uri="s3://b/diag.csv", + database="test_db", + table_prefix="test_prefix", + athena_s3_output="s3://db-root/output/", + table_format="parquet", + dataset_root="s3://legacy-bronze/", + mode="append", + ) + + mock_write_parquet.assert_called_once() + call_args, call_kwargs = mock_write_parquet.call_args + assert call_kwargs['table'] == "test_prefix_diagnosis" + assert call_kwargs['dataset_root'] == "s3://legacy-bronze/" + assert call_kwargs['mode'] == "append" + assert call_kwargs['partition_cols'] == ["study_id", "submission_date", "ingest_file_name"] + mock_write_iceberg.assert_not_called() + + assert result['files_processed'] == 1 + + def test_ingest_table_to_dataset_parquet_requires_dataset_root(self): + """ + Tests that table_format="parquet" without a dataset_root fails fast. + + A Parquet dataset write needs an S3 root path to place the files + under (Iceberg tables carry their own location in the Glue catalog, + so they don't). Failing before any S3/Glue work happens gives the + operator a clear message instead of a partial ingest. + + Input: table_format="parquet" with dataset_root left as None. + Expected: ValueError mentioning dataset_root, raised before any + AWS call is attempted. + """ + with pytest.raises(ValueError, match="dataset_root is required"): + ingest_module.ingest_table_to_dataset( + s3_uri="s3://b/diag.csv", + database="test_db", + table_prefix="test_prefix", + athena_s3_output="s3://db-root/output/", + table_format="parquet", + ) + + def test_ingest_table_to_dataset_rejects_unknown_format(self): + """ + Tests that an unsupported table_format is rejected up front. + + Only "iceberg" and "parquet" writers exist. A typo like "delta" + should raise immediately rather than silently falling through to + one of the writers. + + Input: table_format="delta". + Expected: ValueError naming the allowed formats. + """ + with pytest.raises(ValueError, match="table_format must be 'iceberg' or 'parquet'"): + ingest_module.ingest_table_to_dataset( + s3_uri="s3://b/diag.csv", + database="test_db", + table_prefix="test_prefix", + athena_s3_output="s3://db-root/output/", + table_format="delta", + ) + + @patch('g3dt.ingest.ingest.ingest_table_to_dataset') + def test_ingest_table_to_parquet_dataset_alias_is_deprecated(self, mock_new_fn): + """ + Tests the backwards-compatibility alias for the renamed single-file + function. + + ingest_table_to_parquet_dataset was renamed to ingest_table_to_dataset + when the table_format flag was added (the old name implied Parquet + output, but it always wrote Iceberg). Existing callers still import + the old name, so it must keep working: it should warn with + DeprecationWarning and delegate all arguments to the new function. + + Input: a call to the old name with keyword arguments. + Expected: DeprecationWarning emitted, new function called once with + the same arguments, and its return value passed through. + """ + mock_new_fn.return_value = {"files_processed": 1} + + with pytest.warns(DeprecationWarning, match="ingest_table_to_parquet_dataset is deprecated"): + result = ingest_module.ingest_table_to_parquet_dataset( + s3_uri="s3://b/diag.csv", + database="test_db", + table_prefix="test_prefix", + athena_s3_output="s3://db-root/output/", + ) + + mock_new_fn.assert_called_once_with( + s3_uri="s3://b/diag.csv", + database="test_db", + table_prefix="test_prefix", + athena_s3_output="s3://db-root/output/", + ) + assert result == {"files_processed": 1} + + @patch('g3dt.ingest.ingest.ingest_files_to_dataset') + def test_ingest_files_to_parquet_dataset_alias_is_deprecated(self, mock_new_fn): + """ + Tests the backwards-compatibility alias for the renamed function. + + ingest_files_to_parquet_dataset was renamed to ingest_files_to_dataset + when the table_format flag was added (the old name implied Parquet + output, but the default is now Iceberg). Existing Glue job scripts + still import the old name, so it must keep working: it should warn + with DeprecationWarning and delegate all arguments to the new function. + + Input: a call to the old name with a keyword argument. + Expected: DeprecationWarning emitted, new function called once with + the same arguments, and its return value passed through. + """ + mock_new_fn.return_value = [{"files_processed": 1}] + + with pytest.warns(DeprecationWarning, match="ingest_files_to_parquet_dataset is deprecated"): + result = ingest_module.ingest_files_to_parquet_dataset( + s3_uris=["s3://b/diag.csv"], + database="test_db", + table_prefix="test_prefix", + athena_s3_output="s3://db-root/output/", + ) + + mock_new_fn.assert_called_once_with( + s3_uris=["s3://b/diag.csv"], + database="test_db", + table_prefix="test_prefix", + athena_s3_output="s3://db-root/output/", + ) + assert result == [{"files_processed": 1}] + # --- Tests for XLSX Reading ---