Fixed bug with batched insert_tuple

This commit is contained in:
M1ha 2019-02-15 17:42:04 +05:00
parent ab2fe24f74
commit 2798cbb602
2 changed files with 19 additions and 1 deletions

View File

@ -97,7 +97,8 @@ class Database(InfiDatabase):
buf = BytesIO() buf = BytesIO()
query = 'INSERT INTO `%s`.`%s` (%s) FORMAT TabSeparated\n' \ query = 'INSERT INTO `%s`.`%s` (%s) FORMAT TabSeparated\n' \
% (self.db_name, model_class.table_name(), fields_list) % (self.db_name, model_class.table_name(), fields_list)
buf.write(query.encode('utf-8')) query_enc = query.encode('utf-8')
buf.write(query_enc)
buf.write(tuple_to_csv(first_tuple).encode('utf-8')) buf.write(tuple_to_csv(first_tuple).encode('utf-8'))
# Collect lines in batches of batch_size # Collect lines in batches of batch_size
@ -112,6 +113,7 @@ class Database(InfiDatabase):
yield buf.getvalue() yield buf.getvalue()
# Start a new batch # Start a new batch
buf = BytesIO() buf = BytesIO()
buf.write(query_enc)
lines = 0 lines = 0
# Return any remaining lines in partial batch # Return any remaining lines in partial batch

View File

@ -31,6 +31,22 @@ class CollapsingMergeTreeTest(TestCase):
'str_field': str(i) 'str_field': str(i)
} for i in range(10)], [item.to_dict() for item in qs]) } for i in range(10)], [item.to_dict() for item in qs])
def test_insert_tuples_batch_size(self):
tuple_class = ClickHouseTestModel.get_tuple_class()
data = [
tuple_class(id=i, created_date=date.today(), value=i, str_field=str(i))
for i in range(10)
]
self.db.insert_tuples(ClickHouseTestModel, data, batch_size=2)
qs = ClickHouseTestModel.objects.order_by('id').all()
self.assertListEqual([{
'id': i,
'created_date': date.today(),
'value': i,
'str_field': str(i)
} for i in range(10)], [item.to_dict() for item in qs])
def test_insert_tuples_special_characters(self): def test_insert_tuples_special_characters(self):
tuple_class = ClickHouseTestModel.get_tuple_class() tuple_class = ClickHouseTestModel.get_tuple_class()
data = [tuple_class(id=1, created_date=date.today(), value=1, str_field='\t')] data = [tuple_class(id=1, created_date=date.today(), value=1, str_field='\t')]