Pad missing columns. Fixes #252.

This commit is contained in:
Friedrich Lindenberg 2018-05-29 21:55:25 +03:00
parent f04629825f
commit 881127880a
2 changed files with 15 additions and 1 deletions

View File

@ -12,7 +12,7 @@ from sqlalchemy.exc import NoSuchTableError
from dataset.types import Types from dataset.types import Types
from dataset.util import normalize_column_name, index_name, ensure_tuple from dataset.util import normalize_column_name, index_name, ensure_tuple
from dataset.util import DatasetException, ResultIter, QUERY_STEP from dataset.util import DatasetException, ResultIter, QUERY_STEP
from dataset.util import normalize_table_name from dataset.util import normalize_table_name, pad_chunk_columns
log = logging.getLogger(__name__) log = logging.getLogger(__name__)
@ -127,10 +127,12 @@ class Table(object):
row = self._sync_columns(row, ensure, types=types) row = self._sync_columns(row, ensure, types=types)
chunk.append(row) chunk.append(row)
if len(chunk) == chunk_size: if len(chunk) == chunk_size:
chunk = pad_chunk_columns(chunk)
self.table.insert().execute(chunk) self.table.insert().execute(chunk)
chunk = [] chunk = []
if len(chunk): if len(chunk):
chunk = pad_chunk_columns(chunk)
self.table.insert().execute(chunk) self.table.insert().execute(chunk)
def update(self, row, keys, ensure=None, types=None, return_count=False): def update(self, row, keys, ensure=None, types=None, return_count=False):

View File

@ -102,3 +102,15 @@ def ensure_tuple(obj):
if isinstance(obj, Sequence) and not isinstance(obj, six.string_types): if isinstance(obj, Sequence) and not isinstance(obj, six.string_types):
return tuple(obj) return tuple(obj)
return obj, return obj,
def pad_chunk_columns(chunk):
"""Given a set of items to be inserted, make sure they all have the
same columns by padding columns with None if they are missing."""
columns = set()
for record in chunk:
columns.update(record.keys())
for record in chunk:
for column in columns:
record.setdefault(column, None)
return chunk