diff --git a/src/infi/clickhouse_orm/engines.py b/src/infi/clickhouse_orm/engines.py index 7c9a94c..c26b451 100644 --- a/src/infi/clickhouse_orm/engines.py +++ b/src/infi/clickhouse_orm/engines.py @@ -9,6 +9,7 @@ class MergeTree(Engine): def __init__(self, date_col, key_cols, sampling_expr=None, index_granularity=8192, replica_table_path=None, replica_name=None): + assert type(key_cols) in (list, tuple), 'key_cols must be a list or tuple' self.date_col = date_col self.key_cols = key_cols self.sampling_expr = sampling_expr @@ -54,6 +55,7 @@ class SummingMergeTree(MergeTree): def __init__(self, date_col, key_cols, summing_cols=None, sampling_expr=None, index_granularity=8192, replica_table_path=None, replica_name=None): super(SummingMergeTree, self).__init__(date_col, key_cols, sampling_expr, index_granularity, replica_table_path, replica_name) + assert type is None or type(summing_cols) in (list, tuple), 'summing_cols must be a list or tuple' self.summing_cols = summing_cols def _build_sql_params(self): @@ -63,6 +65,20 @@ class SummingMergeTree(MergeTree): return params +class ReplacingMergeTree(MergeTree): + + def __init__(self, date_col, key_cols, ver_col=None, sampling_expr=None, + index_granularity=8192, replica_table_path=None, replica_name=None): + super(ReplacingMergeTree, self).__init__(date_col, key_cols, sampling_expr, index_granularity, replica_table_path, replica_name) + self.ver_col = ver_col + + def _build_sql_params(self): + params = super(ReplacingMergeTree, self)._build_sql_params() + if self.ver_col: + params.append(self.ver_col) + return params + + class Buffer(Engine): """Here we define Buffer engine Read more here https://clickhouse.yandex/reference_en.html#Buffer diff --git a/tests/test_engines.py b/tests/test_engines.py new file mode 100644 index 0000000..d3d8865 --- /dev/null +++ b/tests/test_engines.py @@ -0,0 +1,64 @@ +import unittest + +from infi.clickhouse_orm.database import Database, DatabaseException +from infi.clickhouse_orm.models import Model +from infi.clickhouse_orm.fields import * +from infi.clickhouse_orm.engines import * + +import logging +logging.getLogger("requests").setLevel(logging.WARNING) + + +class EnginesTestCase(unittest.TestCase): + + def setUp(self): + self.database = Database('test-db') + + def tearDown(self): + self.database.drop_database() + + def _create_and_insert(self, model_class): + self.database.create_table(model_class) + self.database.insert([ + model_class(date='2017-01-01', event_id=23423, event_group=13, event_count=7, event_version=1) + ]) + + def test_merge_tree(self): + class TestModel(SampleModel): + engine = MergeTree('date', ('date', 'event_id', 'event_group')) + self._create_and_insert(TestModel) + + def test_merge_tree_with_sampling(self): + class TestModel(SampleModel): + engine = MergeTree('date', ('date', 'event_id', 'event_group'), sampling_expr='intHash32(event_id)') + self._create_and_insert(TestModel) + + def test_merge_tree_with_granularity(self): + class TestModel(SampleModel): + engine = MergeTree('date', ('date', 'event_id', 'event_group'), index_granularity=4096) + self._create_and_insert(TestModel) + + def test_collapsing_merge_tree(self): + class TestModel(SampleModel): + engine = CollapsingMergeTree('date', ('date', 'event_id', 'event_group'), 'event_version') + self._create_and_insert(TestModel) + + def test_summing_merge_tree(self): + class TestModel(SampleModel): + engine = SummingMergeTree('date', ('date', 'event_group'), ('event_count',)) + self._create_and_insert(TestModel) + + def test_replacing_merge_tree(self): + class TestModel(SampleModel): + engine = ReplacingMergeTree('date', ('date', 'event_id', 'event_group'), 'event_uversion') + self._create_and_insert(TestModel) + + +class SampleModel(Model): + + date = DateField() + event_id = UInt32Field() + event_group = UInt32Field() + event_count = UInt16Field() + event_version = Int8Field() + event_uversion = UInt8Field(materialized='abs(event_version)')