from __future__ import unicode_literals import unittest from infi.clickhouse_orm.database import Database from infi.clickhouse_orm.models import Model, BufferModel from infi.clickhouse_orm.fields import * from infi.clickhouse_orm.engines import * from infi.clickhouse_orm.migrations import MigrationHistory from enum import Enum # Add tests to path so that migrations will be importable import sys, os sys.path.append(os.path.dirname(__file__)) import logging logging.basicConfig(level=logging.DEBUG, format='%(message)s') logging.getLogger("requests").setLevel(logging.WARNING) class MigrationsTestCase(unittest.TestCase): def setUp(self): self.database = Database('test-db', log_statements=True) self.database.drop_table(MigrationHistory) def tearDown(self): self.database.drop_database() def tableExists(self, model_class): query = "EXISTS TABLE $db.`%s`" % model_class.table_name() return next(self.database.select(query)).result == 1 def getTableFields(self, model_class): query = "DESC `%s`.`%s`" % (self.database.db_name, model_class.table_name()) return [(row.name, row.type) for row in self.database.select(query)] def test_migrations(self): # Creation and deletion of table self.database.migrate('tests.sample_migrations', 1) self.assertTrue(self.tableExists(Model1)) self.database.migrate('tests.sample_migrations', 2) self.assertFalse(self.tableExists(Model1)) self.database.migrate('tests.sample_migrations', 3) self.assertTrue(self.tableExists(Model1)) # Adding, removing and altering simple fields self.assertEqual(self.getTableFields(Model1), [('date', 'Date'), ('f1', 'Int32'), ('f2', 'String')]) self.database.migrate('tests.sample_migrations', 4) self.assertEqual(self.getTableFields(Model2), [('date', 'Date'), ('f1', 'Int32'), ('f3', 'Float32'), ('f2', 'String'), ('f4', 'String'), ('f5', 'Array(UInt64)')]) self.database.migrate('tests.sample_migrations', 5) self.assertEqual(self.getTableFields(Model3), [('date', 'Date'), ('f1', 'Int64'), ('f3', 'Float64'), ('f4', 'String')]) # Altering enum fields self.database.migrate('tests.sample_migrations', 6) self.assertTrue(self.tableExists(EnumModel1)) self.assertEqual(self.getTableFields(EnumModel1), [('date', 'Date'), ('f1', "Enum8('dog' = 1, 'cat' = 2, 'cow' = 3)")]) self.database.migrate('tests.sample_migrations', 7) self.assertTrue(self.tableExists(EnumModel1)) self.assertEqual(self.getTableFields(EnumModel2), [('date', 'Date'), ('f1', "Enum16('dog' = 1, 'cat' = 2, 'horse' = 3, 'pig' = 4)")]) # Materialized fields and alias fields self.database.migrate('tests.sample_migrations', 8) self.assertTrue(self.tableExists(MaterializedModel)) self.assertEqual(self.getTableFields(MaterializedModel), [('date_time', "DateTime"), ('date', 'Date')]) self.database.migrate('tests.sample_migrations', 9) self.assertTrue(self.tableExists(AliasModel)) self.assertEqual(self.getTableFields(AliasModel), [('date', 'Date'), ('date_alias', "Date")]) # Buffer models creation and alteration self.database.migrate('tests.sample_migrations', 10) self.assertTrue(self.tableExists(Model4)) self.assertTrue(self.tableExists(Model4Buffer)) self.assertEqual(self.getTableFields(Model4), [('date', 'Date'), ('f1', 'Int32'), ('f2', 'String')]) self.assertEqual(self.getTableFields(Model4Buffer), [('date', 'Date'), ('f1', 'Int32'), ('f2', 'String')]) self.database.migrate('tests.sample_migrations', 11) self.assertEqual(self.getTableFields(Model4), [('date', 'Date'), ('f3', 'DateTime'), ('f2', 'String')]) self.assertEqual(self.getTableFields(Model4Buffer), [('date', 'Date'), ('f3', 'DateTime'), ('f2', 'String')]) self.database.migrate('tests.sample_migrations', 12) self.assertEqual(self.database.count(Model3), 3) data = [item.f1 for item in self.database.select('SELECT f1 FROM $table ORDER BY f1', model_class=Model3)] self.assertListEqual(data, [1, 2, 3]) self.database.migrate('tests.sample_migrations', 13) self.assertEqual(self.database.count(Model3), 4) data = [item.f1 for item in self.database.select('SELECT f1 FROM $table ORDER BY f1', model_class=Model3)] self.assertListEqual(data, [1, 2, 3, 4]) self.database.migrate('tests.sample_migrations', 14) self.assertTrue(self.tableExists(MaterializedModel1)) self.assertEqual(self.getTableFields(MaterializedModel1), [('date_time', 'DateTime'), ('int_field', 'Int8'), ('date', 'Date'), ('int_field_plus_one', 'Int8')]) self.assertTrue(self.tableExists(AliasModel1)) self.assertEqual(self.getTableFields(AliasModel1), [('date', 'Date'), ('int_field', 'Int8'), ('date_alias', 'Date'), ('int_field_plus_one', 'Int8')]) self.database.migrate('tests.sample_migrations', 15) self.assertTrue(self.tableExists(Model4_compressed)) if self.database.has_low_cardinality_support: self.assertEqual(self.getTableFields(Model2LowCardinality), [('date', 'Date'), ('f1', 'LowCardinality(Int32)'), ('f3', 'LowCardinality(Float32)'), ('f2', 'LowCardinality(String)'), ('f4', 'LowCardinality(Nullable(String))'), ('f5', 'Array(LowCardinality(UInt64))')]) else: logging.warning('No support for low cardinality') self.assertEqual(self.getTableFields(Model2), [('date', 'Date'), ('f1', 'Int32'), ('f3', 'Float32'), ('f2', 'String'), ('f4', 'Nullable(String)'), ('f5', 'Array(UInt64)')]) # Several different models with the same table name, to simulate a table that changes over time class Model1(Model): date = DateField() f1 = Int32Field() f2 = StringField() engine = MergeTree('date', ('date',)) @classmethod def table_name(cls): return 'mig' class Model2(Model): date = DateField() f1 = Int32Field() f3 = Float32Field() f2 = StringField() f4 = StringField() f5 = ArrayField(UInt64Field()) # addition of an array field engine = MergeTree('date', ('date',)) @classmethod def table_name(cls): return 'mig' class Model3(Model): date = DateField() f1 = Int64Field() # changed from Int32 f3 = Float64Field() # changed from Float32 f4 = StringField() engine = MergeTree('date', ('date',)) @classmethod def table_name(cls): return 'mig' class EnumModel1(Model): date = DateField() f1 = Enum8Field(Enum('SomeEnum1', 'dog cat cow')) engine = MergeTree('date', ('date',)) @classmethod def table_name(cls): return 'enum_mig' class EnumModel2(Model): date = DateField() f1 = Enum16Field(Enum('SomeEnum2', 'dog cat horse pig')) # changed type and values engine = MergeTree('date', ('date',)) @classmethod def table_name(cls): return 'enum_mig' class MaterializedModel(Model): date_time = DateTimeField() date = DateField(materialized='toDate(date_time)') engine = MergeTree('date', ('date',)) @classmethod def table_name(cls): return 'materalized_date' class MaterializedModel1(Model): date_time = DateTimeField() date = DateField(materialized='toDate(date_time)') int_field = Int8Field() int_field_plus_one = Int8Field(materialized='int_field + 1') engine = MergeTree('date', ('date',)) @classmethod def table_name(cls): return 'materalized_date' class AliasModel(Model): date = DateField() date_alias = DateField(alias='date') engine = MergeTree('date', ('date',)) @classmethod def table_name(cls): return 'alias_date' class AliasModel1(Model): date = DateField() date_alias = DateField(alias='date') int_field = Int8Field() int_field_plus_one = Int8Field(alias='int_field + 1') engine = MergeTree('date', ('date',)) @classmethod def table_name(cls): return 'alias_date' class Model4(Model): date = DateField() f1 = Int32Field() f2 = StringField() engine = MergeTree('date', ('date',)) @classmethod def table_name(cls): return 'model4' class Model4Buffer(BufferModel, Model4): engine = Buffer(Model4) @classmethod def table_name(cls): return 'model4buffer' class Model4_changed(Model): date = DateField() f3 = DateTimeField() f2 = StringField() engine = MergeTree('date', ('date',)) @classmethod def table_name(cls): return 'model4' class Model4Buffer_changed(BufferModel, Model4_changed): engine = Buffer(Model4_changed) @classmethod def table_name(cls): return 'model4buffer' class Model4_compressed(Model): date = DateField() f3 = DateTimeField(codec='Delta,ZSTD(10)') f2 = StringField(codec='LZ4HC') engine = MergeTree('date', ('date',)) @classmethod def table_name(cls): return 'model4' class Model2LowCardinality(Model): date = DateField() f1 = LowCardinalityField(Int32Field()) f3 = LowCardinalityField(Float32Field()) f2 = LowCardinalityField(StringField()) f4 = LowCardinalityField(NullableField(StringField())) f5 = ArrayField(LowCardinalityField(UInt64Field())) engine = MergeTree('date', ('date',)) @classmethod def table_name(cls): return 'mig'