-
Notifications
You must be signed in to change notification settings - Fork 0
/
import.py
48 lines (37 loc) · 1.52 KB
/
import.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
import os
import psycopg2
def pretty_print(msg):
print("=" * len(msg))
print(msg)
print("=" * len(msg))
print()
# Connect to an existing database
conn = psycopg2.connect("dbname='311CI' user='postgres' host='localhost' password='postgres'")
# Open a cursor to perform database operations
cur = conn.cursor()
DIR = './out/'
input_files = [DIR + f for f in os.listdir(DIR) if f.endswith('.csv')]
# Add our SQL function + PostGIS extension
cur.execute(open("extras.sql", "r").read())
pretty_print('Added SQL function + PostGIS extension')
# We need to import the request table first
fname = DIR + 'request.csv'
with open(fname, 'rb') as f:
relation = fname.replace(DIR, '').replace('.csv', '')
cur.execute('TRUNCATE TABLE ' + relation + ' CASCADE;')
SQL = "COPY {} FROM '{}' CSV QUOTE '\"';".format(relation, os.path.abspath(fname))
cur.execute(SQL)
pretty_print('Imported ' + fname + ' to relation ' + relation)
# Copy every other csv to the respective PostgreSQL table
for fname in [x for x in input_files if x != DIR + 'request.csv']:
with open(fname, 'rb') as f:
relation = fname.replace(DIR, '').replace('.csv', '')
cur.execute('TRUNCATE TABLE ' + relation + ' CASCADE;')
SQL = "COPY {} FROM '{}' CSV QUOTE '\"';".format(relation, os.path.abspath(fname))
cur.execute(SQL)
pretty_print('Imported ' + fname + ' to relation ' + relation)
# Make the changes to the database persistent
conn.commit()
# Close communication with the database
cur.close()
conn.close()