#=============================================================
###  CREATED AT     : 04 MARET 2024                         
###  UPDATED AT     : 24 FEBRUARI 2026                      
###  COPYRIGHT      : ANDRI ARIYANTO                        
###  DESCRIPTION    : Module untuk melakukan load data      
#=============================================================

## HDFS ==============================================================================================
from pytz import timezone
from datetime import datetime
from typing import Optional
from pyspark.sql import functions as F

def load_to_hive(
  spark, 
  df, 
  schema: str, 
  table: str, 
  partition: str = None, 
  mode: str = 'append', 
  repartition_number: Optional[int] = None, 
  blocksize: Optional[int] = 256,
  modified_date: bool = True,
  schema_temp: str = "sharingvision",
  count: bool = False):
  """Melakukan penulisan pyspark dataframe sebagai tabel hive di hdfs
  
  Args:
      spark (pyspark.sql.session.SparkSession): spark session
      df (pyspark.sql.dataframe.DataFrame): pyspark dataframe yang akan di tulis sebagai tabel hive
      schema (str): target schema di hdfs
      table (str): target table di hdfs
      partition (str): kolom yang akan dijadikan partisi (dipisahkan '|' jika lebih dari satu kolom)
          (default is None)
      mode (str): cara yang dilakukan untuk penulisan tabel
          (optional is append, overwrite, error, ignore)
          (default is append)
      repartition_number (int | None): jumlah file parquet yang diinginkan dalam satu partisi
          (default is None)      
      blocksize (int | None): ukuran per file dalam satuan megabyte
          (optional is None, 64, 128, 256, 512)
          (default is 256)
      modified_date (bool): menambahkan modified_date pada tabel hasil
          (default is True)
      schema_temp (str): target schema untuk menyimpan temp table di hdfs 
          (*temp table akan terhapus otomatis setelah proses selesai)
          (default is "sharingvision")
      count (bool): menampilkan count data tersimpan
          (default is False)
    
  Returns:
      -
  """
  allowed_values_mode = {"append", "overwrite", "error", "ignore"}

  if mode not in allowed_values_mode:
      raise ValueError(
          f"mode must be one of {sorted(v for v in allowed_values_mode if v is not None)}"
      )
      
  allowed_values_blocksize = {64, 128, 256, 512, None}

  if blocksize not in allowed_values_blocksize:
      raise ValueError(
          f"blocksize must be one of {sorted(v for v in allowed_values_blocksize if v is not None)} or None"
      )
        
  if not isinstance(modified_date, bool):
    raise ValueError("Invalid value for 'modified_date'. Must be either True or False.")
  
  if not isinstance(count, bool):
    raise ValueError("Invalid value for 'count'. Must be either True or False.")

  if schema.startswith("dev_"):
    hdfs_path = "/dev/"
  else:
    hdfs_path = "/user/hive/warehouse/"
    
  if modified_date:         
    df = df.withColumn("modified_date", F.lit(datetime.now(timezone("Asia/Jakarta"))))

  if count:
    df.cache()

  print(f"Write table to Hive: {schema}.{table}")
    
  # pengecekan partitioned table
  if partition:
    
    # kondisi ketika tabel berpartisi
    partition = [col.strip() for col in partition.split("|")]
    list_partition = df.select(partition).distinct().collect()
    
    unit = len(str(len(list_partition)))
    
    first = True
    row_total = 0
    
    # loop untuk penulisan data per partisi sesuai standard blocksize
    for i, row in enumerate(list_partition):
      
      write_mode = mode if first else "append"

      condition = ' AND '.join([f"{col}='{row[col]}'" for col in partition])
      df_write = df.filter(condition)

      repartition_num = repartition_number

      if repartition_num is None and blocksize is not None:
        tag = '_'.join([row[col].strip() for col in partition]).lower()

        spark.sql(f"DROP TABLE IF EXISTS {schema_temp}.{table}_{tag}_allow_deleted")

        # penulisan sample data untuk estimasi ukuran data
        df_write\
          .sample(False,0.1,None)\
          .write\
          .format("parquet")\
          .mode("overwrite")\
          .saveAsTable(f"{schema_temp}.{table}_{tag}_allow_deleted")

        # mendapatkan metadata content summary dari sample table
        path = lambda p: spark._jvm.org.apache.hadoop.fs.Path(p)
        fs = spark._jvm.org.apache.hadoop.fs.FileSystem.get(spark._jsc.hadoopConfiguration())
        hdfs_folder = fs.getContentSummary(path(hdfs_path+f"{schema_temp}.db/{table}_{tag}_allow_deleted/"))

        # kalkulasi jumlah file yang optimal sesuai standard blocksize
        blocksize_number = float(hdfs_folder.getLength()*10)/float(blocksize*1024*1024)
        repartition_num = int(1 if blocksize_number < 1 else blocksize_number)

        if repartition_num:
          df_write = df_write.repartition(repartition_num)
        
        spark.sql(f"DROP TABLE {schema_temp}.{table}_{tag}_allow_deleted")
        
      elif repartition_num:
        # set jumlah file sesuai hasil kalkulasi
        df_write = df_write.repartition(repartition_num)
    
      # penghapusan data partisi jika sudah ada di target tabel
      if spark.sql(f"SHOW TABLES IN {schema} LIKE '{table}'").count() != 0:
        query_drop_partition = f"ALTER TABLE {schema}.{table} DROP IF EXISTS PARTITION({condition.replace(' AND ',',')})"
        spark.sql(query_drop_partition) # menghapus partisi di spark

      # penulisan pyspark dataframe ke hdfs
      df_write\
        .write\
        .format("parquet")\
        .partitionBy(partition)\
        .mode(write_mode)\
        .saveAsTable(f"{schema}.{table}")

      spark.sql(f"ANALYZE TABLE {schema}.{table} PARTITION ({condition.replace(' AND ',',')}) COMPUTE STATISTICS") # kalkulasi data
      spark.sql(f"REFRESH TABLE {schema}.{table}") # refresh tabel di spark
      

      first = False
        
      if count:
        row_count = df_write.count()
        row_total += row_count
        print(f"{str(i+1).ljust(unit, ' ')}) Count Record ({condition.replace(' AND ',',')}) = {row_count}")

    spark.sql(f"MSCK REPAIR TABLE {schema}.{table}") # memperbaiki tabel
    spark.sql(f"REFRESH TABLE {schema}.{table}") # refresh tabel di spark

  else:
    
    repartition_num = repartition_number

    if not repartition_num and blocksize:
      # kondisi ketika tabel berpartisi    
      spark.sql(f"DROP TABLE IF EXISTS {schema_temp}.{table}_allow_deleted")

      # penulisan sample data untuk estimasi ukuran data
      df\
        .sample(False,0.1,None)\
        .write\
        .format("parquet")\
        .mode("overwrite")\
        .saveAsTable(f"{schema_temp}.{table}_allow_deleted")

      # mendapatkan metadata content summary dari sample table
      path = lambda p: spark._jvm.org.apache.hadoop.fs.Path(p)
      fs = spark._jvm.org.apache.hadoop.fs.FileSystem.get(spark._jsc.hadoopConfiguration())
      hdfs_folder = fs.getContentSummary(path(hdfs_path+f"{schema_temp}.db/{table}_allow_deleted/"))

      # kalkulasi jumlah file yang optimal sesuai standard blocksize
      blocksize_number = float(hdfs_folder.getLength()*10)/float(blocksize*1024*1024)
      repartition_num = int(1 if blocksize_number < 1 else blocksize_number)

      if repartition_num:
        df = df.repartition(repartition_num)
        
      spark.sql(f"DROP TABLE {schema_temp}.{table}_allow_deleted")

    elif repartition_num:
      # set jumlah file sesuai hasil kalkulasi
      df = df.repartition(repartition_num)


    # penulisan pyspark dataframe ke hdfs
    df\
      .write\
      .format("parquet")\
      .mode(mode)\
      .saveAsTable(f"{schema}.{table}")

    spark.sql(f"ANALYZE TABLE {schema}.{table} COMPUTE STATISTICS") # kalkulasi data
    spark.sql(f"REFRESH TABLE {schema}.{table}") # refresh tabel di spark
    
    
  if count:
    total_record = df.count()
    print(f"\n====>> Grand Total Record = {total_record}")
    return total_record

  return None

