using utils folder for testing and iceberg format

Author: vitor-pina

src/bronze/get_full_tables.py

@@ -5,6 +5,7 @@
 import os
 import logging
 from typing import List, Union, Optional,Tuple
+from utils.utils import Utils
 
 logging.basicConfig(
     level=logging.INFO,
@@ -13,11 +14,9 @@
         logging.StreamHandler()
     ]
 )
-def get_required_env(env_name:str) -> str:
-    env_value = os.getenv(env_name)
-    if env_value is None:
-        raise ValueError(f"Environment variable {env_name} is not set")
-    return env_value
+
+utils_obj = Utils()
+
 
 def get_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser()
@@ -26,20 +25,6 @@ def get_args() -> argparse.Namespace:
     args = parser.parse_args()
     return args
 
-def get_spark_session(S3_ACCESS_KEY: str,S3_SECRET_KEY: str , S3_ENDPOINT: str) -> SparkSession:
-    
-    spark = SparkSession.builder \
-        .appName("full_table_ingestion") \
-        .config("spark.jars", "/opt/spark/jars/hadoop-aws-3.3.4.jar,/opt/spark/jars/aws-java-sdk-bundle-1.12.375.jar,/opt/spark/jars/delta-spark_2.12-3.2.1.jar,/opt/spark/jars/delta-storage-3.2.1.jar,/opt/spark/jars/delta-kernel-api-3.2.1.jar,/opt/spark/jars/mysql-connector-j-8.3.0.jar") \
-        .config("spark.sql.extensions", "io.delta.sql.DeltaSparkSessionExtension")\
-        .config("spark.sql.catalog.spark_catalog", "org.apache.spark.sql.delta.catalog.DeltaCatalog")\
-        .config("spark.hadoop.fs.s3a.access.key", S3_ACCESS_KEY) \
-        .config("spark.hadoop.fs.s3a.secret.key", S3_SECRET_KEY) \
-        .config("spark.hadoop.fs.s3a.endpoint", S3_ENDPOINT) \
-        .config("spark.hadoop.fs.s3a.path.style.access", "true") \
-        .config("spark.hadoop.fs.s3a.impl", "org.apache.hadoop.fs.s3a.S3AFileSystem") \
-        .getOrCreate()
-    return spark 
 
 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))
@@ -52,16 +37,16 @@ def add_date_partition_columns(df: DataFrame,column_name:str) -> DataFrame:
     return df
 
 def main() -> None:
-    MYSQL_DATABASE = get_required_env("MYSQL_DATABASE")
-    MYSQL_HOST = get_required_env("MYSQL_HOST")
-    MYSQL_PORT = get_required_env("MYSQL_PORT")
-    MYSQL_USER = get_required_env("MYSQL_USER")
-    MYSQL_SECRET = get_required_env("MYSQL_SECRET")
+    MYSQL_DATABASE = utils_obj.get_required_env("MYSQL_DATABASE")
+    MYSQL_HOST = utils_obj.get_required_env("MYSQL_HOST")
+    MYSQL_PORT = utils_obj.get_required_env("MYSQL_PORT")
+    MYSQL_USER = utils_obj.get_required_env("MYSQL_USER")
+    MYSQL_SECRET = utils_obj.get_required_env("MYSQL_SECRET")
     jdbc_url = f"jdbc:mysql://{MYSQL_HOST}:{MYSQL_PORT}/{MYSQL_DATABASE}"
 
-    S3_ACCESS_KEY = get_required_env("S3_ACCESS_KEY")
-    S3_SECRET_KEY = get_required_env("S3_SECRET_KEY")
-    S3_ENDPOINT = get_required_env("S3_ENDPOINT")
+    S3_ACCESS_KEY = utils_obj.get_required_env("S3_ACCESS_KEY")
+    S3_SECRET_KEY = utils_obj.get_required_env("S3_SECRET_KEY")
+    S3_ENDPOINT = utils_obj.get_required_env("S3_ENDPOINT")
     args = get_args()
     S3_SAVEPATH = args.savepath
     undesired_column = args.undesired_column
@@ -79,7 +64,7 @@ def main() -> None:
     "auth_user"
     ]
 
-    spark = get_spark_session(S3_ACCESS_KEY=S3_ACCESS_KEY,S3_SECRET_KEY=S3_SECRET_KEY,S3_ENDPOINT=S3_ENDPOINT)
+    spark = utils_obj.get_spark_session(S3_ACCESS_KEY=S3_ACCESS_KEY,S3_SECRET_KEY=S3_SECRET_KEY,S3_ENDPOINT=S3_ENDPOINT,app_name="Full table ingestion")
     for table in TABLES:
 
         logging.info(f"getting table {table}")
@@ -102,7 +87,7 @@ def main() -> None:
 
             output_path = f"{S3_SAVEPATH}/{table}"
 
-            df.write.format("delta").mode("append").partitionBy("year", "month","day").save(output_path)
+            df.write.format("iceberg").mode("append").partitionBy("year", "month","day").save(output_path)
 
             logging.info(f"Data saved as Delta table to {output_path}")
 

src/utils/__init__.py

@@ -1,7 +1,7 @@
 import os
 from pyspark.sql import SparkSession #type:ignore
 
-class utils:
+class Utils:
     def __init__(self) -> None:
         pass
 
@@ -15,9 +15,13 @@ def get_spark_session(self,S3_ACCESS_KEY: str,S3_SECRET_KEY: str , S3_ENDPOINT:
 
         spark = SparkSession.builder \
             .appName(app_name) \
-            .config("spark.jars", "/opt/spark/jars/hadoop-aws-3.3.4.jar,/opt/spark/jars/aws-java-sdk-bundle-1.12.375.jar,/opt/spark/jars/delta-spark_2.12-3.2.1.jar,/opt/spark/jars/delta-storage-3.2.1.jar,/opt/spark/jars/delta-kernel-api-3.2.1.jar,/opt/spark/jars/mysql-connector-j-8.3.0.jar") \
-            .config("spark.sql.extensions", "io.delta.sql.DeltaSparkSessionExtension")\
-            .config("spark.sql.catalog.spark_catalog", "org.apache.spark.sql.delta.catalog.DeltaCatalog")\
+            .config("spark.jars", 
+                    "/opt/spark/jars/hadoop-aws-3.3.4.jar," 
+                    "/opt/spark/jars/aws-java-sdk-bundle-1.12.375.jar," 
+                    "/opt/spark/jars/mysql-connector-j-8.3.0.jar," 
+                    "/opt/spark/jars/iceberg-spark-runtime-3.5_2.12-1.10.0.jar") \
+            .config("spark.sql.extensions", "org.apache.iceberg.spark.extensions.IcebergSparkSessionExtensions")\
+            .config("spark.sql.catalog.spark_catalog", "org.apache.iceberg.spark.SparkSessionCatalog")\
             .config("spark.hadoop.fs.s3a.access.key", S3_ACCESS_KEY) \
             .config("spark.hadoop.fs.s3a.secret.key", S3_SECRET_KEY) \
             .config("spark.hadoop.fs.s3a.endpoint", S3_ENDPOINT) \