main.py

import os
import pandas as pd
import json
from sqlalchemy import create_engine
from dotenv import load_dotenv

def extract_tabular_data(file_path: str):
    """Extract data from a tabular file_format, with pandas"""
    if file_path.endswith(".csv"):
        return pd.read_csv(file_path)
    elif file_path.endswith(".parquet"):
        return pd.read_parquet(file_path)
    else:
        print("Warning: Invalid file extension. Please try with .csv or .parquet!")
    
def extract_json_data(file_path):
    """Extract and flatten data from a JSON file."""
    if file_path[-5:len(file_path)] == ".json":
        df = pd.read_json(file_path)
        df_energy_source = pd.json_normalize(df["energySource"])
        df["energySource."+df_energy_source.columns] = df_energy_source
        df.drop(columns="energySource", inplace=True)
        return df
    
def transform_electricity_sales_data(raw_data: pd.DataFrame):
    """
    Transform electricity sales to find the total amount of electricity sold
    in the residential and transportation sectors.
    
    To transform the electricity sales data, you'll need to do the following:
    - Drop any records with NA values in the `price` column. Do this inplace.
    - Only keep records with a `sectorName` of "residential" or "transportation".
    - Create a `month` column using the first 4 characters of the values in `period`.
    - Create a `year` column using the last 2 characters of the values in `period`.
    - Return the transformed `DataFrame`, keeping only the columns `year`, `month`, `stateid`, `price` and `price-units`.
    """
    
    raw_data = raw_data[(raw_data["sectorName"]=="residential") | (raw_data["sectorName"]=="transportation")]
    
    period = raw_data["period"].str.split("-", n=1, expand=True)
    raw_data = raw_data.copy()
    
    raw_data.loc[:, "month"] = period[1]
    raw_data.loc[:, "year"] = period[0]
    raw_data.drop(columns=["period", "stateDescription", "sectorid", "sectorName"], inplace=True)

    return raw_data

def transform_electricity_cap_data(raw_data: pd.DataFrame):
    """

    """
    raw_data.columns = ["period", "stateId", "stateDescription", "energySourceid", "energySourceDescription", "energySourceCapability", "energySourceCapabilityUnits"]
    raw_data["energySourceCapability"] = raw_data["energySourceCapability"].astype(float)

    return raw_data

def load(dataframe: pd.DataFrame, table_name: str):
    """Load a DataFrame to MySQL Database"""

    # Create MySQL engine
    dbuser = os.getenv("USER")
    dbpassword = os.getenv("PASSWORD")
    dbname = "electricity_db"
    engine = create_engine(f"mysql+mysqlconnector://{dbuser}:{dbpassword}@localhost/{dbname}")

    dataframe.to_sql(name=table_name, con=engine, index=False, if_exists="replace")

raw_electricity_capability_df = extract_json_data("datasets/electricity_capability_nested.json")
raw_electricity_sales_df = extract_tabular_data("datasets/electricity_sales.csv")

cleaned_electricity_sales_df = transform_electricity_sales_data(raw_electricity_sales_df)
cleaned_electricity_cap_df = transform_electricity_cap_data(raw_electricity_capability_df)

load(cleaned_electricity_cap_df, "electricity_capability")
load(cleaned_electricity_sales_df, "electricity_sales")