## ===================================================================================================
  
import os
import re
import subprocess

def download_hdfs_file(file_name: str, file_type: str, delete_hdfs_file: bool = False):
  """Mengunduh semua file partisi HDFS dari direktori user ke lokal dan melakukan pembersihan opsional

  Args:
      file_name (str): nama file atau direktori target di HDFS
      delete_on_hdfs (bool): menentukan apakah direktori di HDFS akan dihapus setelah proses unduh selesai
          (default is False)

  Returns:
      list[str]: daftar path file lokal yang berhasil diunduh. Mengembalikan list kosong jika gagal atau file tidak ditemukan.
  """
  # Strip known extensions to prevent double extensions (e.g. report.csv.csv)
  known_extensions = {".csv", ".parquet", ".json", ".xlsx"}
  _, ext_found = os.path.splitext(file_name)
  if ext_found.lower() in known_extensions:
      file_name = file_name[: -len(ext_found)]

  def run_cmd(cmd: str):
      proc = subprocess.run(cmd, shell=True, capture_output=True, text=True)
      return proc.returncode == 0, proc.stdout, proc.stderr

  hadoop_user = os.getenv("HADOOP_USER_NAME")
  if not hadoop_user:
      raise RuntimeError("HADOOP_USER_NAME environment variable is not set.")

  # Local directory based on file_type
  local_dir = os.path.join("resources", file_type)
  os.makedirs(local_dir, exist_ok=True)

  # HDFS base path (for csv/parquet/json this is a directory, for excel this is the file)
  hdfs_dir = f"/user/{hadoop_user}/{file_name}"

  # 1) List HDFS dir/file
  ok, out, err = run_cmd(f'hadoop fs -ls "{hdfs_dir}"')
  if not ok:
      raise RuntimeError(f"Error listing HDFS path {hdfs_dir}:\n{err}")

  # 2) Determine which HDFS files to download
  part_files: List[str] = []
    
  if file_type == "excel":
      # Look for .xlsx files in the ls output
      for line in out.strip().splitlines():
          parts = line.split()
          if len(parts) < 8:
              continue
          path = parts[-1]
          if path.endswith(".xlsx"):
              part_files.append(path)

      # Fallback: if user passed the full .xlsx path as file_name,
      # hdfs_dir itself may be the file.
      if not part_files and file_name.endswith(".xlsx"):
          part_files = [hdfs_dir]
  else:
      # For Spark outputs: collect all part-* files under the directory
      for line in out.strip().splitlines():
          m = re.search(r"(/.*part-[^\s]*)", line)
          if m:
              part_files.append(m.group(1))

  if not part_files:
      print(f"No files found to download from HDFS path: {hdfs_dir}")

  part_files.sort()

  # 3) Decide local file extension based on file_type
  ext_map = {
      "csv": ".csv",
      "parquet": ".parquet",
      "json": ".json",
      "excel": ".xlsx",
  }
  ext = ext_map.get(file_type, f".{file_type}")

  # 4) Download and rename
  local_files: List[str] = []
  try:
      if len(part_files) == 1:
          # Single file → use file_name.ext
          target_local = os.path.join(local_dir, f"{file_name}{ext}")
          targets = [(part_files[0], target_local)]
      else:
          # Multiple parts → file_name_part_i.ext
          targets = [
              (hdfs_file, os.path.join(local_dir, f"{file_name}_part_{i}{ext}"))
              for i, hdfs_file in enumerate(part_files)
          ]

      for hdfs_file, local_file in targets:
          # Remove existing local file if any
          if os.path.exists(local_file):
              os.remove(local_file)

          ok, _, err = run_cmd(f'hadoop fs -get "{hdfs_file}" "{local_file}"')
          if not ok:
              raise RuntimeError(f"Failed to copy {hdfs_file}:\n{err}")

          local_files.append(local_file)

      # 5) Optional cleanup on HDFS
      if delete_hdfs_file:
          print(f"Cleaning up HDFS path: {hdfs_dir}")
          ok, _, err = run_cmd(f'hadoop fs -rm -r -skipTrash "{hdfs_dir}"')
          if not ok:
              raise RuntimeError(
                  f"Downloaded to resources/{file_type}: OK, "
                  f"but failed to delete HDFS path:\n{err}"
              )
      else:
          print("Cleanup skipped: delete_hdfs_file is set to False.")

      print(f"Downloaded to resources/{file_type}: OK")

  except Exception as e:
      print(f"Error: {e}")
      print("Operation aborted. HDFS path NOT deleted.")

