]>
Commit | Line | Data |
---|---|---|
1 | #!/usr/bin/env python2 | |
2 | # PyPsd's Internets bot's zip code importer, for the weather commands | |
3 | # This is run separately to PyPsd itself, before if you want weather ZIP codes to work | |
4 | ||
5 | import os | |
6 | import sys | |
7 | import yaml | |
8 | import MySQLdb as db | |
9 | ||
10 | # find config file and zipcode csv file | |
11 | config_filename = 'pypsd.yml' | |
12 | if not os.path.exists(config_filename): | |
13 | config_filename = os.path.join('..', config_filename) | |
14 | if not os.path.exists(config_filename): | |
15 | print "We cannot locate the configuration file pypsd.yml" | |
16 | print 'This script should be run from the root acid/ directory' | |
17 | exit(1) | |
18 | ||
19 | zipcode_csv_file = 'weather-zipcodes.csv' | |
20 | if not os.path.exists(zipcode_csv_file): | |
21 | zipcode_csv_file = os.path.join('data', 'weather-zipcodes.csv') | |
22 | if not os.path.exists(zipcode_csv_file): | |
23 | zipcode_csv_file = os.path.join('pyva', zipcode_csv_file) | |
24 | if not os.path.exists(zipcode_csv_file): | |
25 | print "We cannot locate the zipcode csv file" | |
26 | print 'weather-zipcodes.csv should be in the working dir, in data/ , or in pyva/data/' | |
27 | exit(1) | |
28 | ||
29 | # options | |
30 | drop_table = '--drop' in sys.argv | |
31 | debug = '--debug' in sys.argv | |
32 | ||
33 | # config | |
34 | with open(config_filename, "r") as config_fd: | |
35 | config = yaml.load(config_fd) | |
36 | ||
37 | # make db pointer | |
38 | dbx = db.connect( | |
39 | host=config.get('database').get('host'), | |
40 | user=config.get('database').get('user'), | |
41 | passwd=config.get('database').get('passwd'), | |
42 | db=config.get('database').get('db'), | |
43 | ) | |
44 | dbx.autocommit(True) # no need to have transactions | |
45 | dbp = dbx.cursor() | |
46 | ||
47 | # check if table exists | |
48 | dbp.execute("SELECT count(*) FROM information_schema.tables t WHERE t.table_name='zipcode_citystate';") | |
49 | if dbp.fetchone()[0] == 0: | |
50 | table_exists = False | |
51 | else: | |
52 | table_exists = True | |
53 | rows = 0 # zero rows since table doesn't exist | |
54 | ||
55 | # clearing/setting table information | |
56 | try: | |
57 | # check if we should drop the table, to basically 'refresh'and start again | |
58 | if drop_table: | |
59 | if table_exists: | |
60 | dbp.execute("DROP TABLE zipcode_citystate;") | |
61 | print 'Current ZIP table dropped.' | |
62 | table_exists = False | |
63 | exit() | |
64 | else: | |
65 | print 'ZIP table does not exist, cannot drop.' | |
66 | exit() | |
67 | ||
68 | if table_exists: | |
69 | # count current rows, see if we have any info in there already | |
70 | dbp.execute("SELECT count(*) FROM zipcode_citystate;") | |
71 | rows = int(dbp.fetchone()[0]) | |
72 | if rows > 0: | |
73 | print 'ZIP codes already exist in database. Please run:' | |
74 | print sys.argv[0], '--drop' | |
75 | print 'in order to clear the current ZIP table and populate with new ZIP codes' | |
76 | exit(1) | |
77 | else: | |
78 | print 'ZIP Table does not yet exist.' | |
79 | except Exception, err: | |
80 | print 'Error setting up initial table information: {}'.format(err) | |
81 | raise | |
82 | ||
83 | # create table | |
84 | try: | |
85 | dbp.execute("CREATE TABLE IF NOT EXISTS zipcode_citystate (zipcode INT(5), city VARCHAR(15) NOT NULL, state VARCHAR(2) NOT NULL, UNIQUE KEY (zipcode));") | |
86 | print 'ZIP Table created.' | |
87 | except Exception, err: | |
88 | print 'Error creating database table zipcode_citystate: {}'.format(err) | |
89 | raise | |
90 | ||
91 | # load zipfile contents | |
92 | if rows == 0: | |
93 | with open(zipcode_csv_file, 'rb') as zipfile: | |
94 | print 'Populating ZIP Table' | |
95 | ||
96 | current_row = 0 | |
97 | first_line = True # contains row info, not necessary and breaks our int/float/etc | |
98 | ||
99 | for line in zipfile: | |
100 | if first_line: | |
101 | first_line = False | |
102 | continue | |
103 | ||
104 | # insertion | |
105 | line = line.replace('\n', '').replace('"', '').split(',') | |
106 | if len(line) > 3: | |
107 | zipcode, city, state, latitude, longitude, timezone, dst = line | |
108 | ||
109 | try: | |
110 | dbp.execute("INSERT INTO zipcode_citystate (zipcode, city, state) VALUES (%s,%s,%s);", [int(zipcode), city, state]) | |
111 | if debug: | |
112 | print 'Inserted ZIP code', zipcode, ':', city, ':', state | |
113 | except Exception, err: | |
114 | print 'Error adding ZIP code to database: ({})'.format(err) | |
115 | print line | |
116 | print zipcode, city, state | |
117 | raise | |
118 | ||
119 | # status | |
120 | current_row += 1 | |
121 | if current_row % 1000 == 0: # print every 1k rows | |
122 | print current_row, 'rows inserted' | |
123 | ||
124 | print current_row, 'rows inserted in total' |