"""Operator that transfers a datalake (Athena) table, partition-aware, into a Postgres table with staged hot swap."""
import uuid
from airflow.utils.context import Context
from airflow.providers.postgres.hooks.postgres import PostgresHook
from config import Config
from util import Util
from typing import Any, Literal, List, Dict
from pathlib import Path
from jinja2 import Template
from airflow.models import BaseOperator
from airflow.utils.decorators import apply_defaults
[docs]
class TableToPGOperator(BaseOperator):
ui_color = "#FFA500"
ui_fgcolor = "#000000"
custom_operator_name = "TableToPGOperator"
template_fields_renderers = {
"create_table_sql": "sql",
"copy_table_sql": "sql",
"post_transfer_sql": "sql",
}
template_ext = ".sql"
@apply_defaults
def __init__(
self,
dest_db: str,
dest_table: str,
copy_table_sql: str,
upsert_origin_table: str = None,
on_conflict_sql: str = None,
skip_if_empty: bool = False,
partitions: List[str] = [],
copy_partitions: List[str] = [],
schema: str = "public",
create_table_sql: str = None,
post_transfer_sql: str = None,
create_partition_template: str = "CREATE TABLE {{schema}}.{{table_name}}_{{parname}}_stage PARTITION OF {{schema}}.{{stage_table_name}} FOR VALUES in ('{{parname}}')",
task_suffix: Literal["", "large", "xlarge"] = "",
drop_backup_before_start: bool = True,
drop_backup_at_end: bool = True,
template: str = "",
inlets=[],
map_name: str = "",
table_description: str | None = None,
column_descriptions: Dict[str, str] = {},
data_product_name: str | None = None,
*args,
**kwargs,
):
"""Transfers a datalake table with its partitions to pg table with hot swap
:param dest_db: _description_
:type dest_db: str
:param dest_table: _description_
:type dest_table: str
:param partitions: _description_
:type partitions: list[str]
:param create_table_sql: _description_, defaults to None
:type create_table_sql: str, optional
:param copy_table_sql: _description_
:type copy_table_sql: str, optional
:param post_transfer_sql: _description_, defaults to None; to run on PG db after data transfer complete e.g. CREATE INDEX, it has access to {{uuid}}, {{table_name}}, {{stage_table_name}}, {{schema}} variables predefined
:type post_transfer_sql: str, optional
:param task_suffix: _description_, defaults to ""
:type task_suffix: Literal["", "large", "xlarge"], optional
:param drop_backup_before_start: _description_, defaults to True
:type drop_backup_before_start: bool, optional
:param drop_backup_at_end: _description_, defaults to True
:type drop_backup_at_end: bool, optional
:param upsert_origin_table: _description_, defaults to None; if provided, it will be used to upsert data from the origin table to the destination table. It should be a table name with schema, e.g. "public.my_table" it must exist in the destination database and also must have the proper column indexes and such
:type upsert_origin_table: str, optional
:param on_conflict_sql: _description_, defaults to None; if provided, it will be used in the upsert operation to handle conflicts, e.g. "ON CONFLICT (id) DO UPDATE SET column1 = EXCLUDED.column1, column2 = EXCLUDED.column2"
:type on_conflict_sql: str, optional
:param schema: _description_, defaults to "public"
:type schema: str, optional
:param skip_if_empty: _description_, defaults to False; if True, an empty
copy_table_sql result raises AirflowSkipException (task skipped) instead
of failing the task
:type skip_if_empty: bool, optional
"""
super().__init__(*args, **kwargs)
self.dest_db = dest_db
self.dest_table = dest_table
self.partitions = partitions
self.create_table_sql = create_table_sql
self.copy_table_sql = copy_table_sql
self.post_transfer_sql = post_transfer_sql
self.task_suffix = task_suffix
self.stage_table_name = f"{dest_table}_stage"
self.drop_backup_before_start = drop_backup_before_start
self.drop_backup_at_end = drop_backup_at_end
self.schema = schema
self.query_params = {
"uuid": str(uuid.uuid4().hex[:8]),
"table_name": dest_table,
"stage_table_name": self.stage_table_name,
"schema": self.schema,
}
self.template = template
self.dag_root = Config.DAG_ROOT
self.copy_partitions = copy_partitions
self.create_partition_template = create_partition_template
self.inlets = inlets
if map_name == "":
self.map_name = f"{self.dest_db}.{self.schema}.{self.dest_table}"
self.table_description = table_description
self.column_descriptions = column_descriptions
self.data_product_name = data_product_name
self.upsert_origin_table = upsert_origin_table
self.on_conflict_sql = on_conflict_sql
self.skip_if_empty = skip_if_empty
[docs]
def target_stage_table(self) -> str:
"""returns stage table name with schema
:return: stage table name with schema
:rtype: str
"""
return f"{self.schema}.{self.stage_table_name}"
[docs]
def copy_to_pg(self, partition: str = ""):
import awswrangler as wr
temp_s3_path = f"temp/temp_{self.schema}_{self.dest_db}_{self.dest_table}_temp_{partition}/"
query_params = {**{"partition": partition}, **self.query_params}
sql_to_run = self.read_sql_template(self.copy_table_sql, query_params)
stage_table_name = self.target_stage_table()
if not sql_to_run or not sql_to_run.strip():
raise ValueError(f"copy_table_sql rendered to an empty string for partition '{partition}'. Check your SQL template.")
df_results=wr.athena.read_sql_query(sql=f"""with data as ({sql_to_run}) select * from data limit 10""", database=Config.DATA_LAKE_DB)
if df_results.empty:
msg = f"No data returned from query for partition '{partition}'. Skipping copy to PG."
if self.skip_if_empty:
from airflow.exceptions import AirflowSkipException
raise AirflowSkipException(msg)
raise Exception(msg)
Util.transfer_datalake_db(
temp_key=temp_s3_path,
table_name=stage_table_name,
tgt_query=sql_to_run,
drop_table=False,
task_suffix=self.task_suffix,
dest_db=self.dest_db,
use_pg_copy=True,
upsert_origin_table=self.upsert_origin_table,
on_conflict_sql=self.on_conflict_sql,
)
[docs]
def read_sql_template(self, sql: str, params: Dict[str, str] = {}) -> str:
if sql.lower().endswith(".sql"):
path_to_sql = f"{self.dag_root}/{self.template}/{sql}"
self.log.info(f"opening sql file {path_to_sql}")
sql_file_path = Path(path_to_sql)
sql = sql_file_path.read_text()
sql_jinja_template = Template(sql)
sql_to_run = sql_jinja_template.render(params)
self.log.info(f"{sql_to_run=}")
return sql_to_run
[docs]
def drop_stage_table(self):
db_hook = PostgresHook(postgres_conn_id=self.dest_db)
sql = f"""
begin;
SET search_path TO {self.schema};
DROP TABLE IF EXISTS {self.stage_table_name} CASCADE;
commit;
"""
db_hook.run(sql, autocommit=True)
[docs]
def create_table_with_partitions(self):
params = self.query_params
if self.create_table_sql:
db_hook = PostgresHook(postgres_conn_id=self.dest_db)
sql_text = self.read_sql_template(self.create_table_sql, params)
db_hook.run(sql_text, autocommit=True)
# TODO: @majid the whole if block below,
# and params self.partitions (easily confused with self.copy_partitions) and self.create_partition_template
# are NOT developed, or used anywhere
if (
len(self.partitions)
# and "PARTITION BY".lower() in self.create_table_sql.lower()
):
# create table partisions
partitions_query = ";".join(
[
self.read_sql_template(
self.create_partition_template,
{**params, **{"parname": parname}},
)
for parname in self.partitions
]
)
self.log.info(f"partitions_query:{partitions_query}")
db_hook.run(partitions_query, autocommit=True)
[docs]
def transfer_data_per_partition(self):
if len(self.copy_partitions):
for partition in self.copy_partitions:
self.copy_to_pg(partition=partition)
else:
self.copy_to_pg()
[docs]
def resolve_self_x_com_attributes(self, variable: str, context: Context):
value = getattr(self, variable)
self.log.info("%s is %s", variable, value)
x_com_value = self.resolve_x_com_arg(value, context)
self.log.info("%s after task is %s", variable, x_com_value)
setattr(self, variable, x_com_value)
[docs]
def resolve_x_com_arg(self, value: str, context: Context):
if value and "PlainXComArg" in str(type(value)):
return value.resolve(context)
# otherwise do nothing
return value
[docs]
def sanitize_params(self, context: Context):
"""read params and converts them to xcom results if they are passed from a task
:param context: operator context
:type context: Context
:return: returns params object
:rtype: _type_
"""
params = context["params"]
self.log.info("input params are %s", params)
for key in params:
value = params[key]
self.log.info("%s is %s", key, value)
x_com_value = self.resolve_x_com_arg(value, context)
self.log.info("%s after task is %s", key, x_com_value)
params[key] = x_com_value
return params
[docs]
def get_output_table_urn(self):
from datahub_airflow_plugin.entities import Urn
env = "DEV"
if Config.PRODUCTION:
env = "PROD"
if "STAGE".lower() in self.dest_db.lower():
env = "STG"
# We only support bossplatform pg
# urn:li:dataset:(urn:li:dataPlatform:postgres,bossplatform.data.c_cai,STG)
# urn:li:dataset:(urn:li:dataPlatform:postgres,bossplatform.data.c_cai,STG)
# urn:li:dataset:(urn:li:dataPlatform:postgres,bossplatform.data.c_bbm_audit,STG)
return Urn(
f"urn:li:dataset:(urn:li:dataPlatform:postgres,bossplatform.{self.schema}.{self.dest_table},{env})"
)
[docs]
def get_lineage(self):
from datahub_airflow_plugin.entities import Urn
inlets = []
outlets = []
if self.inlets:
for urn in inlets:
if isinstance(urn, str):
inlets.append(Urn(urn))
else:
inlets.append(urn)
# We only support bossplatform pg
dataset_urn = self.get_output_table_urn()
outlets.append(dataset_urn)
return inlets, outlets
[docs]
def execute(self, context: Context) -> Any:
# add map_index_template
context["params"]["map_name"] = self.map_name
context["map_index_template"] = "{{ params.map_name }}"
for variable in ["partitions", "copy_partitions", "dest_db"]:
self.resolve_self_x_com_attributes(variable, context=context)
params = self.sanitize_params(context=context)
self.query_params = {**params, **self.query_params}
self.log.info("query_params: %s", self.query_params)
db_hook = PostgresHook(postgres_conn_id=self.dest_db)
self.drop_stage_table()
self.create_table_with_partitions()
self.transfer_data_per_partition()
if self.post_transfer_sql:
sql_text = self.read_sql_template(self.post_transfer_sql, self.query_params)
db_hook.run(sql_text, autocommit=True)
Util.pg_swap_with_partitions(
db_hook,
stage_table=self.stage_table_name,
table=self.dest_table,
postfix="_old",
drop_backup_before_start=True,
drop_backup_at_end=True,
schema=self.schema,
)
self.drop_stage_table()
try:
if self.schema == "demo_base":
self.add_datahub_metadata(context)
else:
self.log.info(f"Skipping DataHub metadata for schema '{self.schema}' (only 'demo_base' is tracked)")
self.add_table_metadata()
except Exception as e:
self.log.error(f"Error adding metadata: {e}")
self.log.info("TableToPGOperator completed")