· 8 years ago · Feb 02, 2018, 09:10 AM
1from collections import OrderedDict
2import datetime
3from typing import Iterator, List, Sized, Union
4
5import numpy as np
6import pandas as pd
7from psycopg2.extensions import QuotedString
8from sqlalchemy import and_, exists, MetaData, Table, Column as SAColumn
9
10import logging
11log = logging.getLogger()
12
13
14#: Number of rows to insert per batch transaction
15BATCH_SIZE = 50000
16
17
18def to_python_type(column):
19 if str(column.type) == 'UUID':
20 return str
21 return column.type.python_type
22
23
24def to_str(val):
25 if isinstance(val, bytes):
26 val = val.decode('utf-8')
27 return QuotedString((str(val) or "").encode('utf-8')).getquoted().decode('utf-8')
28
29
30def to_date(value):
31 dt = pd.to_datetime(value)
32 return to_str(dt.date().isoformat() if dt else None)
33
34
35def to_datetime(value):
36 dt = pd.to_datetime(value)
37 return to_str(dt.to_pydatetime().isoformat() if dt else None)
38
39
40class Column:
41 def __init__(self, name: str, python_type: type):
42 """Wrapper to cast Python values for use in ad-hoc SQL.
43
44 Example::
45
46 columns = [Column('id', int), Column('amount', float)]
47
48 :param name: Name of the column.
49 :param python_type: Python type e.g. int, str, float.
50 """
51 self.name = name
52 self.python_type = python_type
53
54 def escape(self, value) -> str:
55 """Escape a value for use in a Postgres ad-hoc SQL statement."""
56 if pd.isnull(value):
57 return 'NULL'
58
59 func = self.python_type
60
61 if isinstance(value, (datetime.datetime, np.datetime64, pd.Timestamp)) or \
62 func in (datetime.date, datetime.datetime):
63 func = to_datetime
64 elif isinstance(value, datetime.date):
65 print(self.name)
66 func = to_date
67 elif issubclass(self.python_type, str):
68 func = to_str
69
70 return func(value)
71
72 def __eq__(self, b):
73 return self.name == b.name and self.python_type == b.python_type
74
75 def __repr__(self):
76 return '{}<name={}, type={}>'.format(
77 self.__class__.__name__, self.name, self.python_type.__name__)
78
79
80class ColumnCollection(OrderedDict):
81 def __init__(self, columns: list):
82 super().__init__([(c.name, c) for c in columns])
83
84
85class BulkInsertFromIterator:
86 def __init__(self, table, data: Iterator, columns: list,
87 batch_size: int=BATCH_SIZE, header: bool=False):
88 """Bulk insert into Postgres from an iterator in fixed-size batches.
89
90 Example::
91
92 bulk = BulkInsertFromIterator(
93 'table.name',
94 iter([[1, 'Python'], [2, 'PyPy', 3]]),
95 [Column('id', int), Column('name', str)]
96 )
97 bulk.execute(db.engine.raw_connection)
98
99 :param table: Name of the table.
100 :param data: Iterable containing the data to insert.
101 :param columns: List of :class:`Column` objects.
102 :param batch_size: Rows to insert per batch.
103 :param header: True if the first row is a header.
104 """
105 self.table = table
106 self.data = data
107 self.columns = columns
108 self.batch_size = batch_size
109 self.header = header
110
111 if isinstance(self.data, list):
112 self.data = iter(self.data)
113
114 if not isinstance(self.data, Iterator):
115 raise TypeError('Expected Iterator, got {}'.format(
116 self.data.__class__))
117
118 if not self.columns:
119 raise ValueError('Columns cannot be empty')
120
121 if isinstance(self.columns[0], tuple):
122 self.columns = [Column(*c) for c in self.columns]
123
124 def batch_execute(self, conn):
125 """Insert data in batches of `batch_size`.
126
127 :param conn: A DB API 2.0 connection object
128 """
129 def batches(data, batch_size):
130 """Return batches of length `batch_size` from any object that
131 supports iteration without knowing length."""
132 rv = []
133 for idx, line in enumerate(data):
134 if idx != 0 and idx % batch_size == 0:
135 yield rv
136 rv = []
137 rv.append(line)
138 yield rv
139
140 columns = ColumnCollection(self.columns)
141 if self.header:
142 self.columns = [columns.get(h) for h in next(self.data)]
143 columns = ColumnCollection(self.columns)
144
145 total = 0
146 query = BulkInsertQuery(self.table, columns)
147 for batch in batches(self.data, self.batch_size):
148 total += query.execute(conn, batch) or 0
149 yield total
150
151 def execute(self, conn):
152 """Execute all batches."""
153 return max(list(self.batch_execute(conn)))
154
155
156class BulkInsertQuery:
157 def __init__(self, table: str, columns):
158 """Execute a multi-row INSERT statement.
159
160 This does not take advantage of parameterized queries, but escapes
161 string values manually in :class:`Column`.
162
163 :param table: Name of the table being inserted into.
164 :param columns: Columns required for type coercion.
165 """
166 self.table = table
167 self.columns = columns
168 self.query = 'INSERT INTO {} ({}) VALUES '.format(
169 table, ', '.join([c for c in columns]))
170
171 def execute(self, conn, rows: list) -> int:
172 """Execute a single multi-row INSERT for `rows`.
173
174 :param conn: Function that returns a database connection
175 :param rows: List of tuples in the same order as :attr:`columns`.
176 """
177 if not len(rows):
178 raise ValueError('No data provided')
179 if len(self.columns) != len(rows[0]):
180 raise ValueError('Expecting {} columns, found {}'.format(
181 len(self.columns), len(rows[0])))
182
183 # Clone the data
184 rows = list(rows)
185
186 conn = conn()
187 cursor = conn.cursor()
188 try:
189 cursor.execute(self.query + ', '.join(self.escape_rows(rows)))
190 conn.commit()
191 finally:
192 cursor.close()
193 conn.close()
194
195 return len(rows)
196
197 def escape_rows(self, rows: list):
198 """Escape values for use in non-parameterized SQL queries.
199
200 :param rows: List of values to escape.
201 """
202 def to_tuple(values):
203 rv = []
204 for column in self.columns:
205 rv.append(self.columns.get(column).escape(values[column]))
206 return tuple(rv)
207
208 for idx, row in enumerate(rows):
209 data = to_tuple(row)
210 rows[idx] = '({})'.format(', '.join(map(str, data)))
211
212 return rows
213
214
215def as_columns(columns) -> List[Column]:
216 rv = []
217 for column in columns:
218 if isinstance(column, Column):
219 rv.append(column)
220 if isinstance(column, tuple):
221 rv.append(Column(*column))
222 if isinstance(column, str):
223 rv.append(Column(column, str))
224 if isinstance(column, SAColumn):
225 rv.append(Column(column.name, to_python_type(column)))
226 return rv
227
228
229def from_sqlalchemy_table(table: Table, data: Iterator, columns: List[str],
230 batch_size: int=BATCH_SIZE) -> BulkInsertFromIterator:
231 """Return a :class:`BulkInsertFromIterator` based on the metadata
232 of a SQLAlchemy table.
233
234 Example::
235
236 batch = from_sqlalchemy_table(
237 Rating.__table__,
238 data,
239 ['rating_id', 'repo_id', 'login_id', 'rating']
240 )
241
242 :param table: A :class:`sqlalchemy.Table` instance.
243 :param data: An iterator.
244 :param columns: List of column names to use.
245 :param batch_size: Number of rows to insert per SQL statement
246 """
247 if not isinstance(table, Table):
248 raise TypeError('Expected sqlalchemy.Table, got {}'.format(table))
249
250 wrapped = []
251 for name in columns:
252 column = table.columns.get(name)
253 wrapped.append(Column(str(column.name), to_python_type(column)))
254
255 return BulkInsertFromIterator(table, data, wrapped, batch_size, False)
256
257
258def create_staging_table(engine, table: Table) -> Table:
259 """Create a copy of the table to store intermediary results.
260
261 Primary keys and other unique constraints are removed.
262
263 :param engine: SQLAlchemy engine
264 :param table: SQLAlchemy table to clone schema from
265 """
266 table = table.tometadata(MetaData(), schema="staging")
267
268 # Remove constraints to prevent errors
269 for column in table.columns:
270 if column.primary_key:
271 column.primary_key = False
272
273 table.indexes = []
274 table.constraints = []
275 table.primary_key = None
276
277 log.info('Creating staging table {}.{}'.format(table.schema, table.name))
278 table.drop(engine, checkfirst=True)
279 table.create(engine)
280
281 return table
282
283
284def stage_and_merge(engine, target: Table, rows: Union[Iterator, Sized]):
285 """Write data to an intermediary staging table before adding to `target`.
286
287 :param engine: A instance of :class:`sqlalchemy.engine.Engine`
288 :param target: Table to write the results
289 :param rows: Data to write
290 """
291 if isinstance(rows, Sized) and len(rows) > 0:
292 log.info('Staging Rows: {}'.format(len(rows)))
293
294 # Drop & recreate the staging table
295 source = create_staging_table(engine, target)
296
297 # Insert data into a temporary staging table prior to copying to the target
298 try:
299 bulk = BulkInsertFromIterator(source, rows, as_columns(source.columns))
300 bulk.execute(engine.raw_connection)
301
302 keys = filter(lambda c: c.primary_key, target.columns)
303 where = map(lambda c: source.c[c.name] == target.c[c.name], keys)
304
305 # Only insert rows that do not exist in the target table
306 query = source.select().distinct().where(~exists().where(and_(*where)))
307 result = engine.execute(target.insert().from_select(source.c, query))
308
309 log.info('Updated Row Count: {}'.format(result.rowcount))
310 finally:
311 source.drop(engine)
312
313
314def stage_and_replace(engine, target: Table, rows: Union[Iterator, Sized]):
315 if isinstance(rows, Sized) and len(rows) > 0:
316 log.info('Staging Rows: {}'.format(len(rows)))
317
318 # Drop & recreate the staging table
319 source = create_staging_table(engine, target)
320
321 # Insert data into a temporary staging table prior to copying to the target
322 try:
323 bulk = BulkInsertFromIterator(source, rows, as_columns(source.columns))
324 bulk.execute(engine.raw_connection)
325
326 # Re-create the target table prior to inserting
327 if target.exists(engine):
328 target.drop(engine)
329 target.create(engine)
330
331 query = source.select().distinct()
332 result = engine.execute(target.insert().from_select(source.c, query))
333
334 log.info('Updated Row Count: {}'.format(result.rowcount))
335 finally:
336 source.drop(engine)
337
338
339def determine_columns(table: Table, rows):
340 columns = as_columns(table.columns)
341 if not isinstance(rows[0], dict):
342 return columns
343 keys = rows[0].keys()
344 return list(filter(lambda c: c.name in keys, columns))
345
346
347def stage_and_update(engine, target: Table, rows: Union[Iterator, Sized]):
348 """Write data to an intermediary staging table before adding to `target`.
349
350 :param engine: A instance of :class:`sqlalchemy.engine.Engine`
351 :param target: Table to write the results
352 :param rows: Data to write
353 """
354 if isinstance(rows, Sized) and len(rows) > 0:
355 log.info('Staging Rows: {}'.format(len(rows)))
356
357 # Drop & recreate the staging table
358 source = create_staging_table(engine, target)
359 columns = determine_columns(source, rows)
360
361 # Insert data into a temporary staging table prior to copying to the target
362 try:
363 bulk = BulkInsertFromIterator(source, rows, columns)
364 bulk.execute(engine.raw_connection)
365
366 keys = filter(lambda c: c.primary_key, target.columns)
367 where = map(lambda c: source.c[c.name] == target.c[c.name], keys)
368
369 # Delete from target table before appending
370 delete = target.delete().where(exists().where(and_(*where)))
371 engine.execute(delete)
372
373 # Copy rows from staging table to target
374 insert = source.select()
375 result = engine.execute(target.insert().from_select(source.c, insert))
376
377 log.info('Updated Row Count: {}'.format(result.rowcount))
378 finally:
379 source.drop(engine)