"""Chunked Postgres-to-Athena transfer operator (pandas read_sql chunks ingested via wr.s3.to_parquet)."""
from config import Config
from util import Util
from typing import Literal
from airflow.models import BaseOperator
from airflow.utils.decorators import apply_defaults
from airflow.providers.postgres.hooks.postgres import PostgresHook
from airflow.exceptions import AirflowException
DEFAULT_DB_POSTGRES = Config.BOSS_PLATFORM_GRAPI_DB
DEFAULT_DB_LAKE = Config.DATA_LAKE_DB
DEFAULT_S3_BUCKET = Config.DATA_LAKE_BUCKET
PG_TO_ATHENA_TABLE_PREFIX = "pg"
S3_SAVE_PATH_PREFIX = "postgres/{root}/{partition}"
# IN THE END, wr.s3.to_parquet: MUST USE `append` WRITE MODE because of chunking
WRITE_MODE_CHUNKING = "append"
TMP_LOGDIR = "tmp/log_ingested_{output_table}"
[docs]
class PostgresToDatalake(BaseOperator):
"""
A general operator to transfer data from pg to athena.
This is a revision from PostgresToDatalake which use pandas object, is sensitive to the memory limits, so cant handle very big dataset.
This approach use df=read_sql(chunksize) and ingest each chunk to s3 with wr.s3.to_parquet(dataset=True)
CAUTION #1: Postgis geom col is not supported yet; so geom values might not be properly transferred!
CAUTION #2: ALL QUERIES that save to same athena table will be chunked (offset) using same chunk indexes (stored in airflow /tmp).
So e.g. SELECT * from cb20 limit 100 offset 100 <> SELECT * from cb20 limit 100 offset 200 will be both counted as chunk=0 and the subsequent transfer queries will be skipped!
Args:
source_sql (str): the SELECT query (CTAS) to run against pg. Here, you should cast, rename and manipulate incoming data into desired forms in Athena
table_name (str): the name of the table to generate in athena. This operator will add `pg_` to the table name automatically
mode (Literal['append', 'overwrite', 'overwrite_partitions'], optional): This is a parameter to pass aws wrangler to replace or append the existing table. Defaults to 'overwrite'.
partition_cols (list[str], optional): The partitioning columns. make sure to set this for the performance. Defaults to None.
dtype (dict, optional): A dict to force AWS Wrangler to use specified data types; also cast completely NULL cols to `string`, not specified
workgroup (str, optional): athena workgroup. Defaults to Config.DEFAULT_WORKGROUP.
athena_targt_db (str, optional): athena target db. Defaults to Config.DATA_LAKE_DB.
postgres_source_db (str, optional): pg source db. Defaults to Config.BOSS_PLATFORM_GRAPI_DB.
cleanup_templog (bool, optional): whether to remove all temp log (checkpoints) afterward!
"""
ui_color = "#A6E6A6"
template_fields = ("source_sql", "table_name", "mode")
@apply_defaults
def __init__(
self,
source_sql: str,
table_name: str,
mode: Literal["append", "overwrite", "overwrite_partitions"] = "overwrite",
partition_cols: list[str] = [],
dtype: dict = {},
workgroup: str = Config.DEFAULT_WORKGROUP,
athena_target_db: str = DEFAULT_DB_LAKE,
postgres_source_db: str = DEFAULT_DB_POSTGRES,
chunksize: int = 64000,
cleanup_templog: bool = True,
*args,
**kwargs,
):
super().__init__(*args, **kwargs)
self.source_sql = source_sql
self.table_name = table_name
self.partition_cols = partition_cols
self.workgroup = workgroup
self.athena_target_db = athena_target_db
self.postgres_source_db = postgres_source_db
self.mode = mode
self.dtype = dtype
self.chunksize = chunksize
self.cleanup_templog = cleanup_templog
[docs]
def execute(
self,
context: dict,
) -> None:
"""
main call to transfer data from pg to athena
:return: None
"""
output_table = None
try:
import awswrangler as wr
import pandas as pd
from glob import glob
from pathlib import Path
output_table = f"{PG_TO_ATHENA_TABLE_PREFIX}_{self.table_name}"
athena_target_db = self.athena_target_db
postgres_source_db = self.postgres_source_db
source_sql = self.source_sql
partition_cols = self.partition_cols
mode = self.mode
dtype = self.dtype
chunksize = self.chunksize
eng = PostgresHook(
postgres_conn_id=postgres_source_db
).get_sqlalchemy_engine()
print(f"source_sql: {source_sql}")
source_sql = source_sql.rsplit(";", 1)[
0
] # remove statement ending ";", if any
FULLSIZE = pd.read_sql(
f"""select count(*) as c from
({source_sql})
as sub
""",
eng,
).iloc[0]["c"]
s3_prefix = S3_SAVE_PATH_PREFIX.format(root=output_table, partition="")
s3_full_path = f"s3://{DEFAULT_S3_BUCKET}/{s3_prefix}"
if FULLSIZE == 0:
print("SKIPPED: source_sql returns no results to transfer!")
if mode == "overwrite":
print("NOTICE: WRITE MODE is overwrite with 0 rows: Removing stale table, if exists...")
Util.remove_datalake_table(tgt_table=output_table, tgt_prefix=s3_prefix)
return
print(f"RECORDS TO LOAD: {FULLSIZE}. Now loading...")
print(f"{s3_full_path=}")
checkpoint_logdir = TMP_LOGDIR.format(output_table=output_table)
Path(checkpoint_logdir).mkdir(parents=True, exist_ok=True)
partition_logdir = "{checkpoint_logdir}/overwrite_partitions/{partition}"
CHECKPOINT_INGESTED_CHUNK_IDX_FILE_SUFFIX = "ingested"
if mode == "overwrite":
print(
"NOTICE: WRITE MODE is overwrite: Removing entire table (catalog metadata, and s3 content), if exists..."
)
Util.remove_datalake_table(tgt_table=output_table, tgt_prefix=s3_prefix)
import os
existing_ingested_idx_files = glob(
f"{checkpoint_logdir}/*.{CHECKPOINT_INGESTED_CHUNK_IDX_FILE_SUFFIX}"
)
print(
f"UNLINKING/REMOVING {len(existing_ingested_idx_files)} checkpoints (saved chunk indexes)..."
)
for ingested_idx_file in existing_ingested_idx_files:
os.unlink(ingested_idx_file)
with eng.connect().execution_options(stream_results=True) as connection:
NUM_CHUNKS = FULLSIZE // chunksize + 1
ingested = set(
int(
s.rsplit("/", 1)[-1].rsplit(
f".{CHECKPOINT_INGESTED_CHUNK_IDX_FILE_SUFFIX}", 1
)[0]
)
for s in glob(
f"{checkpoint_logdir}/*.{CHECKPOINT_INGESTED_CHUNK_IDX_FILE_SUFFIX}"
)
)
# only need to start from last un-ingested chunk
min_to_ingest = min(set(range(NUM_CHUNKS)) - ingested or {0})
print(f"{NUM_CHUNKS=} - {len(ingested)=} ::: {min_to_ingest=}")
offset = min_to_ingest * chunksize
query_offset = f"""
SELECT * FROM ({source_sql}) as sub
OFFSET {offset}
"""
print(query_offset)
for idx, c_df in enumerate(
pd.read_sql(query_offset, connection, chunksize=chunksize)
):
# This is the actual chunk index, with offset considered
idx = idx + min_to_ingest
if idx % 5 == 0:
print(
f"Progress: Transfering chunk of size={chunksize}: {idx=}.... "
)
# FIRST RUN: CHECK IF ALL COLS in partition_cols are included in the query!
if idx == 0:
missing_partition_cols = set(partition_cols) - set(c_df)
if missing_partition_cols != set():
raise AirflowException(
f"ALERT: source_query doesn't include all PARTITION COLS: {missing_partition_cols=}!"
)
## First, clear existing PARTITIONS before attempting to overwrite!
if mode == "overwrite_partitions":
# Details about: wr.s3.to_parquet(mode == 'overwrite_partitions')
# Only deletes the paths of partitions that should be updated and then writes the new partitions files. It’s like a 'partition Upsert'....
par_list = []
for vals in (
c_df[partition_cols]
.drop_duplicates(keep="first")
.sort_values(partition_cols)
.itertuples(index=False)
):
pardir = ""
for col, val in zip(partition_cols, vals):
pardir += f"{col}={val}/"
par_list.append(pardir)
for pardir in par_list:
if Path(pardir).is_dir() is False:
# FIRST TIME seeing this partition >> REMOVE/OVERWRITE PARTITION
Path(
partition_logdir.format(
checkpoint_logdir=checkpoint_logdir,
partition=pardir,
)
).mkdir(parents=True, exist_ok=True)
par_s3_prefix = S3_SAVE_PATH_PREFIX.format(
root=output_table, partition=pardir
)
print(
f"Progress: overwriting partition {pardir=} at {par_s3_prefix=}..."
)
Util.delete_s3_folder(
bucket=DEFAULT_S3_BUCKET, prefix=par_s3_prefix
)
if idx not in ingested:
# actual ingestion to Athena + s3
# Convert unsupported column types before parquet serialization
import json
from uuid import UUID
for col in c_df.columns:
non_null = c_df[col].dropna()
if len(non_null) == 0:
continue
# Convert UUID columns to strings (PyArrow cannot serialize UUID objects)
if non_null.apply(type).eq(UUID).any():
c_df[col] = c_df[col].astype(str)
c_df.loc[c_df[col] == "None", col] = None
# Convert any object-dtype column with mixed/complex types to JSON strings
# (PyArrow can't handle mixed types like bytes+dict+int in one column)
# Using json.dumps for all values ensures valid JSON roundtrip for JSONB columns
elif c_df[col].dtype == "object":
types = set(non_null.apply(type))
if types != {str}:
def _serialize(x):
if x is None or (isinstance(x, float) and pd.isna(x)):
return None
if isinstance(x, bytes):
return x.decode("utf-8", errors="replace")
return json.dumps(x, default=str)
c_df[col] = c_df[col].apply(_serialize)
# HAVE TO SET DTYPES FOR EVERY CHUNK INGESTION i.e. if a chunk has a column all empty WILL CAUSE wr.s3.to_parquet error
# e.g.One of chunks in cb table: UndetectedType("Impossible to infer the equivalent Athena data type for the no_computer_per_cbg column. It is completely empty (only null values) and has a too generic data type (object). Please, cast this columns with a more deterministic data type (e.g. df['no_computer_per_cbg'] = df['no_computer_per_cbg'].astype('string')) or pass the column schema as argument(e.g. dtype={'no_computer_per_cbg': 'string'}")
for col in set(c_df):
if (
(col not in dtype)
and c_df[col].isnull().all()
or (
c_df[col].fillna("").astype(str).str.strip() == ""
).all()
):
# EMPTY COL; set to athena string type!
# print(f"{idx=}: {col=} IS ENTIRELY EMPTY!")
dtype[col] = "string"
c_df[col] = c_df[col].astype("string")
wr.s3.to_parquet(
df=c_df,
path=s3_full_path,
dataset=True,
mode=WRITE_MODE_CHUNKING,
compression="snappy",
database=athena_target_db,
table=output_table,
partition_cols=partition_cols,
dtype=dtype,
)
ingested_chunk_log = f"{checkpoint_logdir}/{idx}.{CHECKPOINT_INGESTED_CHUNK_IDX_FILE_SUFFIX}"
print(f"marking chunk {ingested_chunk_log} as ingested")
open(ingested_chunk_log, "w").close()
print(
f"SUCCESSFUL: ALL OF {FULLSIZE} rows are uploaded to data lake, at external table: {output_table}!"
)
except Exception as ExceptionMsg:
raise AirflowException(
f"ALERT: unidentified/uncaught ERROR: {output_table=}: Error transferring data from PG to Athena: {ExceptionMsg=}"
)
finally:
if self.cleanup_templog:
import shutil
checkpoint_logdir = TMP_LOGDIR.format(output_table=output_table)
print(f"Removing all checkpoint/progress logs at: {checkpoint_logdir}")
shutil.rmtree(checkpoint_logdir, ignore_errors=True)
print(
f"{output_table=}: PostgresToDatalake Operator/attempt to transferred data from pg to athena) is COMPLETED!"
)