## ===================================================================================================

import openpyxl
import subprocess
import os
    
def load_to_file(
    df,
    file_name: str,
    file_type: str = "csv",
    repartition_number: int = 1,
    delimiter: str = ",",
    download: bool = True,
    delete_hdfs_file: bool = True):
  """Melakukan penulisan pyspark dataframe ke HDFS dengan file format tertentu
  
  Args:
      df (pyspark.sql.dataframe.DataFrame): pyspark dataframe yang akan di tulis
      file_name (str): HDFS path / file name tanpa menambahkan ekstensi file
      file_type (str): tipe file yang akan di tulis
          (optional is csv, parquet, json, excel)
          (default is csv)
      repartition_number (int): jumlah file hasil
          (default is 1)      
      delimiter (str): pemisah antar kolom
          (default is ',')
      download (bool): download ke local CML
          (default is True)
      delete_hdfs_file (bool): delete file di hdfs setelah di download ke local CML
          (default is True)
          
  Returns:
      -
  """
  if " " in file_name:
      raise ValueError(f"Invalid file_name: contains space. '{file_name}'")

  # Strip known extensions to prevent double extensions (e.g. report.csv.csv)
  known_extensions = {".csv", ".parquet", ".json", ".xlsx"}
  _, ext_found = os.path.splitext(file_name)
  if ext_found.lower() in known_extensions:
      file_name = file_name[: -len(ext_found)]
      print(f"Warning: extension '{ext_found}' removed from file_name. Using '{file_name}'")

  file_type = file_type.lower()
  allowed_types = {"csv", "parquet", "json", "excel"}

  if file_type not in allowed_types:
      raise ValueError(f"file_type must be one of {sorted(allowed_types)}")

  print(f"Generate {file_type.upper()} file in HDFS...")

  # =======================
  # CSV
  # =======================
  if file_type == "csv":
      (
          df.repartition(repartition_number)
            .write
            .option("delimiter", delimiter)
            .option("header", True)
            .option("nullValue", "NULL")
            .option("emptyValue", "NULL")
            .mode("overwrite")
            .csv(file_name)
      )

  # =======================
  # PARQUET
  # =======================
  elif file_type == "parquet":
      (
          df.repartition(repartition_number)
            .write
            .mode("overwrite")
            .parquet(file_name)
      )

  # =======================
  # JSON
  # =======================
  elif file_type == "json":
      (
          df.repartition(repartition_number)
            .write
            .mode("overwrite")
            .json(file_name)
      )

  # =======================
  # EXCEL (via Pandas)
  # =======================
  elif file_type == "excel":

      if repartition_number == 1:
          # Single file mode (original behavior)
          pdf = df.coalesce(1).toPandas()
          local_path = f"{file_name.split('/')[-1]}.xlsx"
          pdf.to_excel(local_path, index=False)

          try:
              hadoop_user = os.getenv("HADOOP_USER_NAME")
              hdfs_dir = f"/user/{hadoop_user}/{file_name}"
              subprocess.run(
                  ["hdfs", "dfs", "-mkdir", "-p", hdfs_dir],
                  check=True
              )
              subprocess.run(
                  ["hdfs", "dfs", "-put", "-f", local_path, file_name],
                  check=True
              )
          except subprocess.CalledProcessError as e:
              print(f"Error uploading to HDFS: {e}")
              raise
          else:
              if os.path.exists(local_path):
                  os.remove(local_path)

      else:
          # Multiple files mode
          weights = [1.0] * repartition_number
          splits = df.randomSplit(weights)

          try:
              hadoop_user = os.getenv("HADOOP_USER_NAME")
              hdfs_dir = f"/user/{hadoop_user}/{file_name}"
              subprocess.run(
                  ["hdfs", "dfs", "-mkdir", "-p", hdfs_dir],
                  check=True
              )

              for i, split_df in enumerate(splits):
                  pdf = split_df.toPandas()
                  local_path = f"{file_name.split('/')[-1]}_part_{i}.xlsx"
                  pdf.to_excel(local_path, index=False)

                  subprocess.run(
                      ["hdfs", "dfs", "-put", "-f", local_path,
                       f"{hdfs_dir}/{file_name.split('/')[-1]}_part_{i}.xlsx"],
                      check=True
                  )

                  if os.path.exists(local_path):
                      os.remove(local_path)

                  print(f"  Written part {i+1}/{repartition_number}")

          except subprocess.CalledProcessError as e:
              print(f"Error uploading to HDFS: {e}")
              raise

  print(f"{file_type.upper()} file successfully generated")

  if download:
      download_hdfs_file(file_name, file_type, delete_hdfs_file)
  
