#=========================================
###  POSISI : RABU, 18 AGUSTUS 2021
###  PIC    : ANDRI ARIYANTO
###  CASE   : MIGRATE TABLES
#=========================================
##           Import Dependencies           
#=========================================
import re
import string
from time import time, sleep
from pytz import timezone
from datetime import datetime, timedelta
from pyspark.sql import SparkSession, functions as F, types as T

##              Pandas Config              
#=========================================
import pandas as pd
pd.options.display.html.table_schema=True
pd.options.display.max_columns=999
pd.options.display.max_rows=999

##              Spark SHession              
#=========================================
spark = SparkSession\
  .builder\
  .appName("MIGRATE TABLES")\
  .config("spark.sql.crossJoin.enabled", "true")\
  .config("spark.dynamicAllocation.enabled", "false")\
  .config("spark.executor.instances", "4")\
  .config("spark.executor.cores", "5")\
  .config("spark.executor.memory", "8g")\
  .config("spark.yarn.executor.memoryOverhead", "4g")\
  .config("spark.ui.showConsoleProgress", "false")\
  .config("spark.network.timeout", 60)\
  .enableHiveSupport()\
  .getOrCreate()
        
##               Functions                 
#=========================================
def set_timer():
  global START_TIME
  START_TIME = time()

def get_timer(start_time=None):
  if start_time:
    t = datetime(1,1,1)+timedelta(seconds=int(time()-start_time))
  else:
    t = datetime(1,1,1)+timedelta(seconds=int(time()-START_TIME))
  return "{}:{}:{}".format(str(t.hour).zfill(2),str(t.minute).zfill(2),str(t.second).zfill(2))

def get_partition(schema, table):
  try:
    return spark.sql("""
     SHOW PARTITIONS {}.{}
     """.format(schema, table)).sort('partition').collect() # ambil partisi sesuai format
  except:
    return None

def get_partition_format(partition, sep=" AND "):
  list_partition = []
  for value in partition.split('/'):
    split_value = value.split("=")
    list_partition.append("{}='{}'".format(split_value[0], split_value[1]))
  return "{}".format(sep).join(list_partition)  

def get_partition_column(partition):
  list_partition = []
  for value in partition.split('/'):
    split_value = value.split("=")
    list_partition.append(split_value[0])
  return list_partition


def write_migrate_table(source_schema, source_table, target_schema, target_table, partitioned=False, partition=None):
  if partitioned:
    df = spark.read.table("{}.{}".format(source_schema, source_table)).filter("{}".format(get_partition_format(partition))) # load tabel sumber dengan filter partisi
  else:
    df = spark.read.table("{}.{}".format(source_schema, source_table)) # load tabel sumber

  count = df.count()
  if count != 0:
    df = df.withColumn("modified_date", F.lit(datetime.now(timezone("Asia/Jakarta"))))

    path = lambda p: spark._jvm.org.apache.hadoop.fs.Path(p)
    fs = spark._jvm.org.apache.hadoop.fs.FileSystem.get(spark._jsc.hadoopConfiguration())
    
    if partitioned:
      hdfs_folder = fs.getContentSummary(path("/user/hive/warehouse/{}.db/{}/{}".format(source_schema, source_table, partition)))
      repartition_number = float(hdfs_folder.getLength())/float(128*1024*1024)
      
      df = df.repartition(int(1 if repartition_number < 1 else repartition_number))

      if spark.sql("SHOW TABLES IN {} LIKE '{}'".format(target_schema, target_table)).count() != 0:
        query_drop_partition = "ALTER TABLE {}.{} DROP IF EXISTS PARTITION({})".format(target_schema, target_table, get_partition_format(partition, ","))
        spark.sql(query_drop_partition) # menghapus partisi di spark

      df\
        .write\
        .format("parquet")\
        .partitionBy(get_partition_column(partition))\
        .mode("append")\
        .saveAsTable("{}.{}".format(target_schema,target_table))

      spark.sql("MSCK REPAIR TABLE {}.{}".format(target_schema, target_table)) # memperbaiki tabel
      spark.sql("ANALYZE TABLE {}.{} PARTITION ({}) COMPUTE STATISTICS".format(target_schema, target_table, get_partition_format(partition, ","))) # kalkulasi data
    else:
      hdfs_folder = fs.getContentSummary(path("/user/hive/warehouse/{}.db/{}".format(source_schema, source_table)))
      repartition_number = float(hdfs_folder.getLength())/float(128*1024*1024)

      df = df.repartition(int(1 if repartition_number < 1 else repartition_number))

      df\
        .write\
        .format("parquet")\
        .mode("overwrite")\
        .saveAsTable("{}.{}".format(target_schema,target_table))

      spark.sql("ANALYZE TABLE {}.{} COMPUTE STATISTICS".format(target_schema, target_table)) # kalkulasi data
    
    spark.sql("REFRESH TABLE {}.{}".format(target_schema, target_table)) # refresh tabel di spark
  print("Total Record  = {}".format(count))



##                Parameter                
#=========================================
list_table = [
#"bribrain_brimo_segmentation_fact_summary_trx_ebor_merchant",
"bribrain_brimo_segmentation_fact_summary_investation",
"bribrain_brimo_segmentation_fact_summary_trx_prosw_merchant", 
"bribrain_brimo_segmentation_fact_summary_trx_prosw_non_merchant", 
"bribrain_brimo_segmentation_fact_summary_trx_brimo", 
"bribrain_brimo_segmentation_fact_customer_portfolio", 
"bribrain_brimo_segmentation_fact_dataset_inference", 
"bribrain_brimo_segmentation_fact_customer_demography", 
"bribrain_brimo_segmentation_fact_model_result"]

for table in list_table:

  first = True
  set_timer()


  source_schema = "datamart" # schema sumber
  source_table = "backup_202203_{}".format(table) # tabel sumber

  target_schema = "rpt" # schema target
  target_table = "backup_202203_{}".format(table) # tabel target

  ##             Main Process                
  #=========================================

  list_partition = get_partition(source_schema, source_table) # mendapatkan list partisi tabel

  
#  list_partition = list_partition[-1:]
  
  print(list_partition)

  if list_partition == None: # kondisi ketika tabel tidak berpartisi
    print("Not a partitioned table!")
    if first:
      start_time = START_TIME
      first = False
    write_migrate_table(
      source_schema, 
      source_table,
      target_schema, 
      target_table,
      False,
      None) # proses write mode overwrite untuk tabel tidak berpartisi
  else: # kondisi ketika tabel berpartisi
    print("Partitioned table!")
    for row in list_partition: # melakukan loop atas list partisi
      set_timer()
      if first:
        start_time = START_TIME
        first = False
      write_migrate_table(
        source_schema, 
        source_table,
        target_schema,
        target_table,
        True,
        row['partition']) # proses write mode append untuk tabel berpartisi
      str_partition = get_partition_format(row['partition']) # mengubah format partisi menjadi str input
      print("Process took {} for {}\n".format(get_timer(), str_partition))

  ##              ETL Time                 
  #=========================================
  print("All Process took {}".format(get_timer(start_time)))

#spark.stop()