from airflow.utils.decorators import apply_defaults
from airflow.exceptions import AirflowException
from typing import Dict, Optional, Union, List
from operators.iceberg.base_iceberg_table import BaseIcebergTableOperator
from operators.iceberg.remove_iceberg_table import RemoveIcebergTableOperator
import awswrangler as wr
[docs]
class MigrateHiveToIcebergOperator(BaseIcebergTableOperator):
"""
Migrates a Hive table to Iceberg format with data validation and enhanced controls.
This operator creates a new Iceberg table from a source Hive table with advanced options
for schema customization, table properties, and validation controls.
:param hive_table: Source Hive table name
:type hive_table: str
:param iceberg_table: Target Iceberg table name
:type iceberg_table: str
:param location: S3 location for the new Iceberg table
:type location: str
:param partitioning: Optional partitioning configuration, can be:
- String column name (e.g., "state")
- List of column names (e.g., ["state", "year"])
- ARRAY expression (e.g., "ARRAY['bucket(id, 64)', 'state']")
:type partitioning: Optional[Union[str, List[str]]]
:param database: The database containing the Hive and target Iceberg tables
:type database: str
:param table_properties: Additional Iceberg table properties
:type table_properties: Dict[str, str]
:param column_mapping: Column mapping from source to target {source_col: target_col}
:type column_mapping: Dict[str, str]
:param include_columns: Columns to include (if None, includes all)
:type include_columns: Optional[List[str]]
:param exclude_columns: Columns to exclude
:type exclude_columns: Optional[List[str]]
:param where_clause: WHERE clause to filter source data
:type where_clause: Optional[str]
:param dry_run: If True, only validate and print migration steps without execution
:type dry_run: bool
:param optimize_if_successful: Whether to optimize the table after successful creation
:type optimize_if_successful: bool
:param optimize_options: Options for OPTIMIZE command if performed
:type optimize_options: str
:param cleanup_on_failure: Whether to remove the Iceberg table if validation fails
:type cleanup_on_failure: bool
"""
template_fields = BaseIcebergTableOperator.template_fields + [
"hive_table",
"location",
"partitioning",
"optimize_options",
"table_properties",
"where_clause",
]
ui_color = "#D4EFDF" # Light Green shade
DEFAULT_TABLE_PROPERTIES = {
"format": "parquet",
"write_compression": "snappy",
"optimize_rewrite_delete_file_threshold": "10",
"is_external": "false",
}
@apply_defaults
def __init__(
self,
hive_table: str,
iceberg_table: str,
location: str,
partitioning: Optional[Union[str, List[str]]] = None,
database: Optional[str] = None,
table_properties: Optional[Dict[str, str]] = None,
column_mapping: Optional[Dict[str, str]] = None,
include_columns: Optional[List[str]] = None,
exclude_columns: Optional[List[str]] = None,
where_clause: Optional[str] = None,
dry_run: bool = False,
optimize_if_successful: bool = True,
optimize_options: str = "REWRITE DATA USING BIN_PACK",
cleanup_on_failure: bool = True,
overwrite_if_exists: bool = False,
*args,
**kwargs,
):
# Important: Always set verify_exists=False for this operator
# The table might not exist yet, or we may want to overwrite it
super().__init__(
table_name=iceberg_table,
database=database,
verify_exists=False, # Always false for migration operator
*args,
**kwargs,
)
self.hive_table = hive_table
self.location = location
self.partitioning = partitioning
self.table_properties = table_properties or {}
self.column_mapping = column_mapping or {}
self.include_columns = include_columns
self.exclude_columns = exclude_columns or []
self.where_clause = where_clause
self.dry_run = dry_run
self.optimize_if_successful = optimize_if_successful
self.optimize_options = optimize_options
self.cleanup_on_failure = cleanup_on_failure
self.overwrite_if_exists = overwrite_if_exists
[docs]
def execute(self, context: Dict) -> None:
import awswrangler as wr
dry_run_prefix = "[DRY RUN] " if self.dry_run else ""
self.log.info(
"{0}Starting migration from Hive table '{1}' to Iceberg table '{2}' in database '{3}'".format(
dry_run_prefix, self.hive_table, self.table_name, self.database
)
)
try:
# 1. Check if source Hive table exists
if not self.dry_run and not wr.catalog.does_table_exist(
database=self.database,
table=self.hive_table,
):
raise AirflowException(
"Source Hive table {0}.{1} does not exist".format(
self.database, self.hive_table
)
)
# 2. Check if target Iceberg table already exists
if not self.dry_run:
if wr.catalog.does_table_exist(
database=self.database,
table=self.table_name,
):
if not self.overwrite_if_exists:
# If overwrite not enabled, fail with error
raise AirflowException(
"Target Iceberg table {0}.{1} already exists. Set overwrite_if_exists=True to replace it.".format(
self.database, self.table_name
)
)
# If overwrite is enabled, drop and recreate
self.log.info(
"Overwrite enabled - dropping existing table before recreating..."
)
cleanup = RemoveIcebergTableOperator(
task_id="remove_existing_iceberg_table",
table_name=self.table_name,
database=self.database,
skip_if_not_exists=True,
)
cleanup.execute(context)
self.log.info("Existing table dropped, proceeding with migration.")
# 3. Get schema from source table if needed for column mapping/filtering
source_columns = None
if self.column_mapping or self.include_columns or self.exclude_columns:
source_columns = self._get_source_table_columns()
# 4. Build column projection for SELECT statement
column_projection = self._build_column_projection(source_columns)
# 5. Build table properties clause including partitioning
table_properties_clause = self._build_table_properties_clause()
# 6. Build the full SQL for CREATE TABLE
create_table_sql = self._build_create_table_sql(
column_projection, table_properties_clause
)
# If dry run, just log the SQL and return
if self.dry_run:
self.log.info("[DRY RUN] Would execute: \n{0}".format(create_table_sql))
return
# 7. Create Iceberg table with data from Hive
self.log.info(
"Creating Iceberg table with query:\n{0}".format(create_table_sql)
)
self._execute_athena_query(sql=create_table_sql)
# 8. Validate data count matches
self.log.info("Validating row counts between source and target tables")
source_count, target_count = self._validate_row_counts()
# 9. Check validation results
if source_count == target_count:
self.log.info(
"✅ Validation successful: Counts match ({0} rows)".format(
source_count
)
)
# Optimize the newly created table if requested
if self.optimize_if_successful:
self.log.info("Optimizing new Iceberg table")
optimize_sql = "OPTIMIZE {0}.{1} {2}".format(
self.database, self.table_name, self.optimize_options
)
self._execute_athena_query(sql=optimize_sql)
self.log.info("✅ Table optimization completed")
else:
self.log.error("❌ Validation failed: Count mismatch detected!")
self.log.error(
" Difference: {0} rows".format(abs(source_count - target_count))
)
# Clean up the failed table if requested
if self.cleanup_on_failure:
self._cleanup_failed_table(context)
raise AirflowException(
"Migration validation failed: Source count ({0}) does not match "
"target count ({1})".format(source_count, target_count)
)
except Exception as e:
self.log.error("Migration failed: {0}".format(str(e)))
# Attempt cleanup on any exception if requested
if not self.dry_run and self.cleanup_on_failure:
try:
self._cleanup_failed_table(context)
except Exception as cleanup_error:
self.log.error(
"Cleanup also failed: {0}".format(str(cleanup_error))
)
raise AirflowException(
"Failed to migrate Hive table to Iceberg: {0}".format(str(e))
) from e
def _get_source_table_columns(self) -> List[str]:
"""
Get the column names from the source Hive table.
"""
self.log.info(
"Retrieving schema from source Hive table {0}.{1}".format(
self.database, self.hive_table
)
)
if self.dry_run:
self.log.info("[DRY RUN] Would retrieve source table columns")
return ["column1", "column2", "column3"] # Placeholder for dry run
try:
# Get all column names and types from the table
all_cols = wr.catalog.get_table_types(
database=self.database, table=self.hive_table
)
# Get table DDL to identify partition columns
ddl_str = wr.athena.show_create_table(
database=self.database, table=self.hive_table
)
# Extract partition columns using DDLParser
from simple_ddl_parser import DDLParser
parsed_ddl = DDLParser(ddl_str).run(output_mode="athena")
# Extract partition columns if they exist, otherwise use empty list
partition_cols = []
if parsed_ddl and parsed_ddl[0].get("partitioned_by"):
partition_cols = [
item["name"].replace("`", "")
for item in parsed_ddl[0]["partitioned_by"]
]
# Return all column names (both regular and partition columns)
columns = list(all_cols.keys())
self.log.info("Found {0} columns in source table".format(len(columns)))
self.log.info("Partition columns: {0}".format(partition_cols))
return columns
except Exception as e:
self.log.error(
"Failed to retrieve source table columns: {0}".format(str(e))
)
raise AirflowException(
"Could not get column information from source table: {0}".format(str(e))
) from e
def _build_column_projection(
self, source_columns: Optional[List[str]] = None
) -> str:
"""
Build the column projection (SELECT clause) based on column mappings and filters.
"""
# If no customization is needed, use "*" for all columns
if (
not self.column_mapping
and not self.include_columns
and not self.exclude_columns
):
return "*"
# Get columns to include
if self.include_columns:
columns_to_include = self.include_columns
elif source_columns:
# Start with all columns, then remove excluded ones
columns_to_include = [
col for col in source_columns if col not in self.exclude_columns
]
else:
self.log.warning(
"Column filtering requested but couldn't get source columns. Using '*' instead."
)
return "*"
# Apply column mapping
projection_parts = []
for col in columns_to_include:
if isinstance(col, str) and col in self.column_mapping:
projection_parts.append(
"{0} AS {1}".format(self.column_mapping[col], col)
)
else:
projection_parts.append(col)
return ", ".join(projection_parts)
def _build_partition_clause(self) -> str:
"""
Build the partitioning expression (without the property name).
"""
if not self.partitioning:
return ""
# Handle already formatted ARRAY expression
if isinstance(self.partitioning, str) and self.partitioning.upper().startswith(
"ARRAY["
):
return self.partitioning
# Format the array expression based on input type
if isinstance(self.partitioning, (list, tuple)):
# For multiple columns, join them with commas - ensure single quotes around column names
columns_formatted = ", ".join(
["'{0}'".format(col) for col in self.partitioning]
)
return "ARRAY[{0}]".format(columns_formatted)
else:
# For a single column - ensure single quotes around column name
return "ARRAY['{0}']".format(self.partitioning)
def _build_table_properties_clause(self) -> str:
"""
Build the complete table properties clause including partitioning.
"""
properties = {
"table_type": "ICEBERG",
**self.DEFAULT_TABLE_PROPERTIES,
**self.table_properties,
}
# Add partitioning if specified
partition_clause = self._build_partition_clause()
if partition_clause:
# Store the raw array expression without additional quotes
properties["partitioning"] = partition_clause
# Build the properties string, handling different property types
properties_parts = []
for key, value in properties.items():
if key == "partitioning":
# Special handling for partitioning to avoid double quoting
properties_parts.append("{0} = {1}".format(key, value))
elif key == "is_external":
# Boolean property needs lowercase 'false' or 'true' without quotes
properties_parts.append("{0} = {1}".format(key, value.lower()))
elif key in [
"optimize_rewrite_delete_file_threshold",
"write_target_data_file_size_bytes",
] or (isinstance(value, str) and value.isdigit()):
# Numeric properties - no quotes
properties_parts.append("{0} = {1}".format(key, value))
elif "-" in key:
# Quote property names with hyphens
properties_parts.append("\"{0}\" = '{1}'".format(key, value))
else:
# Standard property format
properties_parts.append("{0} = '{1}'".format(key, value))
return ", ".join(properties_parts)
def _build_create_table_sql(
self, column_projection: str, table_properties_clause: str
) -> str:
"""
Build the complete CREATE TABLE SQL statement.
"""
# Add WHERE clause if specified
where_clause = (
"\nWHERE {0}".format(self.where_clause) if self.where_clause else ""
)
sql_template = """
CREATE TABLE {database}.{table_name}
WITH (
{table_properties},
location = '{location}'
)
AS SELECT {column_projection}
FROM {database}.{hive_table}{where_clause}
"""
return sql_template.format(
database=self.database,
table_name=self.table_name,
table_properties=table_properties_clause,
location=self.location,
column_projection=column_projection,
hive_table=self.hive_table,
where_clause=where_clause,
)
def _validate_row_counts(self) -> tuple:
"""
Validate that row counts match between source and target tables.
"""
# Add WHERE clause to source count if specified
where_condition = (
" WHERE {0}".format(self.where_clause) if self.where_clause else ""
)
sql_template = """
SELECT
(SELECT COUNT(*) FROM {database}.{hive_table}{where_condition}) AS hive_count,
(SELECT COUNT(*) FROM {database}.{table_name}) AS iceberg_count
"""
sql = sql_template.format(
database=self.database,
hive_table=self.hive_table,
table_name=self.table_name,
where_condition=where_condition,
)
validate = wr.athena.read_sql_query(sql, database=self.database)
# Get count values
source_count = validate["hive_count"].iloc[0]
target_count = validate["iceberg_count"].iloc[0]
# Log validation results
self.log.info("Migration validation results:")
self.log.info(
" Source Hive table '{0}' count: {1}".format(self.hive_table, source_count)
)
self.log.info(
" Target Iceberg table '{0}' count: {1}".format(
self.table_name, target_count
)
)
return source_count, target_count
def _cleanup_failed_table(self, context: Dict) -> None:
"""
Clean up a failed table migration.
"""
self.log.info("Cleaning up failed migration by removing target table")
remover = RemoveIcebergTableOperator(
task_id="remove_failed_iceberg_table",
table_name=self.table_name,
database=self.database,
skip_if_not_exists=True,
)
remover.execute(context)