import sys import abc from datetime import datetime, timedelta from pyspark.context import SparkContext # type: ignore from awsglue.utils import getResolvedOptions # type: ignore from awsglue.context import GlueContext # type: ignore from awsglue.job import Job # type: ignore from utils.snowflake.snowflake_client import get_snowflake_connection # type: ignore from utils.util import get_data_from_postgresql_and_save_to_s3 # type: ignore import boto3 # type: ignore class RDStoS3Job(abc.ABC): """ Abstract base class representing an AWS Glue Job that gets data from Postgres, copies it into an S3 bucket, and then upserts it into Snowflake. """ def __init__( self, table_name: str, procedure_name: str, task_monitor_name: str, prod_warehouse: str = "SUNO_PROD_GLUE_HOURLY_X_SMALL", staging_warehouse: str = "SUNO_STAGING_GLUE_HOURLY_X_SMALL", print_debug_logs: bool = False, ): self.table_name = table_name self.procedure_name = procedure_name self.task_monitor_name = task_monitor_name self.print_debug_logs = print_debug_logs self.prod_warehouse_name = prod_warehouse self.staging_warehouse_name = staging_warehouse self.environment, self.s3_bucket, self.warehouse, self.db_name = ( self._get_running_info() ) self.start_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S") def debug_print(self, message: str): if self.print_debug_logs: print(message) def _get_running_info(self): """ Gets important information about the job's running environment, depending on whether the job is run in staging or production. Returns the environment (staging or prod) and the S3 bucket, Snowflake warehouse, and Snowflake database to write to. """ PROD_ACCOUNT_ID = "734185074900" STAGING_ACCOUNT_ID = "590183763515" # uses the account id to determine whether we're in prod or staging sts = boto3.client("sts") account_id = sts.get_caller_identity()["Account"] if account_id == PROD_ACCOUNT_ID: environment = "PROD" s3_bucket = "analytics-database-data" warehouse = self.prod_warehouse_name elif account_id == STAGING_ACCOUNT_ID: environment = "STAGING" s3_bucket = "analytics-database-data-staging" warehouse = self.staging_warehouse_name else: raise Exception("Invalid account ID.") # sets the Snowflake database to write to depending on whether we're in prod or staging db_name = f"SUNO_{environment}" return environment, s3_bucket, warehouse, db_name def _initialize_glue_context(self): """ Initializes some context necessary for the Glue Job to be run. """ args = getResolvedOptions(sys.argv, ["JOB_NAME"]) sc = SparkContext() glueContext = GlueContext(sc) spark = glueContext.spark_session job = Job(glueContext) job.init(args["JOB_NAME"], args) return glueContext, spark, job def get_call_snowflake_procedure_sql( self, procedure_name: str, params: list, ): """ Returns a SQL query that calls a Snowflake procedure with the given parameters. """ return f"""call {procedure_name}({", ".join([f"'{param}'" for param in params])})""" def get_table_name(self) -> str: return self.table_name @abc.abstractmethod def get_postgres_sample_query(self) -> str: """ Returns the sample query that should be used to get data from Postgres. """ pass @abc.abstractmethod def get_postgres_transform_sql(self) -> str: """ Returns the transform SQL query applied to the data from Postgres. """ pass @abc.abstractmethod def get_postgres_partition_keys(self) -> list[str]: """ Returns the partition keys that should be used to partition the data in S3. For example, if the data is processed daily, the partition keys may be ["pdate"]. """ pass @abc.abstractmethod def get_snowflake_procedure_sql_parameters(self) -> list[str]: """ Returns the parameters that should be passed to the Snowflake procedure that performs the upsert. """ pass @abc.abstractmethod def get_job_frequency(self) -> str: """ Returns the frequency with which this job is called, e.g. "DAILY" or "HOURLY". """ pass def run_job(self): """ Runs this Glue Job. Gets data from the Postgres table with the table name associated with this job, saves it to S3, and then upserts it into Snowflake. Does not handle any errors. """ # Glue parameters glueContext, spark, job = self._initialize_glue_context() self.debug_print(f"Environment: {self.environment}") self.debug_print(f"S3 Bucket: {self.s3_bucket}") self.debug_print(f"Warehouse: {self.warehouse}") self.debug_print(f"DB Name: {self.db_name}") POSTGRES_CONNECTION_NAME = "analystic-database-connection" row_count = get_data_from_postgresql_and_save_to_s3( db_connection_options={ "useConnectionProperties": "true", "dbtable": self.table_name, "connectionName": POSTGRES_CONNECTION_NAME, "sampleQuery": self.get_postgres_sample_query(), }, s3_connection_options={ "path": f"s3://{self.s3_bucket}/{self.table_name}/", "partitionKeys": self.get_postgres_partition_keys(), }, transform_sql=self.get_postgres_transform_sql(), glueContext=glueContext, spark=spark, ) self.debug_print(f"Saved data to S3. Row count: {row_count}") task_monitor_sql_params = [ self.task_monitor_name, self.start_time, row_count, "SUCCESS", self.get_job_frequency().upper(), ] snowflake_procedure_sql = self.get_call_snowflake_procedure_sql( self.procedure_name, self.get_snowflake_procedure_sql_parameters() ) task_monitor_sql = self.get_call_snowflake_procedure_sql( "TASK_MONITOR_INSERT_PROC", task_monitor_sql_params ) connection = get_snowflake_connection( self.warehouse, self.db_name, self.environment ) cursor = connection.cursor() self.debug_print(f"Calling SQL: {snowflake_procedure_sql}") try: result = cursor.execute(snowflake_procedure_sql) self.debug_print(f"Calling SQL: {task_monitor_sql}") result = cursor.execute(task_monitor_sql) except Exception as e: print(f"Error executing SQL: {e}") raise e self.debug_print("Query executed successfully.") job.commit() def run_backfill_job(self, glueContext, spark): """ Runs this Glue Job. Gets data from the Postgres table with the table name associated with this job, saves it to S3, and then upserts it into Snowflake. Does not handle any errors. """ # Glue parameters self.debug_print(f"Environment: {self.environment}") self.debug_print(f"S3 Bucket: {self.s3_bucket}") self.debug_print(f"Warehouse: {self.warehouse}") self.debug_print(f"DB Name: {self.db_name}") POSTGRES_CONNECTION_NAME = "analystic-database-connection" row_count = get_data_from_postgresql_and_save_to_s3( db_connection_options={ "useConnectionProperties": "true", "dbtable": self.table_name, "connectionName": POSTGRES_CONNECTION_NAME, "sampleQuery": self.get_postgres_sample_query(), }, s3_connection_options={ "path": f"s3://{self.s3_bucket}/{self.table_name}/", "partitionKeys": self.get_postgres_partition_keys(), }, transform_sql=self.get_postgres_transform_sql(), glueContext=glueContext, spark=spark, ) self.debug_print(f"Saved data to S3. Row count: {row_count}") task_monitor_sql_params = [ self.task_monitor_name, self.start_time, row_count, "SUCCESS", self.get_job_frequency().upper(), ] snowflake_procedure_sql = self.get_call_snowflake_procedure_sql( self.procedure_name, self.get_snowflake_procedure_sql_parameters() ) task_monitor_sql = self.get_call_snowflake_procedure_sql( "TASK_MONITOR_INSERT_PROC", task_monitor_sql_params ) connection = get_snowflake_connection( self.warehouse, self.db_name, self.environment ) cursor = connection.cursor() self.debug_print(f"Calling SQL: {snowflake_procedure_sql}") try: result = cursor.execute(snowflake_procedure_sql) self.debug_print(f"Calling SQL: {task_monitor_sql}") result = cursor.execute(task_monitor_sql) except Exception as e: print(f"Error executing SQL: {e}") raise e self.debug_print("Query executed successfully.") class RDStoS3JobCustom(RDStoS3Job): """ Represents an RDS to S3 job. The sample query, transform SQL, partition keys, job frequency, and parameters to be fed to the Snowflake procedure are provided by the caller when this class is instantiated. """ def __init__( self, table_name: str, procedure_name: str, task_monitor_name: str, postgres_sample_query: str, postgres_transform_sql: str, postgres_partition_keys: list[str], snowflake_procedure_sql_parameters: list[str], job_frequency: str, prod_warehouse: str = "SUNO_PROD_GLUE_HOURLY_X_SMALL", staging_warehouse: str = "SUNO_STAGING_GLUE_HOURLY_X_SMALL", print_debug_logs: bool = False, ): RDStoS3Job.__init__( self, table_name, procedure_name, task_monitor_name, prod_warehouse=prod_warehouse, staging_warehouse=staging_warehouse, print_debug_logs=print_debug_logs, ) self.postgres_sample_query = postgres_sample_query self.postgres_transform_sql = postgres_transform_sql self.postgres_partition_keys = postgres_partition_keys self.snowflake_procedure_sql_parameters = snowflake_procedure_sql_parameters self.job_frequency = job_frequency def get_postgres_sample_query(self) -> str: return self.postgres_sample_query def get_postgres_transform_sql(self) -> str: return self.postgres_transform_sql def get_postgres_partition_keys(self) -> list[str]: return self.postgres_partition_keys def get_snowflake_procedure_sql_parameters(self) -> list[str]: return self.snowflake_procedure_sql_parameters def get_job_frequency(self) -> str: return self.job_frequency class HourlyRDStoS3Job(RDStoS3Job): """ Represents an RDS to S3 job that executes hourly. Uses a default sample query and transform SQL. """ def __init__( self, table_name: str, procedure_name: str, task_monitor_name: str, prod_warehouse: str = "SUNO_PROD_GLUE_HOURLY_X_SMALL", staging_warehouse: str = "SUNO_STAGING_GLUE_HOURLY_X_SMALL", print_debug_logs: bool = False, ): RDStoS3Job.__init__( self, table_name, procedure_name, task_monitor_name, prod_warehouse=prod_warehouse, staging_warehouse=staging_warehouse, print_debug_logs=print_debug_logs, ) current_date = datetime.now() one_hour_ago = current_date - timedelta(hours=1) self.start_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S") self.hour_start = one_hour_ago.replace(minute=0, second=0, microsecond=0) self.hour_end = self.hour_start + timedelta(hours=1) self.date = one_hour_ago.strftime("%Y-%m-%d") self.hour = one_hour_ago.strftime("%H") def get_postgres_sample_query(self) -> str: return f"SELECT * FROM {self.table_name} WHERE (updated_at >= '{self.hour_start}' AND updated_at < '{self.hour_end}') OR (created_at >= '{self.hour_start}' AND created_at < '{self.hour_end}')" def get_postgres_transform_sql(self) -> str: return f""" select *, '{self.date}' as pdate, '{self.hour}' as phour from result """ def get_postgres_partition_keys(self) -> list[str]: return ["pdate", "phour"] def get_snowflake_procedure_sql_parameters(self) -> list[str]: stage_path = f"@SUNO_DATABASE_EVENTS/{self.table_name}/pdate={self.date}/phour={self.hour}" return [self.date, self.hour, stage_path] def get_job_frequency(self) -> str: return "HOURLY" class HourlyRDStoS3JobCustomSampleQuery(HourlyRDStoS3Job): """ Represents an RDS to S3 job that executes hourly. Uses a custom sample query. """ def __init__( self, table_name: str, procedure_name: str, task_monitor_name: str, sample_query: str, prod_warehouse: str = "SUNO_PROD_GLUE_HOURLY_X_SMALL", staging_warehouse: str = "SUNO_STAGING_GLUE_HOURLY_X_SMALL", print_debug_logs: bool = False, ): HourlyRDStoS3Job.__init__( self, table_name, procedure_name, task_monitor_name, prod_warehouse=prod_warehouse, staging_warehouse=staging_warehouse, print_debug_logs=print_debug_logs, ) self.sample_query = sample_query def get_postgres_sample_query(self) -> str: return self.sample_query class HourlyRDStoS3JobNoCreatedAtColumnPostgres(HourlyRDStoS3Job): """ Represents an hourly RDS to S3 job to be used if the Postgres table doesn't have a "created_at" column. Uses a slightly different sample query. """ def get_postgres_sample_query(self) -> str: return f"SELECT * FROM {self.table_name} WHERE updated_at >= '{self.hour_start}' AND updated_at < '{self.hour_end}'" class DailyRDStoS3Job(RDStoS3Job): """ Represents an RDS to S3 job that executes daily. Uses a default sample query and transform SQL. """ def __init__( self, table_name: str, procedure_name: str, task_monitor_name: str, prod_warehouse: str = "SUNO_PROD_GLUE_HOURLY_X_SMALL", staging_warehouse: str = "SUNO_STAGING_GLUE_HOURLY_X_SMALL", print_debug_logs: bool = False, ): RDStoS3Job.__init__( self, table_name, procedure_name, task_monitor_name, prod_warehouse=prod_warehouse, staging_warehouse=staging_warehouse, print_debug_logs=print_debug_logs, ) current_date = datetime.now() one_day_ago = current_date - timedelta(days=1) self.day_start = one_day_ago.replace(hour=0, minute=0, second=0, microsecond=0) self.day_end = self.day_start + timedelta(days=1) self.date = one_day_ago.strftime("%Y-%m-%d") def get_postgres_sample_query(self) -> str: return f"SELECT * FROM {self.table_name} WHERE (updated_at >= '{self.day_start}' AND updated_at < '{self.day_end}') OR (created_at >= '{self.day_start}' AND created_at < '{self.day_end}')" def get_postgres_transform_sql(self) -> str: return f""" select *, '{self.date}' as pdate from result """ def get_postgres_partition_keys(self) -> list[str]: return ["pdate"] def get_snowflake_procedure_sql_parameters(self) -> list[str]: stage_path = f"@SUNO_DATABASE_EVENTS/{self.table_name}/pdate={self.date}" return [self.date, stage_path] def get_job_frequency(self) -> str: return "DAILY" class DailyRDStoS3JobCustomSampleQuery(DailyRDStoS3Job): """ Represents an RDS to S3 job that executes daily. Uses a custom sample query. """ def __init__( self, table_name: str, procedure_name: str, task_monitor_name: str, sample_query: str, prod_warehouse: str = "SUNO_PROD_GLUE_HOURLY_X_SMALL", staging_warehouse: str = "SUNO_STAGING_GLUE_HOURLY_X_SMALL", print_debug_logs: bool = False, ): DailyRDStoS3Job.__init__( self, table_name, procedure_name, task_monitor_name, prod_warehouse=prod_warehouse, staging_warehouse=staging_warehouse, print_debug_logs=print_debug_logs, ) self.sample_query = sample_query def get_postgres_sample_query(self) -> str: return self.sample_query class HourlyRDStoS3JobBackfill(RDStoS3Job): """ Represents an RDS to S3 job that can backfill data for a specific date and hour. Accepts p_date and p_hour parameters to process data for that specific hour. """ def __init__( self, table_name: str, procedure_name: str, task_monitor_name: str, p_date: str, p_hour: str, prod_warehouse: str = "SUNO_PROD_GLUE_HOURLY_X_SMALL", staging_warehouse: str = "SUNO_STAGING_GLUE_HOURLY_X_SMALL", print_debug_logs: bool = False, ): RDStoS3Job.__init__( self, table_name, procedure_name, task_monitor_name, prod_warehouse=prod_warehouse, staging_warehouse=staging_warehouse, print_debug_logs=print_debug_logs, ) # Parse the provided date and hour self.date = p_date self.hour = p_hour # Calculate hour_start and hour_end from p_date and p_hour date_obj = datetime.strptime(p_date, "%Y-%m-%d") hour_int = int(p_hour) self.hour_start = date_obj.replace(hour=hour_int, minute=0, second=0, microsecond=0) self.hour_end = self.hour_start + timedelta(hours=1) def get_postgres_sample_query(self) -> str: return f"SELECT * FROM {self.table_name} WHERE (updated_at >= '{self.hour_start}' AND updated_at < '{self.hour_end}') OR (created_at >= '{self.hour_start}' AND created_at < '{self.hour_end}')" def get_postgres_transform_sql(self) -> str: return f""" select *, '{self.date}' as pdate, '{self.hour}' as phour from result """ def get_postgres_partition_keys(self) -> list[str]: return ["pdate", "phour"] def get_snowflake_procedure_sql_parameters(self) -> list[str]: stage_path = f"@SUNO_DATABASE_EVENTS/{self.table_name}/pdate={self.date}/phour={self.hour}" return [self.date, self.hour, stage_path] def get_job_frequency(self) -> str: return "BACKFILL" class HourlyRDStoS3JobBackfillCustomSampleQuery(HourlyRDStoS3JobBackfill): """ Represents an RDS to S3 job that can backfill data for a specific date and hour. Uses a custom sample query instead of the default one. """ def __init__( self, table_name: str, procedure_name: str, task_monitor_name: str, p_date: str, p_hour: str, sample_query: str, prod_warehouse: str = "SUNO_PROD_GLUE_HOURLY_X_SMALL", staging_warehouse: str = "SUNO_STAGING_GLUE_HOURLY_X_SMALL", print_debug_logs: bool = False, ): HourlyRDStoS3JobBackfill.__init__( self, table_name, procedure_name, task_monitor_name, p_date, p_hour, prod_warehouse=prod_warehouse, staging_warehouse=staging_warehouse, print_debug_logs=print_debug_logs, ) self.sample_query = sample_query def get_postgres_sample_query(self) -> str: return self.sample_query def _run_job_with_existing_context(job_instance, glueContext, spark): """ Helper function to run a job using existing GlueContext and Spark session instead of creating new ones. This avoids SparkContext conflicts. """ from utils.util import get_data_from_postgresql_and_save_to_s3 from utils.snowflake.snowflake_client import get_snowflake_connection job_instance.debug_print(f"Environment: {job_instance.environment}") job_instance.debug_print(f"S3 Bucket: {job_instance.s3_bucket}") job_instance.debug_print(f"Warehouse: {job_instance.warehouse}") job_instance.debug_print(f"DB Name: {job_instance.db_name}") POSTGRES_CONNECTION_NAME = "analystic-database-connection" row_count = get_data_from_postgresql_and_save_to_s3( db_connection_options={ "useConnectionProperties": "true", "dbtable": job_instance.table_name, "connectionName": POSTGRES_CONNECTION_NAME, "sampleQuery": job_instance.get_postgres_sample_query(), }, s3_connection_options={ "path": f"s3://{job_instance.s3_bucket}/{job_instance.table_name}/", "partitionKeys": job_instance.get_postgres_partition_keys(), }, transform_sql=job_instance.get_postgres_transform_sql(), glueContext=glueContext, spark=spark, ) job_instance.debug_print(f"Saved data to S3. Row count: {row_count}") task_monitor_sql_params = [ job_instance.task_monitor_name, job_instance.start_time, row_count, "SUCCESS", job_instance.get_job_frequency().upper(), ] snowflake_procedure_sql = job_instance.get_call_snowflake_procedure_sql( job_instance.procedure_name, job_instance.get_snowflake_procedure_sql_parameters() ) task_monitor_sql = job_instance.get_call_snowflake_procedure_sql( "TASK_MONITOR_INSERT_PROC", task_monitor_sql_params ) connection = get_snowflake_connection( job_instance.warehouse, job_instance.db_name, job_instance.environment ) cursor = connection.cursor() job_instance.debug_print(f"Calling SQL: {snowflake_procedure_sql}") try: result = cursor.execute(snowflake_procedure_sql) job_instance.debug_print(f"Calling SQL: {task_monitor_sql}") result = cursor.execute(task_monitor_sql) except Exception as e: print(f"Error executing SQL: {e}") raise e job_instance.debug_print("Query executed successfully.") def backfill_data( glueContext, spark, table_name: str, procedure_name: str, base_sample_query: str, task_monitor_name: str = 'TASK_MONITOR_INSERT_PROC', start_time=None, end_time=None, mode="hourly", prod_warehouse: str = "SUNO_PROD_GLUE_HOURLY_X_SMALL", staging_warehouse: str = "SUNO_STAGING_GLUE_HOURLY_X_SMALL", print_debug_logs: bool = False ): """ Enhanced backfill utility function that supports hourly, daily, and all data backfill for a time range. Creates and uses the appropriate job class based on the mode. Args: glueContext: Glue context spark: Spark session table_name: Name of the table to backfill procedure_name: Snowflake procedure name for upsert task_monitor_name: Task monitor name for tracking base_sample_query: Base SQL query (time filters will be added automatically for hourly/daily modes) start_time: Start datetime for backfill range (datetime object, required for hourly/daily) end_time: End datetime for backfill range (datetime object, required for hourly/daily) mode: "hourly", "daily", or "all" - determines the granularity of backfill prod_warehouse: Production warehouse name staging_warehouse: Staging warehouse name print_debug_logs: Whether to print debug logs """ if mode == "all": # Process all data using the original custom sample query print(f"Running backfill for all data in table {table_name}") # Use HourlyRDStoS3JobCustomSampleQuery for all mode with the base query job = HourlyRDStoS3JobCustomSampleQuery( table_name=table_name, procedure_name=procedure_name, task_monitor_name=task_monitor_name, sample_query=base_sample_query, prod_warehouse=prod_warehouse, staging_warehouse=staging_warehouse, print_debug_logs=print_debug_logs ) # Use existing glueContext and spark instead of run_job() which creates new SparkContext _run_job_with_existing_context(job, glueContext, spark) return # Validate inputs for time-based backfill if start_time is None or end_time is None: raise ValueError("start_time and end_time are required for hourly and daily modes") if not isinstance(start_time, datetime) or not isinstance(end_time, datetime): raise ValueError("start_time and end_time must be datetime objects") if start_time >= end_time: raise ValueError("start_time must be before end_time") if mode not in ["hourly", "daily"]: raise ValueError("mode must be 'hourly', 'daily', or 'all'") # Determine time increment and job class if mode == "hourly": time_delta = timedelta(hours=1) job_class = HourlyRDStoS3JobCustomSampleQuery print(f"Running hourly backfill from {start_time} to {end_time}") else: # daily time_delta = timedelta(days=1) job_class = DailyRDStoS3JobCustomSampleQuery print(f"Running daily backfill from {start_time.date()} to {end_time.date()}") # Loop through time range current_time = start_time total_periods = 0 successful_periods = 0 while current_time < end_time: next_time = current_time + time_delta # Don't exceed end_time if next_time > end_time: next_time = end_time try: # Create time-filtered sample query if "WHERE" in base_sample_query.upper(): # If original query has WHERE clause, add AND condition time_filter = f" AND created_at >= TIMESTAMP '{current_time.strftime('%Y-%m-%d %H:%M:%S')}' AND created_at < TIMESTAMP '{next_time.strftime('%Y-%m-%d %H:%M:%S')}'" # Insert the time filter before ORDER BY, LIMIT, etc. if they exist query_parts = base_sample_query.split() insert_pos = len(query_parts) for i, part in enumerate(query_parts): if part.upper() in ['ORDER', 'GROUP', 'HAVING', 'LIMIT', 'OFFSET']: insert_pos = i break query_before = ' '.join(query_parts[:insert_pos]) query_after = ' '.join(query_parts[insert_pos:]) if insert_pos < len(query_parts) else '' time_filtered_query = query_before + time_filter + (' ' + query_after if query_after else '') else: # If no WHERE clause, add one cleaned_query = base_sample_query.rstrip() if cleaned_query.endswith(';'): cleaned_query = cleaned_query.rstrip(';') time_filtered_query = cleaned_query + f" WHERE created_at >= TIMESTAMP '{current_time.strftime('%Y-%m-%d %H:%M:%S')}' AND created_at < TIMESTAMP '{next_time.strftime('%Y-%m-%d %H:%M:%S')}'" period_desc = f"{current_time.strftime('%Y-%m-%d %H:%M')} to {next_time.strftime('%Y-%m-%d %H:%M')}" print(f"Processing period: {period_desc}") # Create appropriate job instance for this time period job = job_class( table_name=table_name, procedure_name=procedure_name, task_monitor_name=task_monitor_name, sample_query=time_filtered_query, prod_warehouse=prod_warehouse, staging_warehouse=staging_warehouse, print_debug_logs=print_debug_logs ) # Run the job for this time period using existing context _run_job_with_existing_context(job, glueContext, spark) successful_periods += 1 print(f"Successfully completed period: {period_desc}") except Exception as e: print(f"Error processing period {current_time} to {next_time}: {str(e)}") # Continue with next period rather than failing entire job print(f"Continuing with next period...") total_periods += 1 current_time = next_time print(f"Backfill completed. Processed {successful_periods}/{total_periods} periods successfully.") if successful_periods == 0: raise Exception("All periods failed during backfill") elif successful_periods < total_periods: print(f"Warning: {total_periods - successful_periods} periods failed during backfill")