-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathmain.py
More file actions
40 lines (33 loc) · 1.11 KB
/
main.py
File metadata and controls
40 lines (33 loc) · 1.11 KB
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
import pandas as pd
import psycopg2
import yaml
def read_csv(file_name):
return pd.read_csv(file_name)
def upload_to_postgres(df, config):
connection = psycopg2.connect(
host='localhost',
port=config['port'],
user=config['username'],
password=config['password'],
dbname=config['database']
)
cursor = connection.cursor()
# Create table if not exists
cols = ",".join([f"{col} TEXT" for col in df.columns])
create_table_query = f"CREATE TABLE IF NOT EXISTS data ({cols});"
cursor.execute(create_table_query)
connection.commit()
# Insert data
for _, row in df.iterrows():
columns = ','.join(row.keys())
values = ','.join([f"'{val}'" for val in row])
insert_query = f"INSERT INTO data ({columns}) VALUES ({values})"
cursor.execute(insert_query)
connection.commit()
cursor.close()
connection.close()
if __name__ == "__main__":
with open("config.yaml", 'r') as stream:
config = yaml.safe_load(stream)
df = read_csv("data.csv")
upload_to_postgres(df, config['postgresql'])