def load_to_excel(
    df,
    file_name: str,
    sheet_name=None,
    repartition_number: int = 1,
    download: bool = True,
    delete_hdfs_file: bool = True):
  """Melakukan penulisan pyspark dataframe ke HDFS dengan file format excel (multi-sheet)
  
  Args:
      df (pyspark.sql.dataframe.DataFrame | list[pyspark.sql.dataframe.DataFrame]): 
          pyspark dataframe atau list of pyspark dataframe yang akan di tulis.
          Jika berupa list, setiap dataframe akan ditulis sebagai sheet terpisah.
      file_name (str): HDFS path / file name tanpa menambahkan ekstensi file
      sheet_name (str | list[str] | None): nama sheet atau list nama sheet.
          - Jika None, sheet akan diberi nama "Sheet1", "Sheet2", dst.
          - Jika str (single), hanya berlaku untuk dataframe tunggal.
          - Jika list, jumlahnya harus sama dengan jumlah dataframe.
          (default is None)
      repartition_number (int): jumlah file excel hasil
          - Jika 1, menghasilkan satu file excel
          - Jika > 1, data pada setiap sheet akan di-split menjadi N file excel
          (default is 1)
      download (bool): download ke local CML
          (default is True)
      delete_hdfs_file (bool): delete file di hdfs setelah di download ke local CML
          (default is True)
          
  Returns:
      -
  """
  import pandas as pd

  if " " in file_name:
      raise ValueError(f"Invalid file_name: contains space. '{file_name}'")

  # Strip known extensions to prevent double extensions (e.g. report.xlsx.xlsx)
  known_extensions = {".csv", ".parquet", ".json", ".xlsx"}
  _, ext_found = os.path.splitext(file_name)
  if ext_found.lower() in known_extensions:
      file_name = file_name[: -len(ext_found)]
      print(f"Warning: extension '{ext_found}' removed from file_name. Using '{file_name}'")

  # Normalize df to list
  if not isinstance(df, list):
      df_list = [df]
  else:
      df_list = df

  if len(df_list) == 0:
      raise ValueError("df must contain at least one dataframe.")

  # Normalize sheet_name to list
  if sheet_name is None:
      sheet_names = [f"Sheet{i+1}" for i in range(len(df_list))]
  elif isinstance(sheet_name, str):
      if len(df_list) == 1:
          sheet_names = [sheet_name]
      else:
          raise ValueError(
              "sheet_name is a single string but multiple dataframes were provided. "
              "Please provide a list of sheet names."
          )
  elif isinstance(sheet_name, list):
      if len(sheet_name) != len(df_list):
          raise ValueError(
              f"Length of sheet_name ({len(sheet_name)}) must match "
              f"the number of dataframes ({len(df_list)})."
          )
      sheet_names = sheet_name
  else:
      raise ValueError("sheet_name must be a str, list of str, or None.")

  print(f"Generate EXCEL file in HDFS...")

  try:
      hadoop_user = os.getenv("HADOOP_USER_NAME")
      hdfs_dir = f"/user/{hadoop_user}/{file_name}"
      subprocess.run(
          ["hdfs", "dfs", "-mkdir", "-p", hdfs_dir],
          check=True
      )

      if repartition_number == 1:
          # Single file mode (original behavior)
          local_path = f"{file_name.split('/')[-1]}.xlsx"

          with pd.ExcelWriter(local_path, engine="openpyxl") as writer:
              for i, (spark_df, sname) in enumerate(zip(df_list, sheet_names)):
                  print(f"  Writing sheet [{i+1}/{len(df_list)}]: {sname}")
                  pdf = spark_df.coalesce(1).toPandas()
                  pdf.to_excel(writer, sheet_name=sname, index=False)

          subprocess.run(
              ["hdfs", "dfs", "-put", "-f", local_path, file_name],
              check=True
          )

          if os.path.exists(local_path):
              os.remove(local_path)

      else:
          # Multiple files mode: split each sheet's data across N files
          weights = [1.0] * repartition_number

          # Pre-split all DataFrames
          splits_per_sheet = []
          for spark_df in df_list:
              splits_per_sheet.append(spark_df.randomSplit(weights))

          # Write N files, each containing all sheets with their respective chunk
          for part_idx in range(repartition_number):
              local_path = f"{file_name.split('/')[-1]}_part_{part_idx}.xlsx"

              with pd.ExcelWriter(local_path, engine="openpyxl") as writer:
                  for sheet_idx, sname in enumerate(sheet_names):
                      pdf = splits_per_sheet[sheet_idx][part_idx].toPandas()
                      pdf.to_excel(writer, sheet_name=sname, index=False)

              subprocess.run(
                  ["hdfs", "dfs", "-put", "-f", local_path,
                   f"{hdfs_dir}/{file_name.split('/')[-1]}_part_{part_idx}.xlsx"],
                  check=True
              )

              if os.path.exists(local_path):
                  os.remove(local_path)

              print(f"  Written part {part_idx+1}/{repartition_number}")

  except subprocess.CalledProcessError as e:
      print(f"Error uploading to HDFS: {e}")
      raise

  print(f"EXCEL file successfully generated")

  if download:
      download_hdfs_file(file_name, "excel", delete_hdfs_file)

