turning redudant code into function

Author: vitor-pina

src/bronze/incremental_load.py

@@ -1,5 +1,6 @@
 from datetime import datetime
 from pyspark.sql import SparkSession # type: ignore
+from pyspark.sql import DataFrame #type:ignore
 import pyspark.sql.functions as F # type: ignore
 import pyspark.sql.types as T # type: ignore
 import argparse
@@ -32,6 +33,18 @@ def get_args() -> argparse.Namespace:
     args = parser.parse_args()
     return args
 
+
+def get_metadata(metadatapath: str, spark: SparkSession, table:str) -> str | bool:
+    metadatapath = f"{metadatapath}/{table}/last_updated_date"
+    customSchema = T.StructType([      
+        T.StructField("table_name", T.StringType(), True),
+        T.StructField("last_date", T.StringType(), True)
+    ])
+    row = spark.read.csv(metadatapath,schema=customSchema).filter( F.col("table_name") == table).first()
+    if row is None or row["last_date"] is None:
+        return False
+    return str(row["last_date"])
+
 def update_metadata(metadatapath: str ,spark: SparkSession,table:str,last_date:str) -> bool:
     try:
         tmp_df = spark.createDataFrame(
@@ -68,8 +81,14 @@ def get_spark_session(S3_ACCESS_KEY: str,S3_SECRET_KEY: str , S3_ENDPOINT: str)
 
 
 
+def add_ingestion_metadata_column(df: DataFrame,table: str) -> DataFrame:
+    tmp_df = df.withColumn("ingestion_date", F.current_timestamp()).withColumn("source_name", F.lit(table))
+    return tmp_df
 
-
+def add_date_partition_columns(df: DataFrame,column_name:str) -> DataFrame:
+    df = df.withColumn("year", F.year(F.col(column_name)))\
+        .withColumn("month", F.month(F.col(column_name)))
+    return df
 
 def full_initial_ingestion(spark: SparkSession, table: str, savepath: str, jdbc_url:str, MYSQL_USER:str, MYSQL_SECRET:str) -> Tuple[bool, str]:
     current_year = datetime.now().year
@@ -82,6 +101,8 @@ def full_initial_ingestion(spark: SparkSession, table: str, savepath: str, jdbc_
 
     for year in years:
         for month in months:
+            if (year > current_year) or (year == current_year and month > current_month):
+                break 
 
             if (year == 2019 and month ==1):
                 query = f"(SELECT * FROM {table} WHERE YEAR(created) = {year} AND MONTH(created) = {month}) AS limited_table"
@@ -93,22 +114,15 @@ def full_initial_ingestion(spark: SparkSession, table: str, savepath: str, jdbc_
                     .option("driver", "com.mysql.cj.jdbc.Driver") \
                     .option("dbtable", query) \
                     .load() 
-
-                df = df.withColumn("created", F.date_format("created", "yyyy-MM-dd'T'HH:mm:ss.SSSSSS")) \
-                       .withColumn("year", F.year(F.col("created"))) \
-                       .withColumn("month", F.month(F.col("created")))
-                
-                df = df.withColumn("ingestion_date", F.current_timestamp()) \
-                   .withColumn("source_name", F.lit(table))
+                df = df.withColumn("created", F.date_format("created", "yyyy-MM-dd'T'HH:mm:ss.SSSSSS"))
+                df = add_date_partition_columns(df=df, column_name="created")
+                df = add_ingestion_metadata_column(df=df,table=table)
 
                 df.write.format("delta") \
                     .mode("overwrite") \
                     .partitionBy("year", "month") \
                     .save(path)
                 continue
-            
-            if (year > current_year) or (year == current_year and month > current_month):
-                break 
             else:
 
                 query = f"(SELECT * FROM {table} WHERE YEAR(created) = {year} AND MONTH(created) = {month}) AS limited_table"
@@ -122,42 +136,35 @@ def full_initial_ingestion(spark: SparkSession, table: str, savepath: str, jdbc_
                     .load()
                 if (year == current_year and month == current_month):
                     last_update =  datetime.now().isoformat()
-                incremental_df = new_df.withColumn("created", F.date_format("created", "yyyy-MM-dd'T'HH:mm:ss.SSSSSS")) \
-                                       .withColumn("year", F.year(F.col("created"))) \
-                                       .withColumn("month", F.month(F.col("created")))
+                incremental_df = new_df.withColumn("created", F.date_format("created", "yyyy-MM-dd'T'HH:mm:ss.SSSSSS")) 
+                incremental_df = add_date_partition_columns(df=incremental_df, column_name="created")
 
-                incremental_df = incremental_df.withColumn("ingestion_date", F.current_timestamp()).withColumn("source_name", F.lit(table))
+                incremental_df = add_ingestion_metadata_column(df=incremental_df,table=table)
 
                 # Append new partitions directly
                 incremental_df.write.format("delta").mode("append").partitionBy("year", "month").save(path)
     return (True, last_update)
 
-def get_metadata(metadatapath: str, spark: SparkSession, table:str) -> str | bool:
-    metadatapath = f"{metadatapath}/{table}/last_updated_date"
-    customSchema = T.StructType([      
-        T.StructField("table_name", T.StringType(), True),
-        T.StructField("last_date", T.StringType(), True)
-    ])
-    row = spark.read.csv(metadatapath,schema=customSchema).filter( F.col("table_name") == table).first()
-    if row is None or row["last_date"] is None:
-        return False
-    return str(row["last_date"])
 
 def delta_load(spark: SparkSession, jdbc_url:str, MYSQL_USER:str, MYSQL_SECRET:str,last_updated:str,table:str,savepath: str) -> Tuple[bool, str]:
+    
     path = f"{savepath}/{table}"
     query = f"(SELECT * FROM {table} WHERE created >= '{last_updated}') AS limited_table"
+    
     logging.info(query)            
+    
     new_df = spark.read.format("jdbc") \
         .option("url", jdbc_url) \
         .option("user", MYSQL_USER) \
         .option("password", MYSQL_SECRET) \
         .option("driver", "com.mysql.cj.jdbc.Driver") \
         .option("dbtable", query) \
         .load()
+    
     last_update =  datetime.now().isoformat()
-    incremental_df = new_df.withColumn("created", F.date_format("created", "yyyy-MM-dd'T'HH:mm:ss.SSSSSS")) \
-                           .withColumn("year", F.year(F.col("created"))) \
-                           .withColumn("month", F.month(F.col("created")))
+    incremental_df = new_df.withColumn("created", F.date_format("created", "yyyy-MM-dd'T'HH:mm:ss.SSSSSS")) 
+    incremental_df = add_date_partition_columns(df=incremental_df,column_name="created")
+    incremental_df = add_ingestion_metadata_column(df=incremental_df,table=table)
 
     incremental_df = incremental_df.withColumn("ingestion_date", F.current_timestamp()).withColumn("source_name", F.lit(table))
 
@@ -185,20 +192,27 @@ def main() -> None:
     metadata = args.metadatapath
     is_full_ingestion_flag = args.first_ingestion_flag
     table = args.table
+
+
     spark = get_spark_session(S3_ACCESS_KEY,S3_SECRET_KEY,S3_ENDPOINT)
+    
     if is_full_ingestion_flag == 1:
         result = full_initial_ingestion(spark,table,savepath,jdbc_url,MYSQL_USER,MYSQL_SECRET)
         logging.info(result)
         if result[0]:
             update_metadata(metadatapath=metadata,spark=spark,table=table,last_date=result[1])
+    
     if is_full_ingestion_flag == 0:
         last_date = get_metadata(metadatapath=metadata,spark=spark,table=table)
         if last_date == False:
             raise Exception("No date Found")
+        
         last_date = str(last_date)
         result = delta_load(spark=spark,jdbc_url=jdbc_url,MYSQL_USER=MYSQL_USER,MYSQL_SECRET=MYSQL_SECRET,last_updated=last_date,table=table,savepath=savepath)
+        
         if result[0]:
              update_metadata(metadatapath=metadata,spark=spark,table=table,last_date=result[1])
+    
     spark.stop()
 
 if __name__=="__main__":