"""Operator that runs Athena CTAS (Create Table As Select) via awswrangler, with partition-aware insert mode."""
from typing import Dict, Optional
from config import Config
from util import Util
from airflow.models import BaseOperator
from airflow.utils.decorators import apply_defaults
from airflow.exceptions import AirflowException
[docs]
class AthenaCTASOperator(BaseOperator):
"""
An operator to execute CTAS (Create Table As Select) operations in Athena.
This operator wraps awswrangler's create_ctas_table functionality to create new tables
from SELECT queries in Athena. It handles cleaning up existing tables and managing
partitioning/bucketing configurations.
:param prefix: The S3 prefix where the table data will be stored
:type prefix: str
:param query: The SELECT query to use for creating the new table
:type query: str
:param target_table: Name of the target table to create
:type target_table: str
:param partition_by: Columns to partition the table by. Defaults to [].
:type partition_by: list[str]
:param bucket: S3 bucket name. Defaults to Config.DATA_LAKE_BUCKET.
:type bucket: str
:param source_db: Source database name. Defaults to Config.DATA_LAKE_DB.
:type source_db: str
:param target_db: Target database name. Defaults to Config.DATA_LAKE_DB.
:type target_db: str
:param workgroup: Athena workgroup name. Defaults to Config.DEFAULT_WORKGROUP.
:type workgroup: str
:param clean: Whether to remove existing table before creation. Defaults to True.
:type clean: bool
:param data_source: Data source name. Defaults to "AwsDataCatalog".
:type data_source: str
:param bucketing_info: Tuple of (columns, bucket_count) for table bucketing.
:type bucketing_info: tuple[list[str], int]
:param partition_key: Partition column name for partition-aware mode.
When set (along with partition_val), the operator will drop the partition
and INSERT INTO if the table exists, or CTAS on first run.
:type partition_key: str
:param partition_val: Partition value for partition-aware mode (templatable).
Used with partition_key for drop_athena_partition and row-count logging.
:type partition_val: str
:param column_comments: Optional dict mapping column names to their descriptions.
If provided, updates Glue table metadata after CTAS completes.
:type column_comments: Dict[str, str]
:param table_description: Optional description for the table.
If provided, updates Glue table metadata after CTAS completes.
:type table_description: str
"""
ui_color = "#89CFF0"
template_fields = ("prefix", "query", "target_table", "table_description", "partition_val")
@apply_defaults
def __init__(
self,
prefix: str,
query: str,
target_table: str,
partition_by: list[str] = [],
bucket: str = Config.DATA_LAKE_BUCKET,
source_db: str = Config.DATA_LAKE_DB,
target_db: str = Config.DATA_LAKE_DB,
workgroup: str = Config.DEFAULT_WORKGROUP,
clean: bool = True,
data_source: str = "AwsDataCatalog",
bucketing_info: tuple[list[str], int] = None,
partition_key: Optional[str] = None,
partition_val: Optional[str] = None,
column_comments: Optional[Dict[str, str]] = None,
table_description: Optional[str] = None,
*args,
**kwargs,
):
super().__init__(*args, **kwargs)
self.prefix = prefix
self.query = query
self.target_table = target_table
self.partition_by = partition_by
self.bucket = bucket
self.source_db = source_db
self.target_db = target_db
self.workgroup = workgroup
self.clean = clean
self.data_source = data_source
self.bucketing_info = bucketing_info
self.partition_key = partition_key
self.partition_val = partition_val
self.column_comments = column_comments
self.table_description = table_description
[docs]
def execute(self, context: dict) -> None:
"""
Execute the CTAS operation.
When partition_key and partition_val are set, operates in partition-aware
mode: drops only the target partition (if table exists) and INSERT INTO,
or CTAS on first run. This allows multiple states to coexist.
Args:
context: Airflow context dict
Raises:
AirflowException: If there's an error during execution
"""
try:
import awswrangler as wr
if self.partition_key and self.partition_val:
self._execute_partition_aware(wr)
else:
self._execute_full_replace(wr)
# Update Glue metadata if column_comments or table_description provided
if self.column_comments or self.table_description:
self.log.info(
f"Updating Glue metadata for {self.target_db}.{self.target_table}"
)
Util.update_glue_table_metadata(
table_name=self.target_table,
comments=self.column_comments or {},
table_description=self.table_description,
database_name=self.target_db,
)
self.log.info("Successfully updated Glue table metadata")
except Exception as e:
raise AirflowException(
f"Failed to execute CTAS operation for table {self.target_table}: {str(e)}"
)
def _execute_full_replace(self, wr) -> None:
"""Original behavior: drop entire table + CTAS."""
if self.clean:
self.log.info(f"Cleaning existing table {self.target_table}")
Util.remove_datalake_table(
tgt_table=self.target_table,
tgt_db=self.target_db,
)
self.log.info(f"Creating table {self.target_table} using CTAS")
self.log.info(f"Query: {self.query}")
wr.athena.create_ctas_table(
sql=self.query,
database=self.source_db,
ctas_database=self.target_db,
data_source=self.data_source,
ctas_table=self.target_table,
storage_format="PARQUET",
write_compression="snappy",
s3_output=f"s3://{self.bucket}/{self.prefix}",
wait=True,
workgroup=self.workgroup,
partitioning_info=self.partition_by,
bucketing_info=self.bucketing_info,
)
self.log.info(f"Successfully created table {self.target_table}")
def _execute_partition_aware(self, wr) -> None:
"""Partition-aware: drop partition + INSERT INTO, or CTAS on first run."""
# Jinja may render numeric-looking values (e.g. FIPS "37") as int;
# Glue APIs require strings.
partition_val = str(self.partition_val)
table_exists = wr.catalog.does_table_exist(
database=self.target_db, table=self.target_table
)
# If table exists, verify it has the expected partition key.
# Legacy unpartitioned tables must be dropped and recreated.
if table_exists:
import boto3
glue = boto3.client("glue")
table_def = glue.get_table(
DatabaseName=self.target_db, Name=self.target_table
)
partition_keys = {
p["Name"]
for p in table_def["Table"].get("PartitionKeys", [])
}
if self.partition_key not in partition_keys:
self.log.info(
f"Table {self.target_table} exists but lacks partition key "
f"'{self.partition_key}'; dropping and recreating."
)
Util.remove_datalake_table(
tgt_table=self.target_table, tgt_db=self.target_db
)
table_exists = False
if table_exists:
self.log.info(
f"Dropping partition {self.partition_key}={partition_val} "
f"from {self.target_table}"
)
Util.drop_athena_partition(
partition_key=self.partition_key,
partition_val=partition_val,
parent_tbname=self.target_table,
)
self.log.info(f"Inserting into {self.target_table}")
wr.athena.start_query_execution(
sql=f"INSERT INTO {self.target_table} {self.query}",
database=self.source_db,
wait=True,
workgroup=self.workgroup,
)
else:
self.log.info(
f"Creating table {self.target_table} using CTAS (first run)"
)
# Clean stale S3 data that may remain from a previous unpartitioned table
s3_path = f"s3://{self.bucket}/{self.prefix}"
self.log.info(f"Cleaning S3 path before CTAS: {s3_path}")
wr.s3.delete_objects(s3_path)
wr.athena.create_ctas_table(
sql=self.query,
database=self.source_db,
ctas_database=self.target_db,
data_source=self.data_source,
ctas_table=self.target_table,
storage_format="PARQUET",
write_compression="snappy",
s3_output=f"s3://{self.bucket}/{self.prefix}",
wait=True,
workgroup=self.workgroup,
partitioning_info=[self.partition_key],
bucketing_info=self.bucketing_info,
)
rows = wr.athena.read_sql_query(
f"SELECT COUNT(*) AS cnt FROM {self.target_table} "
f"WHERE {self.partition_key} = '{partition_val}'",
database=self.target_db,
)
self.log.info(
f"{self.target_table}: {int(rows['cnt'].iloc[0]):,} rows "
f"({self.partition_key}={partition_val})"
)
[docs]
class PostgresCTASOperator(BaseOperator):
"""
An operator to execute CTAS (Create Table As Select) operations in Postgres.
This operator wraps the Postgres CREATE TABLE AS SELECT statement to create new tables
from SELECT queries in Postgres.
ATTENTION: This operator is DESTRUCTIVE and will remove the existing table before creating
a new one if clean=True. <3
:param query: The SELECT query to use for creating the new table
:type query: str
:param target_table: Name of the target table to create
:type target_table: str
:param target_db: Name of the target database. Defaults to Config.DEFAULT_DB_POSTGRES.
:type target_db: str
:param target_schema: Target schema name. Defaults to "data".
:type target_schema: str
:param clean: Whether to remove existing table before creation. Defaults to True.
:type clean: bool
"""
ui_color = "#BEEFED"
template_fields = ("query",)
@apply_defaults
def __init__(
self,
query: str,
target_table: str,
target_db: str,
target_schema: str = None,
clean: bool = True,
*args,
**kwargs,
):
super().__init__(*args, **kwargs)
self.query = query
self.target_table = target_table
self.target_db = target_db
self.target_schema = target_schema
self.clean = clean
[docs]
def execute(self, context: dict) -> None:
"""
Execute the CTAS operation.
Args:
context: Airflow context dict
Raises:
AirflowException: If there's an error during execution
"""
try:
from airflow.providers.postgres.hooks.postgres import PostgresHook
hook = PostgresHook(postgres_conn_id=self.target_db)
# Use the resolved value from params if target_schema isn't set on the instance.
# This is done to avoid timing issues with dynamic mapping.
# If target_schema is not set, it will raise an error.
# This is to ensure that the operator can be used in dynamic mapping.
# If target_schema is set, it will be used directly.
# If target_schema is not set, it will be resolved from params.
resolved_schema = self.target_schema or self.params.get("target_schema")
if not resolved_schema:
raise ValueError(
"target_schema must be provided via dynamic mapping or as a default."
)
# Drop existing table if requested
if self.clean:
self.log.info(
f"Dropping existing table {self.target_table} in {resolved_schema} (if it exists)"
)
hook.run(f"DROP TABLE IF EXISTS {resolved_schema}.{self.target_table}")
# Execute CTAS operation
self.log.info(
f"Creating table {self.target_table} in {resolved_schema} using CTAS"
)
self.log.info(f"Query: {self.query}")
hook.run(
f"CREATE TABLE {resolved_schema}.{self.target_table} AS {self.query}"
)
self.log.info(f"Successfully created table {self.target_table}")
except Exception as e:
raise AirflowException(
f"Failed to execute CTAS operation for table {self.target_table}: {str(e)}"
)