## POSTGRESQL ========================================================================================
  
import psycopg2

def postgresql_execute_command(database, ip, port, username, password, query):
  """Menjalankan query di postgresql
  
  Args:
      database (str): target database di posgresql
      ip (str): target ip di posgresql
      port (str): target port di posgresql
      username (str): username di posgresql
      password (str): password di posgresql
      query (str): query yang akan di eksekusi
    
  Returns:
      -
  """
  
  # make a new connection
  conn = psycopg2.connect(
    host=ip,
    port=port,
    database=database,
    user=username,
    password=password
  )

  cur = conn.cursor()  # creating a cursor

  # executing queries
  cur.execute(query)

  # commit the changes
  conn.commit() # This right here

  # close connection
  cur.close()
  conn.close()

def postgresql_execute_fetchall(database, ip, port, username, password, query):
  """Menjalankan query di postgresql
  
  Args:
      database (str): target database di posgresql
      ip (str): target ip di posgresql
      port (str): target port di posgresql
      username (str): username di posgresql
      password (str): password di posgresql
      query (str): query yang akan di eksekusi
    
  Returns:
      list: data keluaran dari query yang di eksekusi
  """
  
  # make a new connection
  conn = psycopg2.connect(
    host=ip,
    port=port,
    database=database,
    user=username,
    password=password
  )
  cur = conn.cursor()  # creating a cursor

  # executing queries
  cur.execute(query)

  # get response data
  result = cur.fetchall()

  # close connection
  cur.close()
  conn.close()

  return result
  
