/cartodb-utils.py
Last active Mar 13, 2017
| import os | |
| import urllib | |
| import urllib2 | |
| import base64 | |
| import json | |
| import sys | |
| import argparse | |
| try: | |
| import requests | |
| except ImportError: | |
| print 'The requests package is required: http://docs.python-requests.org/en/latest/user/install/#install' | |
| sys.exit() | |
| class CartoDB: | |
| def __init__(self, options): | |
| # do stuff | |
| self.options = options | |
| self.api_url = "https://%s.cartodb.com/api/v2/sql" % (self.options['u']) | |
| self.import_url = "https://%s.cartodb.com/api/v1/imports/?api_key=%s" % (self.options['u'], self.options['k']) | |
| self.new_tables = [] | |
| self.internal_columns = ['created_at', 'updated_at', 'the_geom', 'the_geom_webmercator', 'cartodb_id'] | |
| self.type_map = {'string':'text', 'boolean': 'boolean', 'date': 'timestamp', 'number':'numeric'} | |
| def _log(self, message): | |
| if self.options['verbose'] == True: | |
| print message | |
| def _error(self, error): | |
| print error | |
| sys.exit() | |
| def sql_api(self, sql): | |
| # execute sql request over API | |
| params = { | |
| 'api_key' : self.options["k"], | |
| 'q' : sql | |
| } | |
| r = requests.get(self.api_url, params=params) | |
| return r.json() | |
| def upload(self): | |
| # import a file | |
| # see https://gist.github.com/lbosque/5876697 | |
| # returns new table name | |
| r = requests.post(self.import_url, files={'file': open(self.options['f'], 'rb')}) | |
| data = r.json() | |
| if data['success']!=True: | |
| self._error("Upload failed") | |
| complete = False | |
| last_state = '' | |
| while not complete: | |
| import_status_url = "https://%s.cartodb.com/api/v1/imports/%s?api_key=%s" % (self.options['u'], data['item_queue_id'], self.options['k']) | |
| req = urllib2.Request(import_status_url) | |
| response = urllib2.urlopen(req) | |
| d = json.loads(str(response.read())) | |
| if last_state!=d['state']: | |
| last_state=d['state'] | |
| if d['state']=='uploading': | |
| self._log('Uploading file...') | |
| elif d['state']=='importing': | |
| self._log('Importing data...') | |
| elif d['state']=='complete': | |
| complete = True | |
| self._log('Table "%s" created' % d['table_name']) | |
| if d['state']=='failure': | |
| self._error(d['get_error_text']['what_about']) | |
| self.new_tables.append(d['table_name']) | |
| return d['table_name'] | |
| def columns(self, table): | |
| sql = "SELECT * FROM %s LIMIT 0" % table | |
| data = self.sql_api(sql) | |
| return data['fields'] | |
| def add_column(self, table, name, coltype): | |
| sql = "ALTER TABLE %s ADD COLUMN %s %s" % (table, name, coltype) | |
| data = self.sql_api(sql) | |
| return True | |
| def overwrite(self, append=False): | |
| # upload new data | |
| new_table = self.upload() | |
| source_columns = self.columns(new_table) | |
| target_columns = self.columns(self.options['t']) | |
| insert_cols = {} | |
| alter_cols = [] | |
| for c in source_columns.keys(): | |
| if c in self.internal_columns: | |
| source_columns.pop(c, None) | |
| else: | |
| if c not in target_columns.keys(): | |
| insert_cols[c] = self.type_map[source_columns[c]['type']] | |
| alter_cols.append(c) | |
| else: | |
| insert_cols[c] = self.type_map[target_columns[c]['type']] | |
| for c in alter_cols: | |
| self.add_column(self.options['t'], c, insert_cols[c]) | |
| select_list = [] | |
| for c,t in insert_cols.items(): | |
| select_list.append( "%s::%s" % (c,t)) | |
| sql = "INSERT INTO %s (the_geom, %s) " % (self.options['t'], ','.join(insert_cols.keys())) | |
| sql += "SELECT the_geom, %s FROM %s; " % (','.join(select_list), new_table) | |
| sql += "DROP TABLE %s" % new_table | |
| self._log("Writing data to %s and droppping %s" % (self.options['t'],new_table)) | |
| if append==False: | |
| sql = "DELETE FROM %s; %s " % (self.options['t'], sql) | |
| data = self.sql_api(sql) | |
| if 'error' in data.keys(): | |
| self._log('Overwrite failed, cleaning-up') | |
| sql = "DROP TABLE %s" % new_table | |
| self.sql_api(sql) | |
| return False | |
| else: | |
| return True | |
| def drop_table(self): | |
| # drop a table ' | |
| self._log("Dropping table %s" % self.options['t']) | |
| sql = "DROP TABLE %s" % self.options['t'] | |
| data = self.sql_api(sql) | |
| if 'error' in data.keys(): | |
| self._error(data['error']) | |
| return True | |
| def clear_rows(self): | |
| # clear all rows from a table | |
| self._log("Deleting all rows") | |
| sql = "DELETE FROM %s" % self.options['t'] | |
| data = self.sql_api(sql) | |
| if 'error' in data.keys(): | |
| self._error(data['error']) | |
| return True | |
| def export_table(self): | |
| self._log("Exporting new %s" % self.options['m']) | |
| params = {"format": self.options['m'], "api_key": self.options["k"],"q": "SELECT * FROM %s" % self.options["t"]} | |
| r = requests.get(self.api_url, params=params, stream=True) | |
| with open(self.options['l'], 'wb') as fd: | |
| for chunk in r.iter_content(10): | |
| fd.write(chunk) | |
| return True | |
| def clean_table(self): | |
| # clean up table for speed | |
| self._log("Cleaning up unused space") | |
| sql = "VACUUM FULL %s" % self.options['t'] | |
| data = self.sql_api(sql) | |
| if 'error' in data.keys(): | |
| self._error(data['error']) | |
| self._log("Optimizing existing indexes") | |
| sql = "ANALYZE %s" % self.options['t'] | |
| data = self.sql_api(sql) | |
| if 'error' in data.keys(): | |
| self._error(data['error']) | |
| return True | |
| if __name__ == "__main__": | |
| SUPPORTED_METHODS = { | |
| 'import' : { | |
| "description": "Import a file to create a new table", | |
| "requirements": ["f","k","u"], | |
| "example": "python cartodb-utils.py import -f myfile.csv -k myapikey -u myusername" | |
| }, | |
| 'overwrite' : { | |
| "description": "Overwrite an existing table with data from a file", | |
| "requirements": ["f","k","u","t"], | |
| "example": "python cartodb-utils.py overwrite -f myfile.csv -t some_existing_table -k myapikey -u myusername" | |
| }, | |
| 'append' : { | |
| "description": "Append rows to an existing table from a file", | |
| "requirements": ["f","k","u","t"], | |
| "example": "python cartodb-utils.py append -f myfile.csv -t some_existing_table -k myapikey -u myusername" | |
| }, | |
| 'clear' : { | |
| "description": "Clear all rows from an existing table", | |
| "requirements": ["k","u","t"], | |
| "example": "python cartodb-utils.py clear -t some_existing_table -k myapikey -u myusername" | |
| }, | |
| 'drop' : { | |
| "description": "Completely drop an existing table", | |
| "requirements": ["k","u","t"], | |
| "example": "python cartodb-utils.py drop -t some_existing_table -k myapikey -u myusername" | |
| }, | |
| 'export' :{ | |
| "description": "Export an existing table to a local file (default GeoJSON)", | |
| "requirements": ["k","u","t","l"], | |
| "example": "python cartodb-utils.py export -t some_existing_table -m CSV -l local_file.csv -k myapikey -u myusername" | |
| }, | |
| 'clean' : { | |
| "description": "Vacuum and analyze a table for speed", | |
| "requirements": ["k","u","t"], | |
| "example": "python cartodb-utils.py clean -t some_existing_table -k myapikey -u myusername" | |
| } | |
| } | |
| parser = argparse.ArgumentParser(description="CartoDB Python Utility") | |
| parser.add_argument('method', nargs="?", help='e.g. %s' % ','.join(SUPPORTED_METHODS.keys())) | |
| parser.add_argument('-f', '--file', dest='f', help='Source file') | |
| parser.add_argument('-l', '--local', dest='l', help='Local file') | |
| parser.add_argument('-m', '--format', default="GeoJSON", dest='m', help='Export file format') | |
| parser.add_argument('-u', '--user', dest='u', help='CartoDB username') | |
| parser.add_argument('-k', '--key', dest='k', help='CartoDB account API Key') | |
| parser.add_argument('-t', '--target', dest='t', help='Target table') | |
| parser.add_argument('-v', '--verbose', dest='verbose', default=False, action="store_true", help='Verbose output if included') | |
| args = parser.parse_args() | |
| options = vars(args) | |
| def success(message): | |
| print 'SUCCESS', message | |
| def failure(message): | |
| print 'FAILURE', message | |
| m = args.method.lower() | |
| if m in SUPPORTED_METHODS.keys(): | |
| for d in SUPPORTED_METHODS[m]["requirements"]: | |
| if options[d] is None: | |
| print "Arguement -%s is required\n\n%s\n\ndescription:\t%s\nrequired args:\t%s\nexample:\t%s" % (d,m,SUPPORTED_METHODS[m]['description'],SUPPORTED_METHODS[m]['requirements'],SUPPORTED_METHODS[m]['example']) | |
| sys.exit() | |
| cartodb = CartoDB(options) | |
| if args.method.lower() == 'import': | |
| new_table = cartodb.upload() | |
| success(new_table) | |
| if args.method.lower() == 'overwrite': | |
| status = cartodb.overwrite() | |
| if status == True: | |
| success('Table data replaced') | |
| if args.method.lower() == 'append': | |
| status = cartodb.overwrite(True) | |
| if status == True: | |
| success('Data appended to table') | |
| if args.method.lower() == 'clean': | |
| status = cartodb.clean_table() | |
| if status == True: | |
| success('Cleaned table') | |
| if args.method.lower() == 'drop': | |
| status = cartodb.drop_table() | |
| if status == True: | |
| success('Dropped table') | |
| if args.method.lower() == 'clear': | |
| status = cartodb.clear_rows() | |
| if status == True: | |
| success('Cleared all rows from table') | |
| if args.method.lower() == 'export': | |
| status = cartodb.export_table() | |
| if status == True: | |
| success('Exported table to %s' % options['l']) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment