NRL-1268 configure timestamp and partition by date

Author: jackleary

terraform/account-wide-infrastructure/modules/glue/glue.tf

@@ -52,6 +52,7 @@ resource "aws_glue_job" "glue_job" {
     "--source_path"                     = "s3://${aws_s3_bucket.source-data-bucket.id}/"     # Specify the source S3 path
     "--target_path"                     = "s3://${aws_s3_bucket.target-data-bucket.id}/logs" # Specify the destination S3 path
     "--job_name"                        = "${var.name_prefix}-glue-job"
+    "--partition_cols"                  = "date"
     "--enable-continuous-log-filter"    = "true"
     "--enable-metrics"                  = "true"
     "--extra-py-files"                  = "s3://${aws_s3_bucket.code-bucket.id}/src.zip"

terraform/account-wide-infrastructure/modules/glue/src/main.py

@@ -6,7 +6,9 @@
 from transformations import dtype_conversion, flatten_df, logSchema
 
 # Get arguments from AWS Glue job
-args = getResolvedOptions(sys.argv, ["job_name", "source_path", "target_path"])
+args = getResolvedOptions(
+    sys.argv, ["job_name", "source_path", "target_path", "partition_cols"]
+)
 
 # Start Glue context
 sc = SparkContext()

terraform/account-wide-infrastructure/modules/glue/src/transformations.py

@@ -1,15 +1,21 @@
-from pyspark.sql.functions import col, to_timestamp
+from pyspark.sql.functions import (
+    col,
+    from_unixtime,
+    regexp_replace,
+    to_date,
+    to_timestamp,
+)
 from pyspark.sql.types import (
     BooleanType,
+    DoubleType,
     StringType,
     StructField,
     StructType,
-    TimestampType,
 )
 
 logSchema = StructType(
     [
-        StructField("time", TimestampType(), True),
+        StructField("time", DoubleType(), True),
         StructField("index", StringType(), True),
         StructField("host", StringType(), True),
         StructField("source", StringType(), True),
@@ -80,8 +86,15 @@ def flatten(schema, prefix=""):
 
 
 def dtype_conversion(df):
-    df = df.withColumn(
-        "event_timestamp",
-        to_timestamp(df["event_timestamp"], "yyyy-MM-dd HH:mm:ss,SSSXXX"),
+    df = (
+        df.withColumn(
+            "event_timestamp_cleaned", regexp_replace(col("event_timestamp"), ",", ".")
+        )
+        .withColumn(
+            "event_timestamp",
+            to_timestamp(col("event_timestamp_cleaned"), "yyyy-MM-dd HH:mm:ss.SSSZ"),
+        )
+        .withColumn("time", from_unixtime(col("time")).cast("timestamp"))
+        .withColumn("date", to_date(col("time")))
     )
-    return df
+    return df.drop("event_timestamp_cleaned")