def postgresql_drop_table(database, schema, table, ip, port, username, password):
  """Melakukan query drop table di postgresql
  
  Args:
      database (str): target database di posgresql
      schema (str): target schema di posgresql yang akan di hapus
      table (str): target table di posgresql yang akan di hapus
      ip (str): target ip di posgresql
      port (str): target port di posgresql
      username (str): username di posgresql
      password (str): password di posgresql
    
  Returns:
      -
  """
  
  # pembentukan query drop table
  query = "DROP TABLE IF EXISTS {}.{}".format(schema, table)
  
  # eksekusi query drop table
  postgresql_execute_command(database, ip, port, username, password, query)
     
def postgresql_truncate_table(database, schema, table, ip, port, username, password):
  """Melakukan query truncate table di postgresql
  
  Args:
      database (str): target database di posgresql
      schema (str): target schema di posgresql yang akan di truncate
      table (str): target table di posgresql yang akan di truncate
      ip (str): target ip di posgresql
      port (str): target port di posgresql
      username (str): username di posgresql
      password (str): password di posgresql
    
  Returns:
      -
  """
  
  # query truncate table
  query = "TRUNCATE TABLE {}.{}".format(schema, table)
  
  # eksekusi query truncate table
  postgresql_execute_command(database, ip, port, username, password, query)
  
