pikto/GPT3-Dataset-Generator

1import pandas as pd2from sqlalchemy import create_engine, exc, engine3from snowflake.sqlalchemy import URL4import streamlit as st5 6 7def connect_to_snowflake(8    username: str,9    password: str,10    account: str,11    warehouse: str,12    database: str,13    schema: str,14) -> engine:15    """16    Connect to Snowflake using the specified credentials.17    Parameters:18    - username (str): Snowflake username19    - password (str): Snowflake password20    - account (str): Snowflake account name21    - warehouse (str): Snowflake warehouse name22    - database (str): Snowflake database name23    - schema (str): Snowflake schema name24    Returns:25    - Engine: SQLAlchemy Engine object for the connection26    """27 28    try:29        conn = create_engine(30            URL(31                user=username,32                password=password,33                account=account,34                warehouse=warehouse,35                database=database,36                schema=schema,37            )38        )39        return conn40    except exc.SQLAlchemyError as err:41        st.error(f"Error connecting to Snowflake: {err}")42        return None43 44 45def load_data_to_snowflake(df: pd.DataFrame, conn: engine, table: str) -> None:46    """47    Load data from a CSV file into a table in Snowflake.48    Parameters:49    - filepath (str): Path to the CSV file50    - engine (Engine): SQLAlchemy Engine object for the connection51    - table (str): Snowflake table name52    Returns:53    - None54    """55    try:56        # Load data to Snowflake57        df.to_sql(table, conn, if_exists="replace", index=False)58        st.success("Data loaded to Snowflake successfully")59        st.snow()60    except Exception as err:61        print(f"Error loading data to Snowflake: {err}")62 63 64def connect_to_postgres(65    username: str, password: str, host: str, port: str, database: str66) -> engine:67    """68    Connect to PostgreSQL using the specified credentials.69    Parameters:70    - username (str): PostgreSQL username71    - password (str): PostgreSQL password72    - host (str): PostgreSQL host name73    - port (str): PostgreSQL port74    - database (str): PostgreSQL database name75    Returns:76    - Engine: SQLAlchemy Engine object for the connection77    """78    try:79        conn = create_engine(80            f"postgresql://{username}:{password}@{host}:{port}/{database}"81        )82        return conn83    except exc.SQLAlchemyError as err:84        st.error(f"Error connecting to PostgreSQL: {err}")85        return None86 87 88def load_data_to_postgres(df: pd.DataFrame, conn: engine, table: str) -> None:89    """90    Load data from a CSV file into a table in PostgreSQL.91    Parameters:92    - df (pd.DataFrame): DataFrame containing the data to load93    - conn (engine): SQLAlchemy Engine object for the connection94    - table (str): PostgreSQL table name95    Returns:96    - None97    """98    try:99        # Load data to PostgreSQL100        df.to_sql(table, conn, if_exists="replace", index=False)101        st.success("Data loaded to PostgreSQL successfully")102        st.balloons()103    except Exception as err:104        st.error(f"Error loading data to PostgreSQL: {err}")105 106 107def main():108    st.title("Load Data to Databases")109 110    # Data to load to database(s)111    df = pd.read_csv("philox-testset-1.csv")112 113    # Get user input for data storage option114    storage_option = st.selectbox(115        "Select data storage option:", ["Snowflake", "PostgreSQL"]116    )117 118    @st.cache(allow_output_mutation=True)119    def reset_form_fields():120        user = ""121        password = ""122        account = ""123        warehouse = ""124        database = ""125        schema = ""126        table = ""127        host = ""128        port = ""129 130    if storage_option == "Snowflake":131        st.subheader("Enter Snowflake Credentials")132        # Get user input for Snowflake credentials133        user = st.text_input("Username:", value="TONY")134        password = st.text_input("Password:", type="password")135        account = st.text_input("Account:", value="jn27194.us-east4.gcp")136        warehouse = st.text_input("Warehouse:", value="NAH")137        database = st.text_input("Database:", value="SNOWVATION")138        schema = st.text_input("Schema:", value="PUBLIC")139        table = st.text_input("Table:")140 141        # Load the data to Snowflake142        if st.button("Load data to Snowflake"):143            if (144                user145                and password146                and account147                and warehouse148                and database149                and schema150                and table151            ):152                conn = connect_to_snowflake(153                    username=user,154                    password=password,155                    account=account,156                    warehouse=warehouse,157                    database=database,158                    schema=schema,159                )160                if conn:161                    load_data_to_snowflake(df, conn, table)162            else:163                st.warning("Please enter all Snowflake credentials")164 165    elif storage_option == "PostgreSQL":166        st.subheader("Enter PostgreSQL Credentials")167        # Get user input for PostgreSQL credentials168        user = st.text_input("Username:", value="postgres")169        password = st.text_input("Password:", type="password")170        host = st.selectbox("Host:", ["localhost", "other"])171        if host == "other":172            host = st.text_input("Enter host:")173        port = st.text_input("Port:", value="5432")174        database = st.text_input("Database:", value="snowvation")175        table = st.text_input("Table:")176 177        # Load the data to PostgreSQL178        if st.button("Load data to PostgreSQL"):179            if user and password and host and port and database and table:180                conn = connect_to_postgres(181                    username=user,182                    password=password,183                    host=host,184                    port=port,185                    database=database,186                )187                if conn:188                    load_data_to_postgres(df, conn, table)189            else:190                st.warning("Please enter all PostgreSQL credentials and table name")191 192    # Reset form fields when storage_option changes193    reset_form_fields()194 195 196if __name__ == "__main__":197    main()