-
Notifications
You must be signed in to change notification settings - Fork 0
/
db_manager.py
99 lines (91 loc) · 3.3 KB
/
db_manager.py
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
import os
from dotenv import load_dotenv
from sqlalchemy import text, DDL, event
import sqlalchemy
from typing import List
load_dotenv('.env')
POSTGRES_PASSWORD = os.environ['POSTGRES_PASSWORD']
CONNECTION_STRING = f"postgresql://postgres:{POSTGRES_PASSWORD}@localhost:5432/ShareMySGHouse"
NUM_OF_TABLES = 5
engine = sqlalchemy.create_engine(
CONNECTION_STRING
)
db = engine.connect()
def drop_all_tables(db: sqlalchemy.engine.Connection) -> None:
db.execute(text(
'''
DROP TABLE IF EXISTS credit_cards;
DROP TABLE IF EXISTS rental;
DROP TABLE IF EXISTS house_ratings;
DROP TABLE IF EXISTS houses;
DROP TABLE IF EXISTS users;
'''
))
def create_all_tables(db: sqlalchemy.engine.Connection) -> None:
db.execute(text(
'''CREATE TABLE IF NOT EXISTS users(
email VARCHAR(50) PRIMARY KEY,
fname VARCHAR(32) NOT NULL,
lname VARCHAR(32) NOT NULL,
age INTEGER NOT NULL,
password VARCHAR(32))'''
))
db.execute(text(
'''CREATE TABLE IF NOT EXISTS houses(
id INTEGER PRIMARY KEY,
location VARCHAR(32) NOT NULL,
price FLOAT NOT NULL,
num_room INT NOT NULL,
owner_email VARCHAR(50),
FOREIGN KEY (owner_email) REFERENCES users(email))'''
))
db.execute(text(
'''CREATE TABLE IF NOT EXISTS credit_cards(
type VARCHAR(30) NOT NULL,
number VARCHAR(32) PRIMARY KEY,
email VARCHAR(50) NOT NULL,
FOREIGN KEY (email) REFERENCES users(email))'''
))
db.execute(text(
'''CREATE TABLE IF NOT EXISTS rental(
email VARCHAR(50),
houseid INTEGER,
num_of_days INTEGER NOT NULL,
date DATE NOT NULL,
FOREIGN KEY (email) REFERENCES users(email),
FOREIGN KEY (houseid) REFERENCES houses(id),
PRIMARY KEY(email, houseid))'''
))
db.execute(text(
'''SET datestyle = dmy
'''
))
db.commit()
print("ALL TABLES CREATED SUCCESSFULLY")
def load_dummy_data(db: sqlalchemy.engine.Connection) -> None:
db.execute(text("".join(open("data/users.sql", "r").readlines())))
db.execute(text("".join(open("data/houses.sql", "r").readlines())))
db.execute(text("".join(open("data/credit_cards.sql", "r").readlines())))
db.execute(text("".join(open("data/rental.sql", "r").readlines())))
db.commit()
print("ALL DATA POPULATED")
def check_table_exist(db: sqlalchemy.engine.Connection) -> bool:
res = db.execute(text(
'''SELECT count(*) FROM INFORMATION_SCHEMA.TABLES WHERE TABLE_SCHEMA = 'public'
'''
))
return res.scalar() == NUM_OF_TABLES
def setup_database(db: sqlalchemy.engine.Connection) -> None:
# if (check_table_exist(db)):
# return
# if tables do not exist
# create them
drop_all_tables(db)
create_all_tables(db)
load_dummy_data(db)
def execute_sql(db: sqlalchemy.engine.Connection, command: str) -> List[sqlalchemy.engine.row.Row]:
return db.execute(text(command)).fetchall()
def execute_update(db: sqlalchemy.engine.Connection, command: str):
db.execute(text(command))
return
# print(type(execute_sql(db,"SELECT * FROM INFORMATION_SCHEMA.TABLES WHERE TABLE_SCHEMA = 'public'")[0]))