def load_to_postgresql(df, database, schema, table, ip, port, username, password, mode="append", strategy=None, ddl=None, count=True):
  """Melakukan penulisan pyspark dataframe sebagai tabel di postgresql
  
  Args:
      df (pyspark.sql.dataframe.DataFrame): pyspark dataframe yang akan di tulis sebagai tabel posgresql
      database (str): target database di posgresql
      schema (str): target schema di posgresql
      table (str): target table di posgresql
      ip (str): target ip di posgresql
      port (str): target port di posgresql
      username (str): username di posgresql
      password (str): password di posgresql
      mode (str): cara yang dilakukan untuk penulisan tabel
          (optional is append, overwrite, error, ignore)
          (default is append)
      strategy (str): proses yang dilakukan sebelum penulisan tabel
          (optional is None, truncate, drop)
          (default is None)
      ddl (str): query create untuk melakukan pembuatan struktur tabel
          (default is None)
      count (bool): menampilkan count data tersimpan
          (default is True)
    
  Returns:
      -
  """
  if not isinstance(count, bool):
    raise ValueError("Invalid value for 'count'. Must be either True or False.")
      
  print("Write table to PostgreSQL: {}.{}".format(schema, table))
    
  if strategy:
    
    if strategy.lower() == "truncate":
      # menjalankan query truncate table
      postgresql_truncate_table(database, schema, table, ip, port, username, password)
    elif strategy.lower() == "drop":
      # menjalankan query drop table
      postgresql_drop_table(database, schema, table, ip, port, username, password)
      
  if ddl:
    # menjalankan query dll
    postgresql_execute_command(database, ip, port, username, password, ddl)
    
  # penulisan pyspark dataframe ke postgresql
  df\
    .write\
    .mode(mode)\
    .format("jdbc")\
    .option("driver", "org.postgresql.Driver")\
    .option("url", "jdbc:postgresql://{}:{}/{}".format(ip, port, database))\
    .option("user", username)\
    .option("password", password)\
    .option("dbtable", ".".join([schema, table]))\
    .save()

  if count:
    print("\n====>> Total Record =", df.count())
      
## MYSQL =============================================================================================

import mysql.connector as ms

def mysql_execute_command(schema, ip, port, username, password, query):
  """Menjalankan query di mysql
  
  Args:
      schema (str): target schema di mysql
      ip (str): target ip di mysql
      port (str): target port di mysql
      username (str): username di mysql
      password (str): password di mysql
      query (str): query yang akan di eksekusi
    
  Returns:
      -
  """
  # make a new connection  
  conn = ms.connect(
    host=ip,
    port=port,
    database=schema,
    user=username,
    password=password
  )

  cur = conn.cursor() # creating a cursor

  # executing queries
  cur.execute(query)

  # commit the changes
  conn.commit() # This right here

  # close connection
  cur.close()
  conn.close()
  
