"""
================================================================================
Athena & S3 Management Operators
================================================================================
This module provides custom Airflow operators for managing Athena (Glue Data
Catalog) tables and their backing S3 data, using **awswrangler** and **boto3**.
These operators assume **instance/role credentials** (no explicit keys).
Operators
---------
1) DeleteAthenaTableOperator
-------------------------
- Validates that a Glue/Athena table exists.
- If `delete_s3_objects=False`: drops the Glue table only.
- If `delete_s3_objects=True`: discovers the table's S3 location, drops the table, then deletes all S3 objects under that prefix (recursive).
Safety:
- Fails fast if the table does not exist.
- Tolerates tables with no discovered S3 location when `delete_s3_objects=True` (will still drop the table and log a warning about missing location).
Typical use:
- Cleanup CTAS outputs, ephemeral staging tables, or decommissioned datasets.
Planned Operators (future in this module)
-----------------------------------------
- AthenaCTASOperator (create via SELECT with optional compression/partitioning) (currently in ctas module)
- AthenaRefreshPartitionsOperator (MSCK REPAIR / crawler-less refresh)
- AthenaCreateExternalTableOperator (DDL in Glue)
- AthenaRepairTableLocationsOperator (normalize or re-point S3 prefixes)
Usage Patterns
--------------
- Single table cleanup::
DeleteAthenaTableOperator(
task_id="drop_users_daily",
database="shoreline",
table="users_daily",
delete_s3_objects=True,
)
- Map across many tables (TaskFlow)::
@task
def list_to_drop():
return [
{"database": "shoreline", "table": "t1", "delete_s3_objects": True},
{"database": "shoreline", "table": "t2", "delete_s3_objects": False},
]
DeleteAthenaTableOperator.partial(task_id="drop_batch").expand_kwargs(list_to_drop())
Notes
-----
- Deleting S3 data is **irreversible**. Use `dry_run=True` to preview.
- If your tables use special formats (e.g., Iceberg/Hudi), the *primary* location
is discovered from Glue's StorageDescriptor.Location. Additional data/metadata
paths (like separate metadata dirs) aren’t removed unless they live under the
same prefix. Extend helpers if you need format-aware cleanup.
- Requires awswrangler >= 3.x.
================================================================================
"""
import logging
from functools import lru_cache
from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
from config import Config
from airflow.exceptions import AirflowException
from airflow.models.baseoperator import BaseOperator
from util import Util
####################################################################################################
## Lazy Load Imports for Heavy Libraries
####################################################################################################
@lru_cache(maxsize=None)
def _import_libs():
"""
Lazy import of heavyweight libraries (awswrangler, boto3, botocore).
Ensures they're only loaded when an operator executes, not when the DAG parses.
Returns a tuple (wr, boto3, botocore).
"""
try:
import importlib
wr = importlib.import_module("awswrangler")
boto3 = importlib.import_module("boto3")
botocore = importlib.import_module("botocore")
return wr, boto3, botocore
except ModuleNotFoundError as e:
raise AirflowException(
f"Required AWS library missing: {e.name}. "
"Ensure boto3, botocore, and awswrangler are installed in the Airflow environment."
)
####################################################################################################
## Helpers
####################################################################################################
# Panic Handling for S3 Deletes
def _guard_panic_prefix(prefix_url: str, *, logger=None) -> None:
"""
Refuse deletion when the prefix is unsafe:
- bucket root (e.g., s3://my-bucket/ or s3://my-bucket)
- exact match (after normalization) to any value in Config.PANIC_PREFIXES
"""
logger = logger or logging.getLogger(__name__)
norm = _norm_prefix(prefix_url)
# 1) Block bucket root deletes
b, k = _parse_s3_url(prefix_url)
if k.strip("/") == "":
raise AirflowException(
f"Refusing to delete bucket root 's3://{b}/' (unsafe). Panic exit."
)
# 2) Block exact matches to configured panic list
# PANIC_PREFIXES holds bare prefix names ("curated", "raw", ...) — compare
# against the key path, not a normalized s3:// URL
panic_values = {
x.strip("/")
for x in getattr(Config, "PANIC_PREFIXES", [])
if isinstance(x, str)
}
if k.strip("/") in panic_values:
raise AirflowException(
f"Refusing to clean up unsafe root prefix '{norm}' (matches Config.PANIC_PREFIXES). Panic exit."
)
logger.info("Prefix passed panic guard: %s", norm)
# S3 URL Normalization
def _norm_prefix(url: str) -> str:
"""
Canonicalize an S3 URL to 's3://bucket/prefix/' (always trailing slash; empty key -> '/').
"""
bucket, key = _parse_s3_url(url)
key = _ensure_trailing_slash(key)
return f"s3://{bucket}/{key}"
def _ensure_trailing_slash(prefix: str) -> str:
"""Return prefix that always ends with '/'."""
return prefix if prefix.endswith("/") else prefix + "/"
def _parse_s3_url(url: str) -> Tuple[str, str]:
"""
Parse s3://bucket/prefix -> (bucket, prefix). Raises on invalid URLs.
"""
if not url or not url.startswith("s3://"):
raise AirflowException(f"Invalid or missing S3 URL: {url!r}")
path = url[len("s3://") :]
if "/" not in path:
return path, ""
bucket, key = path.split("/", 1)
return bucket, key
# Glue Activities
def _glue_get_table(
database: str,
table: str,
*,
glue=None,
botocore_mod=None,
region_name: Optional[str] = "us-east-1",
) -> dict:
if glue is None or botocore_mod is None:
wr, boto3_mod, botocore_mod = _import_libs()
glue = glue or boto3_mod.client("glue", region_name=region_name)
try:
resp = glue.get_table(DatabaseName=database, Name=table)
return resp["Table"]
except glue.exceptions.EntityNotFoundException:
raise AirflowException(f"Glue/Athena table not found: {database}.{table}")
except botocore_mod.exceptions.BotoCoreError as e:
raise AirflowException(f"Glue get_table failed: {e}") from e
def _glue_table_exists(
database: str,
table: str,
*,
glue=None,
botocore_mod=None,
region_name: Optional[str] = "us-east-1",
) -> bool:
if glue is None or botocore_mod is None:
_, boto3_mod, botocore_mod = _import_libs()
glue = glue or boto3_mod.client("glue", region_name=region_name)
try:
glue.get_table(DatabaseName=database, Name=table)
return True
except glue.exceptions.EntityNotFoundException:
return False
except botocore_mod.exceptions.BotoCoreError as e:
raise AirflowException(f"Glue exists check failed: {e}") from e
def _glue_delete_table(
database: str,
table: str,
*,
glue=None,
botocore_mod=None,
region_name: Optional[str] = "us-east-1",
) -> None:
if glue is None or botocore_mod is None:
_, boto3_mod, botocore_mod = _import_libs()
glue = glue or boto3_mod.client("glue", region_name=region_name)
try:
glue.delete_table(DatabaseName=database, Name=table)
except glue.exceptions.EntityNotFoundException:
pass
except botocore_mod.exceptions.BotoCoreError as e:
raise AirflowException(f"Glue delete_table failed: {e}") from e
def _discover_table_location(table_dict: dict) -> Optional[str]:
"""
Extract StorageDescriptor.Location from Glue table dict if available.
"""
try:
loc = (table_dict.get("StorageDescriptor") or {}).get("Location")
return loc or None
except Exception:
return None
# Delete S3 Prefix
def _delete_s3_prefix(
prefix_url: str, *, dry_run: bool = False, logger=None, wr_mod=None, boto3_mod=None
) -> int:
if wr_mod is None or boto3_mod is None:
wr_mod, boto3_mod, _ = _import_libs()
logger = logger or logging.getLogger(__name__)
# Make sure we don't delete anything dangerous
_guard_panic_prefix(prefix_url, logger=logger)
bucket, key = _parse_s3_url(prefix_url)
key = _ensure_trailing_slash(key)
if dry_run:
s3 = boto3_mod.client("s3")
paginator = s3.get_paginator("list_objects_v2")
total = 0
for page in paginator.paginate(Bucket=bucket, Prefix=key):
total += len(page.get("Contents", []))
logger.info(
"[dry_run] Would delete %d object(s) under s3://%s/%s", total, bucket, key
)
return total
result = wr_mod.s3.delete_objects(f"s3://{bucket}/{key}")
# Normalize to an integer count
if isinstance(result, int):
deleted = result
elif isinstance(result, dict):
# Just in case a dict slips through (boto3-style), try to count
deleted = len(result.get("Deleted", [])) if "Deleted" in result else 0
else:
deleted = 0 # Fallback
# TODO TO FIX: deleted count (more than 0) but showed 0 in logging!
logger.info("Deleted %s object(s) under s3://%s/%s", deleted, bucket, key)
return int(deleted)
####################################################################################################
## Ingestion Helpers (URL + S3 listing)
####################################################################################################
def _read_from_url(input_file_path: str, *, pd_mod) -> "pd.DataFrame":
"""
Read a file from a non-S3 URL into a pandas DataFrame, matching the legacy DAG behavior.
"""
if input_file_path.endswith(".csv"):
return pd_mod.read_csv(input_file_path, dtype=str)
elif input_file_path.endswith(".xls") or input_file_path.endswith(".xlsx"):
return pd_mod.read_excel(input_file_path, dtype=str)
else:
raise AirflowException(f"File type not supported yet: {input_file_path!r}")
def _list_s3_files(
bucket_name: str,
prefix: str,
*,
boto3_mod,
logger: logging.Logger,
filter_fun,
):
"""
List all files in an S3 bucket with a given prefix and return a sorted list (latest first),
mirroring the original ingest_raw_file DAG implementation.
"""
s3_client = boto3_mod.client("s3")
try:
paginator = s3_client.get_paginator("list_objects_v2")
pages = paginator.paginate(Bucket=bucket_name, Prefix=prefix)
data = []
for page in pages:
if "Contents" in page:
files_filtered = filter(filter_fun, page["Contents"])
data.extend(files_filtered)
files_filtered_sorted = sorted(
data,
key=lambda obj: int(obj["LastModified"].strftime("%s")),
reverse=True,
)
logger.info(
"Found %d file(s) with prefix %s", len(files_filtered_sorted), prefix
)
return files_filtered_sorted
except Exception as e:
logger.error("Error listing S3 files: %s", str(e))
raise AirflowException(
f"Error listing S3 files under {bucket_name}/{prefix}: {e}"
) from e
####################################################################################################
## DeleteAthenaTableOperator
####################################################################################################
[docs]
class DeleteAthenaTableOperator(BaseOperator):
"""
Drop a Glue/Athena table and (optionally) delete its S3 data prefix.
Behavior
--------
1) Validates that the table exists in the Glue Data Catalog
- If not found and `succeed_if_missing=False`: raises an error.
- If not found and `succeed_if_missing=True`: logs and exits successfully
2) If `delete_s3_objects=True`: discovers the S3 location (StorageDescriptor.Location),
drops the table, then recursively deletes all objects in that prefix.
3) If `delete_s3_objects=False`: drops the table only.
Parameters
----------
:param database: Glue database (Athena schema).
:type database: str
:param table: Glue/Athena table name.
:type table: str
:param delete_s3_objects: If True, also delete the S3 location recursively.
:type delete_s3_objects: bool
:param dry_run: If True, do not modify anything; log intended actions.
:type dry_run: bool
:param region_name: Optional AWS region (boto3/awslambda resolves from environment/instance profile if not set).
:type region_name: Optional[str]
:param succeed_if_missing: If True, do not fail if the table is missing; just log a warning.
:type succeed_if_missing: bool
Template fields
---------------
- database
- table
- delete_s3_objects
- dry_run
- succeed_if_missing
Returns
-------
dict
{
"database": "...",
"table": "...",
"existed": True,
"deleted_catalog": True,
"s3_prefix": "s3://bucket/prefix/" | None,
"deleted_objects": <int> | 0,
"dry_run": <bool>,
"succeed_if_missing": <bool>
}
Guide for Use (DAG Mapping)
---------------------------
- Single table cleanup:
DeleteAthenaTableOperator(
task_id="drop_users_daily",
database="shoreline",
table="users_daily",
delete_s3_objects=True,
)
- Map across many tables (TaskFlow ):
@task
def list_to_drop():
return [
{"database": "shoreline", "table": "t1", "delete_s3_objects": True},
{"database": "shoreline", "table": "t2", "delete_s3_objects": False},
]
DeleteAthenaTableOperator.partial(task_id="drop_batch").expand_kwargs(list_to_drop())
"""
ui_color = "#B00B00"
template_fields: Sequence[str] = (
"database",
"table",
"delete_s3_objects",
"dry_run",
"succeed_if_missing",
)
def __init__(
self,
*,
database: str = Config.DATA_LAKE_DB,
table: str,
delete_s3_objects: bool = False,
dry_run: bool = False,
succeed_if_missing: bool = False,
region_name: Optional[str] = "us-east-1",
**kwargs,
):
super().__init__(**kwargs)
self.database = database
self.table = table
# keep raw values: these are template_fields, and coercing here (e.g. bool())
# would collapse a "{{ ... }}" string to True before Airflow renders it
self.delete_s3_objects = delete_s3_objects
self.dry_run = dry_run
self.succeed_if_missing = succeed_if_missing
self.region_name = region_name
################################################################################################
## Helpers
################################################################################################
@staticmethod
def _to_bool(value) -> bool:
# rendered template_fields can arrive as strings when a DAG does not use
# render_template_as_native_obj; bool("False") would be True
if isinstance(value, str):
return value.strip().lower() in ("true", "1", "t", "yes", "y")
return bool(value)
[docs]
def execute(self, context):
# template_fields are rendered by now — safe to coerce
self.delete_s3_objects = self._to_bool(self.delete_s3_objects)
self.dry_run = self._to_bool(self.dry_run)
self.succeed_if_missing = self._to_bool(self.succeed_if_missing)
# Create clients at runtime
# Each function will fail over to lazy import if needed
wr, boto3_mod, botocore_mod = _import_libs()
glue = (
boto3_mod.client("glue", region_name=self.region_name)
if self.region_name
else boto3_mod.client("glue")
)
self.log.info(
"Preparing to drop Glue/Athena table: %s.%s (delete_s3=%s, dry_run=%s)",
self.database,
self.table,
self.delete_s3_objects,
self.dry_run,
)
# 1) Validate exists (fail or succeed-noop if configured)
exists = _glue_table_exists(
self.database,
self.table,
glue=glue,
botocore_mod=botocore_mod,
region_name=self.region_name,
)
if not exists:
msg = f"Table does not exist: {self.database}.{self.table}"
if self.succeed_if_missing:
self.log.info(
"%s — succeed_if_missing=True, treating as successful no-op.", msg
)
return {
"database": self.database,
"table": self.table,
"existed": False,
"deleted_catalog": False,
"s3_prefix": None,
"deleted_objects": 0,
"dry_run": self.dry_run,
"succeed_if_missing": self.succeed_if_missing,
}
raise AirflowException(msg)
# 2) Discover S3 location (if we might need it)
s3_prefix = None
if self.delete_s3_objects:
tdict = _glue_get_table(
self.database,
self.table,
glue=glue,
botocore_mod=botocore_mod,
region_name=self.region_name,
)
s3_prefix = _discover_table_location(tdict)
if s3_prefix:
# normalize for logging
try:
b, k = _parse_s3_url(s3_prefix)
s3_prefix = f"s3://{b}/{_ensure_trailing_slash(k)}"
except Exception:
self.log.warning(
"Table has non-S3 or unparsable location: %r", s3_prefix
)
else:
self.log.warning(
"No StorageDescriptor.Location on %s.%s; will drop table but cannot remove S3 data.",
self.database,
self.table,
)
# 3) If dry_run: just report planned actions
if self.dry_run:
self.log.info(
"[dry_run] Would drop catalog table %s.%s", self.database, self.table
)
deleted_objects = 0
if self.delete_s3_objects and s3_prefix:
deleted_objects = _delete_s3_prefix(
s3_prefix,
dry_run=True,
logger=self.log,
wr_mod=wr,
boto3_mod=boto3_mod,
)
return {
"database": self.database,
"table": self.table,
"existed": True,
"deleted_catalog": False,
"s3_prefix": s3_prefix,
"deleted_objects": deleted_objects,
"dry_run": True,
"succeed_if_missing": self.succeed_if_missing,
}
# 4) Drop the catalog table
self.log.info("Dropping Glue table %s.%s ...", self.database, self.table)
_glue_delete_table(
self.database,
self.table,
glue=glue,
botocore_mod=botocore_mod,
region_name=self.region_name,
)
# 5) Optionally delete S3 data
deleted_objects = 0
if self.delete_s3_objects and s3_prefix:
self.log.info("Deleting S3 objects under %s ...", s3_prefix)
deleted_objects = _delete_s3_prefix(
s3_prefix,
dry_run=False,
logger=self.log,
wr_mod=wr,
boto3_mod=boto3_mod,
)
self.log.info(
"Completed delete for %s.%s (deleted_catalog=True, deleted_objects=%d).",
self.database,
self.table,
deleted_objects,
)
return {
"database": self.database,
"table": self.table,
"existed": True,
"deleted_catalog": True,
"s3_prefix": s3_prefix,
"deleted_objects": deleted_objects,
"dry_run": False,
"succeed_if_missing": self.succeed_if_missing,
}
####################################################################################################
## S3ToAthenaParquetOperator
####################################################################################################
[docs]
class S3ToAthenaParquetOperator(BaseOperator):
"""
=========================
S3ToAthenaParquetOperator
=========================
Ingest raw data from S3 (or an external HTTP/HTTPS source) into the Ready data
lake as partitioned Parquet files and register the resulting Glue/Athena tables.
This operator is the BaseOperator equivalent of the historical
`ingest_raw_file` TaskFlow DAG. It provides a controlled, extensible, and
map-friendly ingestion mechanism built entirely on the shared Athena/S3 helper
functions in this module.
Primary Responsibilities
------------------------
1. **S3 Cleanup (safe)**
- Computes the destination folder from the provided `output_path`.
- Uses the shared, panic-protected `_delete_s3_prefix()` helper to remove
existing objects under that prefix.
- Prevents accidental bucket-root deletes and respects `Config.PANIC_PREFIXES`.
2. **File Ingestion**
- Supports three input modes:
* **S3 folder** (prefix ending with "/"):
- Lists all files with extensions (.csv, .xls, .xlsx, .parquet).
- Ingests *each file* into its own Athena table with an auto-generated
suffix derived from the file’s stem.
* **Single S3 file**:
- Ingests into a single Athena table.
* **External URL** (HTTP/HTTPS):
- Reads via pandas (CSV or Excel).
- Writes Parquet to S3 with awswrangler.
- Recreates a Glue/Athena table for the output data.
3. **Table Creation & Glue Catalog Management**
- Creates tables in the configured Athena/Glue database (`Config.DATA_LAKE_DB`).
- Automatically builds canonical table names:
r_{state_abbr}up_{output_table}
or (for folder ingestion):
r_{state_abbr}up_{output_table}_{file_stem}
- Optionally drops the base table before ingestion (`drop_table_if_exists=True`).
4. **Row Count Verification (non-S3 URL ingestion)**
- When ingesting via HTTP/HTTPS, performs an Athena query to confirm the
number of imported rows, logging results for transparency.
Safety & Consistency
--------------------
- All S3 deletion uses the shared `_delete_s3_prefix()` helper, ensuring:
* guard-rails against bucket-root or panic prefixes,
* consistent logging and dry-run capability.
- Uses lazy imports via `_import_libs()` for awswrangler/boto3 to remain DAG-parse safe.
- Inherits Airflow template fields for all user-provided paths, enabling
`.expand_kwargs()` and dynamic ingestion.
Typical Usage
-------------
S3ToAthenaParquetOperator(
task_id="ingest_raw",
database=Config.DATA_LAKE_DB,
bucket=Config.DATA_LAKE_BUCKET,
output_table="proposal_uploads",
state_abbr="va",
output_path="s3://{DEST_BUCKET}/raw/uploads/Virginia/unzipped",
input_path="s3://463002504035-shoreline/uploads/VA/proposals/2025-09-11/",
)
Ideal For
---------
- Ad-hoc ingestion of CSV/XLSX/Parquet into the data lake.
- Bulk-folder ingestion where each file becomes a separate table.
- Automated ingestion pipelines where inputs vary by tenant, state, or date.
- Eliminating legacy TaskFlow DAG logic while preserving full functionality.
Returns
-------
dict containing:
{
"database": "...",
"bucket": "...",
"output_path": "s3://bucket/path",
"input_path": "...",
"tables_created": [
{"table": "...", "source": "...", "prefix": "..."},
...
],
}
"""
ui_color = "#006BB3"
template_fields: Sequence[str] = (
"database",
"bucket",
"output_table",
"state_abbr",
"output_path",
"input_path",
)
def __init__(
self,
*,
database: str = Config.DATA_LAKE_DB,
bucket: str = Config.DATA_LAKE_BUCKET,
output_table: str,
state_abbr: str,
output_path: str,
input_path: str,
drop_table_if_exists: bool = True,
region_name: Optional[str] = "us-east-1",
column_mapping: Optional[Dict[str, Union[str, List[str]]]] = None,
**kwargs: Any,
):
super().__init__(**kwargs)
self.database = database
self.bucket = bucket
self.output_table = output_table
self.state_abbr = state_abbr
self.output_path = output_path
self.input_path = input_path
self.drop_table_if_exists = bool(drop_table_if_exists)
self.region_name = region_name
self.column_mapping = column_mapping
[docs]
def execute(self, context):
# Lazy imports for heavy libs
wr, boto3_mod, _ = _import_libs()
# Import pandas lazily as well
import importlib
pd_mod = importlib.import_module("pandas")
database = self.database or Config.DATA_LAKE_DB
bucket = self.bucket or Config.DATA_LAKE_BUCKET
# Format the output path to resolve {DEST_BUCKET} / {BUCKET} placeholders
output_s3_path = (self.output_path or "").format(
DEST_BUCKET=bucket,
BUCKET=bucket,
)
if not output_s3_path.startswith("s3://"):
raise AirflowException(
f"output_path must be an S3 URL starting with 's3://', got: {output_s3_path!r}"
)
# Normalize trailing slash handling
if output_s3_path.endswith("/"):
output_s3_path = output_s3_path[:-1]
base_output_table_name = self.output_table
state_abbr = self.state_abbr
base_athena_table = str.lower(f"r_{state_abbr}up_{base_output_table_name}")
# Compute the raw folder (prefix within bucket) and delete it first,
# via the shared panic-safe helper.
raw_table_folder = output_s3_path.replace(f"s3://{bucket}/", "")
cleanup_prefix = f"s3://{bucket}/{raw_table_folder}"
self.log.info(
"Cleaning up existing S3 data under %s",
cleanup_prefix,
)
try:
deleted = _delete_s3_prefix(
cleanup_prefix,
dry_run=False,
logger=self.log,
wr_mod=wr,
boto3_mod=boto3_mod,
)
self.log.info(
"Deleted %d existing object(s) under %s",
deleted,
cleanup_prefix,
)
except AirflowException as e:
# issue warn
self.log.warning(
"Failed to clean up S3 prefix %s (continuing): %s",
cleanup_prefix,
e,
)
save_path = f"{output_s3_path}/{base_output_table_name}"
tables_created = []
# Drop Glue table if requested (base table only, non-folder case)
if self.drop_table_if_exists:
self.log.info(
"Dropping Glue/Athena table if exists: %s.%s",
database,
base_athena_table,
)
try:
wr.catalog.delete_table_if_exists(
database=database, table=base_athena_table
)
except Exception as e:
self.log.warning(
"Failed to drop existing table %s.%s (continuing): %s",
database,
base_athena_table,
e,
)
input_path = self.input_path
# S3 path ingestion
if input_path.startswith("s3://"):
raw_table_prefix = save_path.replace(f"s3://{bucket}/", "")
# Folder case
if input_path.endswith("/"):
self.log.info(
"Input path is an S3 folder; ingesting all supported files in the folder: %s",
input_path,
)
bucket_name, prefix = _parse_s3_url(input_path)
prefix = _ensure_trailing_slash(prefix)
def _ext_filter(obj):
key = obj["Key"]
return (
key.endswith(".csv")
or key.endswith(".xls")
or key.endswith(".xlsx")
or key.endswith(".parquet")
)
files = _list_s3_files(
bucket_name=bucket_name,
prefix=prefix,
boto3_mod=boto3_mod,
logger=self.log,
filter_fun=_ext_filter,
)
for file_obj in files:
key = file_obj["Key"]
path_to_file = f"s3://{bucket_name}/{key}"
file_stem = key.split("/")[-1].split(".")[0].replace(" ", "_")
athena_table = str.lower(
f"r_{state_abbr}up_{base_output_table_name}_{file_stem}"
)
table_prefix = f"{raw_table_prefix}/{file_stem}"
self.log.info(
"Ingesting file %s to table %s (prefix=%s)...",
path_to_file,
athena_table,
table_prefix,
)
Util.s3_file_to_parquet(
path_to_file,
table_prefix,
athena_table,
column_mapping=self.column_mapping,
)
self.log.info(
"Ingested file %s to table %s", path_to_file, athena_table
)
tables_created.append(
{
"table": athena_table,
"source": path_to_file,
"prefix": table_prefix,
}
)
# Single S3 file case
else:
self.log.info("Input path is a single S3 file: %s", input_path)
athena_table = base_athena_table
Util.s3_file_to_parquet(
input_path,
raw_table_prefix,
athena_table,
column_mapping=self.column_mapping,
)
if base_output_table_name.endswith(".xlsx"):
self.log.info(
"Excel file identified; more than one sheet may be ingested. "
"Final table names may be suffixed with the sheet name."
)
self.log.info(
"Ingested S3 file %s to table %s", input_path, athena_table
)
tables_created.append(
{
"table": athena_table,
"source": input_path,
"prefix": raw_table_prefix,
}
)
# Non-S3 URL ingestion (HTTP/HTTPS)
else:
self.log.info(
"Input path is a non-S3 URL; reading with pandas and writing Parquet via awswrangler: %s",
input_path,
)
df = _read_from_url(input_path, pd_mod=pd_mod)
Util.apply_column_mapping(df, self.column_mapping)
self.log.info(
"Writing DataFrame to Parquet at %s (database=%s, table=%s)",
save_path,
database,
base_athena_table,
)
wr.s3.to_parquet(
df=df,
path=save_path,
database=database,
dataset=True,
index=False,
mode="overwrite",
table=base_athena_table,
dtype={col: "string" for col in df.columns},
)
# Verify row count via Athena
source_df = wr.athena.read_sql_query(
sql=f"SELECT COUNT(*) AS c FROM {base_athena_table}",
database=database,
)
row_count = int(source_df["c"][0])
self.log.info(
"%d row(s) imported into %s.%s",
row_count,
database,
base_athena_table,
)
tables_created.append(
{
"table": base_athena_table,
"source": input_path,
"prefix": save_path.replace(f"s3://{bucket}/", ""),
"rows": row_count,
}
)
result = {
"database": database,
"bucket": bucket,
"output_path": output_s3_path,
"input_path": input_path,
"tables_created": tables_created,
}
self.log.info("S3ToAthenaParquetOperator result: %s", result)
return result