Spaces:
Runtime error
Runtime error
File size: 6,262 Bytes
c2d582c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 |
import pandas as pd
from sqlalchemy import create_engine, exc, engine
from snowflake.sqlalchemy import URL
import streamlit as st
def connect_to_snowflake(
username: str,
password: str,
account: str,
warehouse: str,
database: str,
schema: str,
) -> engine:
"""
Connect to Snowflake using the specified credentials.
Parameters:
- username (str): Snowflake username
- password (str): Snowflake password
- account (str): Snowflake account name
- warehouse (str): Snowflake warehouse name
- database (str): Snowflake database name
- schema (str): Snowflake schema name
Returns:
- Engine: SQLAlchemy Engine object for the connection
"""
try:
conn = create_engine(
URL(
user=username,
password=password,
account=account,
warehouse=warehouse,
database=database,
schema=schema,
)
)
return conn
except exc.SQLAlchemyError as err:
st.error(f"Error connecting to Snowflake: {err}")
return None
def load_data_to_snowflake(df: pd.DataFrame, conn: engine, table: str) -> None:
"""
Load data from a CSV file into a table in Snowflake.
Parameters:
- filepath (str): Path to the CSV file
- engine (Engine): SQLAlchemy Engine object for the connection
- table (str): Snowflake table name
Returns:
- None
"""
try:
# Load data to Snowflake
df.to_sql(table, conn, if_exists="replace", index=False)
st.success("Data loaded to Snowflake successfully")
st.snow()
except Exception as err:
print(f"Error loading data to Snowflake: {err}")
def connect_to_postgres(
username: str, password: str, host: str, port: str, database: str
) -> engine:
"""
Connect to PostgreSQL using the specified credentials.
Parameters:
- username (str): PostgreSQL username
- password (str): PostgreSQL password
- host (str): PostgreSQL host name
- port (str): PostgreSQL port
- database (str): PostgreSQL database name
Returns:
- Engine: SQLAlchemy Engine object for the connection
"""
try:
conn = create_engine(
f"postgresql://{username}:{password}@{host}:{port}/{database}"
)
return conn
except exc.SQLAlchemyError as err:
st.error(f"Error connecting to PostgreSQL: {err}")
return None
def load_data_to_postgres(df: pd.DataFrame, conn: engine, table: str) -> None:
"""
Load data from a CSV file into a table in PostgreSQL.
Parameters:
- df (pd.DataFrame): DataFrame containing the data to load
- conn (engine): SQLAlchemy Engine object for the connection
- table (str): PostgreSQL table name
Returns:
- None
"""
try:
# Load data to PostgreSQL
df.to_sql(table, conn, if_exists="replace", index=False)
st.success("Data loaded to PostgreSQL successfully")
st.balloons()
except Exception as err:
st.error(f"Error loading data to PostgreSQL: {err}")
def main():
st.title("Load Data to Databases")
# Data to load to database(s)
df = pd.read_csv("philox-testset-1.csv")
# Get user input for data storage option
storage_option = st.selectbox(
"Select data storage option:", ["Snowflake", "PostgreSQL"]
)
@st.cache(allow_output_mutation=True)
def reset_form_fields():
user = ""
password = ""
account = ""
warehouse = ""
database = ""
schema = ""
table = ""
host = ""
port = ""
if storage_option == "Snowflake":
st.subheader("Enter Snowflake Credentials")
# Get user input for Snowflake credentials
user = st.text_input("Username:", value="TONY")
password = st.text_input("Password:", type="password")
account = st.text_input("Account:", value="jn27194.us-east4.gcp")
warehouse = st.text_input("Warehouse:", value="NAH")
database = st.text_input("Database:", value="SNOWVATION")
schema = st.text_input("Schema:", value="PUBLIC")
table = st.text_input("Table:")
# Load the data to Snowflake
if st.button("Load data to Snowflake"):
if (
user
and password
and account
and warehouse
and database
and schema
and table
):
conn = connect_to_snowflake(
username=user,
password=password,
account=account,
warehouse=warehouse,
database=database,
schema=schema,
)
if conn:
load_data_to_snowflake(df, conn, table)
else:
st.warning("Please enter all Snowflake credentials")
elif storage_option == "PostgreSQL":
st.subheader("Enter PostgreSQL Credentials")
# Get user input for PostgreSQL credentials
user = st.text_input("Username:", value="postgres")
password = st.text_input("Password:", type="password")
host = st.selectbox("Host:", ["localhost", "other"])
if host == "other":
host = st.text_input("Enter host:")
port = st.text_input("Port:", value="5432")
database = st.text_input("Database:", value="snowvation")
table = st.text_input("Table:")
# Load the data to PostgreSQL
if st.button("Load data to PostgreSQL"):
if user and password and host and port and database and table:
conn = connect_to_postgres(
username=user,
password=password,
host=host,
port=port,
database=database,
)
if conn:
load_data_to_postgres(df, conn, table)
else:
st.warning("Please enter all PostgreSQL credentials and table name")
# Reset form fields when storage_option changes
reset_form_fields()
if __name__ == "__main__":
main() |