-
Notifications
You must be signed in to change notification settings - Fork 25
Expand file tree
/
Copy pathhug-postgresql-example.py
More file actions
66 lines (60 loc) · 2.34 KB
/
hug-postgresql-example.py
File metadata and controls
66 lines (60 loc) · 2.34 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
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
"""Example connect and run write to postgresql."""
import hug
import psycopg2
"""
# don't forget to create postgres sql t database first:
# $ createdb t
# then connect to psql
# $ psql -h localhost -d test 2 ↵
# psql (9.5.3)
# Type "help" for help.
# then create test user
# t=# create user t with password 'test';
# CREATE ROLE
# t=# grant all privileges on database test to test;
# GRANT
# t=#
"""
@hug.get('/test')
def test_connect():
"""Test connection to db."""
# Get database credentials from configs.py or environment variables
try:
from configs import db_name, db_user, db_host, db_password
except ImportError:
import os
db_name = os.getenv('DB_NAME', 't')
db_user = os.getenv('DB_USER', 't')
db_host = os.getenv('DB_HOST', 'localhost')
db_password = os.getenv('DB_PASSWORD', '')
if not db_password:
return {'error': 'Database password not configured. Set DB_PASSWORD environment variable or configs.py'}
psycopg2.connect(f"dbname='{db_name}' user='{db_user}' host='{db_host}' password='{db_password}'")
return ('connected successfully to db! ready for queries.')
@hug.get('/checktable')
def test_write(user='t', table='testtable'):
"""Test write to DB."""
# Get database credentials from configs.py or environment variables
try:
from configs import db_name, db_user, db_host, db_password
except ImportError:
import os
db_name = os.getenv('DB_NAME', 't')
db_user = os.getenv('DB_USER', 't')
db_host = os.getenv('DB_HOST', 'localhost')
db_password = os.getenv('DB_PASSWORD', '')
if not db_password:
return {'error': 'Database password not configured. Set DB_PASSWORD environment variable or configs.py'}
conn = psycopg2.connect(f"dbname='{db_name}' user='{db_user}' host='{db_host}' password='{db_password}'")
print('connected successfully to db! ready for queries.')
cur = conn.cursor()
# Fix SQL injection vulnerability by using parameterized queries
cur.execute("SELECT exists(SELECT relname FROM pg_class WHERE relname=%s)", (table,))
exists = cur.fetchone()[0]
print(exists)
cur.close()
conn.close()
if exists:
return 'THIS TABLE EXISTS'
else:
return 'This Table does not exist'