def mysql_execute_fetchall(schema, ip, port, username, password, query):
  """Menjalankan query di mysql
  
  Args:
      schema (str): target schema di mysql
      ip (str): target ip di mysql
      port (str): target port di mysql
      username (str): username di mysql
      password (str): password di mysql
      query (str): query yang akan di eksekusi
    
  Returns:
      list: data keluaran dari query yang di eksekusi
  """
  
  # make a new connection
  conn = ms.connect(
    host=ip,
    port=port,
    database=schema,
    user=username,
    password=password
  )

  cur = conn.cursor() # creating a cursor

  # executing queries
  cur.execute(query)

  # get response data
  result = cur.fetchall()

  # close connection
  cur.close()
  conn.close()

  return result
  
def mysql_drop_table(schema, table, ip, port, username, password):
  """Melakukan query drop table di mysql
  
  Args:
      schema (str): target schema di mysql yang akan di hapus
      table (str): target table di mysql yang akan di hapus
      ip (str): target ip di mysql
      port (str): target port di mysql
      username (str): username di mysql
      password (str): password di mysql
    
  Returns:
      -
  """
  
  # pembentukan query drop table  
  query = "DROP TABLE IF EXISTS {}.{}".format(schema, table)
  
  # eksekusi query drop table
  mysql_execute_command(schema, ip, port, username, password, query)
        
def mysql_truncate_table(schema, table, ip, port, username, password):
  """Melakukan query truncate table di mysql
  
  Args:
      schema (str): target schema di mysql yang akan di truncate
      table (str): target table di mysql yang akan di truncate
      ip (str): target ip di mysql
      port (str): target port di mysql
      username (str): username di mysql
      password (str): password di mysql
    
  Returns:
      -
  """
  
  # pembentukan query truncate table
  query = "TRUNCATE TABLE {}.{}".format(schema, table)
  
  # eksekusi query truncate table
  mysql_execute_command(schema, ip, port, username, password, query)
  
def load_to_mysql(df, schema, table, ip, port, username, password, mode="append", strategy=None, ddl=None, count=True):
  """Melakukan penulisan pyspark dataframe sebagai tabel di mysql
  
  Args:
      df (pyspark.sql.dataframe.DataFrame): pyspark dataframe yang akan di tulis sebagai tabel mysql
      schema (str): target schema di mysql
      table (str): target table di mysql
      ip (str): target ip di mysql
      port (str): target port di mysql
      username (str): username di mysql
      password (str): password di mysql
      mode (str): cara yang dilakukan untuk penulisan tabel
          (optional is append, overwrite, error, ignore)
          (default is append)
      strategy (str): proses yang dilakukan sebelum penulisan tabel
          (optional is None, truncate, drop)
          (default is None)
      ddl (str): query create untuk melakukan pembuatan struktur tabel
          (default is None)
      count (bool): menampilkan count data tersimpan
          (default is True)
    
  Returns:
      -
  """
  if not isinstance(count, bool):
    raise ValueError("Invalid value for 'count'. Must be either True or False.")

  print("Write table to MySQL: {}.{}".format(schema, table))
      
  if strategy:
    
    if strategy.lower() == "truncate":
      # menjalankan query truncate table
      mysql_truncate_table(schema, table, ip, port, username, password)
    elif strategy.lower() == "drop":
      # menjalankan query drop table
      mysql_drop_table(schema, table, ip, port, username, password)
      
  if ddl:
    # menjalankan query dll
    mysql_execute_command(schema, ip, port, username, password, ddl)
    
  # penulisan pyspark dataframe ke mysql    
  df\
    .write\
    .mode(mode)\
    .format("jdbc")\
    .option("driver", "com.mysql.cj.jdbc.Driver")\
    .option("url","jdbc:mysql://"+ip+":"+port+"/"+schema+"?user="+username+"&password="+password+"&characterEncoding=UTF-8&connectionCollation=utf8mb4_unicode_ci&useUnicode=yes&serverTimezone=Asia%2FJakarta&sessionVariables=character_set_server%3Dutf8mb4,character_set_client%3Dutf8mb4,character_set_connection%3Dutf8mb4,character_set_results%3Dutf8mb4,collation_connection%3Dutf8mb4_unicode_ci")\
    .option("useSSL", "false")\
    .option("dbtable", ".".join([schema, table]))\
    .save()

  if count:
    print("\n====>> Total Record =", df.count())

  