diff --git a/datum/cli.py b/datum/cli.py new file mode 100644 index 0000000..1f4186f --- /dev/null +++ b/datum/cli.py @@ -0,0 +1,39 @@ +import click +import datum + +import logging +log = logging.getLogger(__name__) + +@click.group() +def cli(): + pass + +@cli.command() +@click.option('--connection', '-d', help='The database connection string', required=True) +@click.argument('table') +def truncate(table, connection): + db = datum.connect(connection) + db.table(table).delete() + +@cli.command() +@click.option('--connection', '-d', help='The database connection string', required=True) +@click.argument('csvfile', type=click.File('rU')) +@click.argument('table') +def load(csvfile, table, connection): + db = datum.connect(connection) + db.table(table).load(csvfile) + +@cli.command() +@click.option('--connection', '-d', help='The database connection string', required=True) +@click.argument('sql') +def execute(sql, connection): + db = datum.connect(connection) + rows = db.execute(sql) + if datum.util.isiterable(rows): + import csv, sys + writer = csv.writer(sys.stdout) + writer.writerow(rows.header) + writer.writerows(rows) + +if __name__ == '__main__': + cli() diff --git a/datum/oracle_stgeom/database.py b/datum/oracle_stgeom/database.py index e653185..96542d4 100644 --- a/datum/oracle_stgeom/database.py +++ b/datum/oracle_stgeom/database.py @@ -16,7 +16,7 @@ def __init__(self, parent): self.password = p['password'] self.name = p['db_name'] if self.name: self.name = self.name.lower() - + dsn = '{user}/{password}@{host}'.format(**self.__dict__) if self.name: dsn += '/' + self.name @@ -28,6 +28,13 @@ def __init__(self, parent): def execute(self, stmt): self._c.execute(stmt) + + class SimpleView (list): + """ + A trivial wrapper on `list` to allow attaching a header + """ + pass + try: rows = self._c.fetchall() # Return rowcount for non-SELECT operations @@ -36,6 +43,10 @@ def execute(self, stmt): # Unpack single values if len(rows) > 0 and len(rows[0]) == 1: rows = [x[0] for x in rows] + + rows = SimpleView(rows) + rows.header = [field[0] for field in self._c.description] + return rows def close(self): @@ -43,7 +54,7 @@ def close(self): # def table(self, name): # # Check for a schema - # if '.' in + # if '.' in # return self.parent.table(name) @property @@ -67,7 +78,7 @@ def tables(self): def _dictify(self, geom_field=None): ''' - Turns query results into a list of dictionaries. This reads from the + Turns query results into a list of dictionaries. This reads from the cursor because calling fetchall() on rows breaks the geometry LOB. ''' fields = [x[0].lower() for x in self._c.description] @@ -104,7 +115,7 @@ def read(self, table, fields, geom_field=None, dictify=True, where=None, limit=N fields = list(fields) # Make a copy if fields != ['*']: if geom_field: - fields.append(self._wkt_getter(geom_field)) + fields.append(self._wkt_getter(geom_field)) fields_joined = ', '.join(fields) table = table.upper() stmt = "SELECT {} FROM {}".format(fields_joined, table) @@ -147,14 +158,14 @@ def save(self): def bulk_insert(self, table, rows, geom_field=None, from_srid=None, \ multi_geom=True, chunk_size=None): ''' - Inserts dictionary row objects in the the database + Inserts dictionary row objects in the the database Args: list of row dicts, table name ''' fields = rows[0].keys() if geom_field: non_geom_fields = [x for x in fields if x != geom_field] fields_joined = ', '.join(fields) - stmt = "INSERT INTO {} ({}) VALUES ".format(table, fields_joined) + stmt = "INSERT INTO {} ({}) VALUES ".format(table, fields_joined) len_rows = len(rows) if chunk_size is None or len_rows < chunk_size: diff --git a/datum/oracle_stgeom/table.py b/datum/oracle_stgeom/table.py index 17f1e3b..5eb0acf 100644 --- a/datum/oracle_stgeom/table.py +++ b/datum/oracle_stgeom/table.py @@ -5,7 +5,7 @@ from datum.oracle_stgeom.util import WktTransformer import cx_Oracle -# These are strings because one type (OBJECTVAR) isn't importable from +# These are strings because one type (OBJECTVAR) isn't importable from # the cx_Oracle module. FIELD_TYPE_MAP = { 'NUMBER': 'num', @@ -98,7 +98,7 @@ def _get_geom_type(self): bitand(eflags, 2), bitand(eflags, 4) + bitand(eflags, 8), bitand(eflags, 16), - bitand(eflags, 262144) + bitand(eflags, 262144) from sde.layers where owner = '{}' and @@ -179,7 +179,7 @@ def non_geom_fields(self): def _get_wkt_selector(self, to_srid=None): assert self.geom_field geom_field_t = geom_field = self.geom_field - # SDE.ST_Transform doesn't work when the datums differ. Unfortunately, + # SDE.ST_Transform doesn't work when the datums differ. Unfortunately, # 4326 <=> 2272 is one of those. Using Shapely + PyProj for now. # if to_srid and to_srid != self.srid: # geom_field_t = "SDE.ST_Transform({}, {})"\ @@ -209,7 +209,7 @@ def count(self): def read(self, fields=None, aliases=None, geom_field=None, to_srid=None, return_geom=True, limit=None, where=None, sort=None): - # If no geom_field was specified and we're supposed to return geom, + # If no geom_field was specified and we're supposed to return geom, # get it from the object. geom_field = geom_field or (self.geom_field if return_geom else None) @@ -234,13 +234,13 @@ def read(self, fields=None, aliases=None, geom_field=None, to_srid=None, stmt += " WHERE ROWNUM <= {}".format(limit) self._c.execute(stmt) - + # Handle aliases # fields = [re.sub('.+ AS ', '', x, flags=re.IGNORECASE) for x in fields] if aliases: fields = [aliases[x] if x in aliases else x for x in fields] - fields_lower = [x.lower() for x in fields] + fields_lower = [x.lower() for x in fields] if geom_field: geom_field_i = fields.index(geom_field) rows = [] @@ -264,7 +264,7 @@ def read(self, fields=None, aliases=None, geom_field=None, to_srid=None, for row in rows: geom = self._remove_m_value(row[geom_field_i]) row[geom_field_i] = geom - + # TODO if the WKT geom is single but the geom_type for the table # is multi, we may want to convert it. Seems to be working for now # though. @@ -279,7 +279,7 @@ def read(self, fields=None, aliases=None, geom_field=None, to_srid=None, for row in rows: geom = row[geom_field_l] geom_t = tsf.transform(geom) - row[geom_field_l] = geom_t + row[geom_field_l] = geom_t return rows @@ -347,21 +347,27 @@ def write(self, rows, from_srid=None, chunk_size=None): Args: list of row dicts, table name, ordered field names Originally this formed one big INSERT statement with a chunks of x - rows, but it's considerably faster to use the cx_Oracle `executemany` + rows, but it is considerably faster to use the cx_Oracle `executemany` function. See methods 1 and 2 below. TODO: it might be faster to call NEXTVAL on the DB sequence for OBJECTID rather than use the SDE helper function. """ - if len(rows) == 0: + try: + len_rows = len(rows) + except TypeError: + rows = tuple(rows) + len_rows = len(rows) + + if len_rows == 0: return # Get fields from the row because some fields from self.fields may be # optional, such as an autoincrementing PK. fields = rows[0].keys() # Sort so LOB fields are at the end - fields = sorted(fields, key=lambda x: 'lob' in self.metadata[x]['type']) - + fields = sorted(fields, key=lambda x: 'lob' in self.metadata[x.lower()]['type']) + geom_field = self.geom_field geom_type = self.geom_type srid = from_srid or self.srid @@ -412,10 +418,10 @@ def write(self, rows, from_srid=None, chunk_size=None): # # SELECT 1 FROM DUAL; # fields_joined = ', '.join(fields) # stmt = "INSERT ALL {} SELECT 1 FROM DUAL" - + # # We always have to pass in a value for OBJECTID (or whatever the SDE # # PK field is; sometimes it's something like OBJECTID_3). Check to see - # # if the user passed in a value for object ID (not likely), otherwise + # # if the user passed in a value for object ID (not likely), otherwise # # hardcode the sequence incrementor into the prepared statement. # if self.objectid_field in fields: # into_clause = "INTO {} ({}) VALUES ({{}})".format(self.name, \ @@ -423,38 +429,13 @@ def write(self, rows, from_srid=None, chunk_size=None): # else: # incrementor = "SDE.GDB_UTIL.NEXT_ROWID('{}', '{}')".format(self._owner, self.name) # into_clause = "INTO {} ({}, {}) VALUES ({{}}, {})".format(self.name, fields_joined, self.objectid_field, incrementor) - - # METHOD 2: executemany (not working with SDE.ST_Geometry call) - placeholders = [] - # Create placeholders for prepared statement - for field in fields: - type_ = type_map[field] - if type_ == 'geom': - placeholders.append('SDE.ST_Geometry(:{}, {})'\ - .format(field, self.srid)) - elif type_ == 'date': - # Insert an ISO-8601 timestring - placeholders.append("TO_TIMESTAMP(:{}, 'YYYY-MM-DD\"T\"HH24:MI:SS.FF\"+00:00\"')".format(field)) - else: - placeholders.append(':' + field) - - # Inject the object ID field if it's missing from the supplied rows - stmt_fields = list(fields) - if self.objectid_field and self.objectid_field not in fields: - stmt_fields.append(self.objectid_field) - incrementor = "SDE.GDB_UTIL.NEXT_ROWID('{}', '{}')"\ - .format(self._owner, self.name) - placeholders.append(incrementor) - # Prepare statement - placeholders_joined = ', '.join(placeholders) - stmt_fields_joined = ', '.join(stmt_fields) - stmt = "INSERT INTO {} ({}) VALUES ({})".format(self.name, \ - stmt_fields_joined, placeholders_joined) + # METHOD 2: executemany (not working with SDE.ST_Geometry call) + stmt = self.build_insert_statement(fields, type_map) self._c.prepare(stmt) # END OF METHODS - + len_rows = len(rows) if chunk_size is None or len_rows < chunk_size: iterations = 1 @@ -503,6 +484,42 @@ def write(self, rows, from_srid=None, chunk_size=None): # print(self._c.getbatcherrors()) self._save() + def build_insert_statement(self, input_fields, type_map): + placeholders = [] + + # Build up an exact-case mapping of field names + db_fields = self.fields + icase = {} + for field in input_fields: + db_field = [n for n in db_fields if field.lower() == n.lower()][0] + icase[field] = db_field + + # Create placeholders for prepared statement + for field in input_fields: + type_ = type_map[field] + if type_ == 'geom': + placeholders.append('SDE.ST_Geometry(:{}, {})'\ + .format(field, self.srid)) + elif type_ == 'date': + # Insert an ISO-8601 timestring + placeholders.append("TO_TIMESTAMP(:{}, 'YYYY-MM-DD\"T\"HH24:MI:SS.FF\"+00:00\"')".format(field)) + else: + placeholders.append(':' + field + '_') + + # Inject the object ID field if it's missing from the supplied rows + stmt_fields = [icase[field] for field in input_fields] + if self.objectid_field and self.objectid_field not in input_fields: + stmt_fields.append(self.objectid_field) + incrementor = "SDE.GDB_UTIL.NEXT_ROWID('{}', '{}')"\ + .format(self._owner, self.name) + placeholders.append(incrementor) + # Prepare statement + placeholders_joined = ', '.join(placeholders) + stmt_fields_joined = ', '.join(stmt_fields) + stmt = 'INSERT INTO {} ({}) VALUES ({})'.format(self.name, \ + stmt_fields_joined, placeholders_joined) + return stmt + def delete(self, cascade=False): """Delete all rows.""" name = self._name_p diff --git a/datum/postgis/table.py b/datum/postgis/table.py index bb5b6a3..5f53d1b 100644 --- a/datum/postgis/table.py +++ b/datum/postgis/table.py @@ -1,6 +1,7 @@ from collections import OrderedDict +from itertools import chain import re -from datum.util import dbl_quote +from datum.util import dbl_quote, chunks_of from psycopg2 import ProgrammingError @@ -8,6 +9,7 @@ 'integer': 'num', 'numeric': 'num', 'double precision': 'num', + 'real': 'num', 'text': 'text', 'character varying': 'text', 'date': 'date', @@ -30,6 +32,10 @@ def __init__(self, parent): def __str__(self): return 'Table: {}'.format(self.name) + @property + def schema(self): + return self._parent.schema + @property def name(self): return self._parent.name @@ -38,9 +44,11 @@ def name(self): def _name_p(self): """The table name prepared for SQL queries.""" name = self.name.lower() + schema = self.schema.lower() if self.schema else None + # Handle schema prefixes - if '.' in name: - return '.'.join([dbl_quote(x) for x in name.split('.')]) + if schema: + return '.'.join([dbl_quote(x) for x in (schema, name)]) else: return dbl_quote(name) @@ -97,10 +105,10 @@ def _get_srid(self): def _get_geom_type(self): stmt = """ - SELECT type - FROM geometry_columns - WHERE f_table_schema = 'public' - AND f_table_name = '{}' + SELECT type + FROM geometry_columns + WHERE f_table_schema = 'public' + AND f_table_name = '{}' and f_geometry_column = '{}'; """.format(self.name, self.geom_field) return self._exec(stmt)[0]['type'] @@ -133,7 +141,7 @@ def read(self, fields=None, aliases=None, geom_field=None, \ else: fields = [dbl_quote(x) for x in fields] if geom_field and return_geom: - fields.append(self._wkt_getter(geom_field, to_srid=to_srid)) + fields.append(self._wkt_getter(geom_field, to_srid=to_srid)) fields_joined = ', '.join(fields) stmt = "SELECT {} FROM {}".format(fields_joined, table_name) else: @@ -158,7 +166,7 @@ def read(self, fields=None, aliases=None, geom_field=None, \ def delete(self, cascade=False): """Delete all rows.""" - name = dbl_quote(self.name) + name = self._name_p # RESTART IDENTITY resets sequence generators. stmt = "TRUNCATE {} RESTART IDENTITY".format(name) stmt += ' CASCADE' if cascade else '' @@ -216,22 +224,35 @@ def _save(self): def write(self, rows, from_srid=None, chunk_size=None): """ - Inserts dictionary row objects in the the database + Inserts dictionary row objects in the the database Args: list of row dicts, table name, ordered field names """ - if len(rows) == 0: + # Split the rows into chunks of the given size, and pull off the first + # one. + if chunk_size is None: + first_chunk, chunks = list(rows), [] + else: + chunks = chunks_of(rows, chunk_size) + try: + first_chunk = list(next(chunks)) + except StopIteration: + return + + # Pull out the first row from the first chunk, if it is not empty. + if len(first_chunk) == 0: return + row = first_chunk[0] # Get fields from the row because some fields from self.fields may be # optional, such as autoincrementing integers. - fields = rows[0].keys() + fields = row.keys() geom_field = self.geom_field srid = from_srid or self.srid - row_geom_type = re.match('[A-Z]+', rows[0][geom_field]).group() \ + row_geom_type = re.match('[A-Z]+', row[geom_field]).group() \ if geom_field else None table_geom_type = self.geom_type if geom_field else None - # Do we need to cast the geometry to a MULTI type? (Assuming all rows + # Do we need to cast the geometry to a MULTI type? (Assuming all rows # have the same geom type.) if geom_field: if self.geom_type.startswith('MULTI') and \ @@ -244,7 +265,7 @@ def write(self, rows, from_srid=None, chunk_size=None): type_map = OrderedDict() for field in fields: try: - type_map[field] = [x['type'] for x in self.metadata if x['name'] == field][0] + type_map[field] = [x['type'] for x in self.metadata if x['name'].lower() == field.lower()][0] except IndexError: raise ValueError('Field `{}` does not exist'.format(field)) type_map_items = type_map.items() @@ -252,25 +273,11 @@ def write(self, rows, from_srid=None, chunk_size=None): fields_joined = ', '.join(fields) stmt = "INSERT INTO {} ({}) VALUES ".format(self.name, fields_joined) - len_rows = len(rows) - if chunk_size is None or len_rows < chunk_size: - iterations = 1 - else: - iterations = int(len_rows / chunk_size) - iterations += (len_rows % chunk_size > 0) # round up - - # Make list of value lists - for i in range(0, iterations): + for chunk in chain([first_chunk], chunks): val_rows = [] cur_stmt = stmt - if chunk_size: - start = i * chunk_size - end = min(len_rows, start + chunk_size) - else: - start = i - end = len_rows - for row in rows[start:end]: + for row in chunk: val_row = [] for field, type_ in type_map_items: if type_ == 'geom': @@ -280,7 +287,7 @@ def write(self, rows, from_srid=None, chunk_size=None): else: val = self._prepare_val(row[field], type_) - val_row.append(val) + val_row.append(val) val_rows.append(val_row) # Execute diff --git a/datum/table.py b/datum/table.py index 160187b..b90598c 100644 --- a/datum/table.py +++ b/datum/table.py @@ -65,11 +65,22 @@ def fields(self): """Returns a list of field names.""" return self._child.fields + def load(self, *infiles, **write_kwargs): + # See Table.write for the vaild write_kwargs + import csv + + if len(infiles) == 0: + infiles = (sys.stdin,) + + for infile in infiles: + reader = csv.DictReader(infile) + self.write(reader, **write_kwargs) + def read(self, fields=None, aliases=None, geom_field=None, to_srid=None, \ return_geom=True, limit=None, where=None, sort=None): """ Read rows from the database. - + ``` Parameters ---------- diff --git a/datum/util.py b/datum/util.py index 4ca72fc..cb523cc 100644 --- a/datum/util.py +++ b/datum/util.py @@ -1,4 +1,9 @@ from functools import partial +from itertools import islice, chain +try: + from itertools import izip_longest as zip_longest +except ImportError: + from itertools import zip_longest from six.moves.urllib.parse import urlparse def dbl_quote(text): @@ -15,3 +20,33 @@ def parse_url(url): 'db_name': p.path[1:] if p.path else None, } return comps + +def chunks_of(iterable, size): + """ + Return chunks of a max size of the iterable + Thanks to http://stackoverflow.com/a/24527424 + """ + if not size: + raise ValueError('Size must be an integer greater than 0') + + try: + iterator = iter(iterable) + except TypeError: + raise TypeError('First argument must be an iterable object, not {}' + .format(type(iterable).__name__)) + + # Pull a single element off of the iterator to ensure + # that it is not empty. This will stop looping when the + # iterator is depleted. + for first in iterator: + # Yield an iterable filled out with the apprpriate + # number of remaining elements from the iterator. + yield chain([first], islice(iterator, size - 1)) + +def isiterable(obj): + try: + iter(obj) + except TypeError: + return False + else: + return True diff --git a/requirements.txt b/requirements.txt index b6e34eb..a4304d9 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1 +1,2 @@ six==1.10.0 +click==6.6 diff --git a/setup.py b/setup.py index a582c0c..ab9e54c 100644 --- a/setup.py +++ b/setup.py @@ -1,4 +1,15 @@ from setuptools import setup +import os + + +def get_packages(package): + """ + Return root package and all sub-packages. + """ + return [dirpath + for dirpath, dirnames, filenames in os.walk(package) + if os.path.exists(os.path.join(dirpath, '__init__.py'))] + setup(name='datum', version='0.1', @@ -7,11 +18,14 @@ author='City of Philadelphia', author_email='maps@phila.gov', license='MIT', - packages=['datum'], - install_requires=['six==1.10.0'], + packages=get_packages('datum'), + install_requires=['six==1.10.0', 'click==6.6'], extras_require={ - 'oracle_stgeom': ['cx-Oracle==5.2.1', 'pyproj==1.9.5.1'], + 'oracle_stgeom': ['cx-Oracle==5.2.1', 'pyproj==1.9.5.1', 'shapely==1.5.17'], 'postgis': ['psycopg2==2.6.1'], }, + entry_points={ + 'console_scripts': ['datum=datum.cli:cli'] + }, # entry_points={'console_scripts': ['ais=ais:manager.run']}, zip_safe=False)