Merge branch 'raw_migrations' of https://github.com/carrotquest/infi.clickhouse_orm into carrotquest-raw_migrations

docs/schema_migrations.md

@@ -59,6 +59,35 @@ A compound migration operation for altering a buffer table and its underlying on
 Applying this migration operation to a regular table has the same effect as an `AlterTable` operation.
 
 
+**RunPython**
+
+A migration operation that runs python function inside migration.
+ 
+    def forward(database):
+        database.insert([
+             TestModel(field=1)
+        ])
+ 
+    operations = [
+        RunPython(forward),
+    ]
+
+
+**RunSQL**
+
+A migration operation that runs raw SQL queries inside migration. 
+SQL parameter can be a string or array of SQL-query strings
+Example:
+  
+    operations = [
+        RunSQL('INSERT INTO `test_table` (field) VALUES (1)'),
+        RunSQL([
+          'INSERT INTO `test_table` (field) VALUES (2)',
+          'INSERT INTO `test_table` (field) VALUES (3)'
+         ])
+    ]
+
+
 Running Migrations
 ------------------
 

src/infi/clickhouse_orm/migrations.py

@@ -1,3 +1,5 @@
+import six
+
 from .models import Model, BufferModel
 from .fields import DateField, StringField
 from .engines import MergeTree
@@ -114,6 +116,37 @@ class DropTable(Operation):
         database.drop_table(self.model_class)
 
 
+class RunPython(Operation):
+    '''
+    A migration operation that executes given python function on database
+    '''
+    def __init__(self, func):
+        assert callable(func), "'func' parameter must be function"
+        self._func = func
+
+    def apply(self, database):
+        logger.info('    Executing python operation %s', self._func.__name__)
+        self._func(database)
+
+
+class RunSQL(Operation):
+    '''
+    A migration operation that executes given SQL on database
+    '''
+
+    def __init__(self, sql):
+        if isinstance(sql, six.string_types):
+            sql = [sql]
+
+        assert isinstance(sql, list), "'sql' parameter must be string or list of strings"
+        self._sql = sql
+
+    def apply(self, database):
+        logger.info('    Executing raw SQL operations')
+        for item in self._sql:
+            database.raw(item)
+
+
 class MigrationHistory(Model):
     '''
     A model for storing which migrations were already applied to the containing database.

tests/sample_migrations/0012.py

@@ -0,0 +1,9 @@
+from infi.clickhouse_orm import migrations
+
+operations = [
+    migrations.RunSQL("INSERT INTO `mig` (date, f1, f3, f4) VALUES ('2016-01-01', 1, 1, 'test') "),
+    migrations.RunSQL([
+        "INSERT INTO `mig` (date, f1, f3, f4) VALUES ('2016-01-02', 2, 2, 'test2') ",
+        "INSERT INTO `mig` (date, f1, f3, f4) VALUES ('2016-01-03', 3, 3, 'test3') ",
+    ])
+]

tests/sample_migrations/0013.py

@@ -0,0 +1,15 @@
+import datetime
+
+from infi.clickhouse_orm import migrations
+from test_migrations import Model3
+
+
+def forward(database):
+    database.insert([
+        Model3(date=datetime.date(2016, 1, 4), f1=4, f3=1, f4='test4')
+    ])
+
+
+operations = [
+    migrations.RunPython(forward)
+]

tests/test_migrations.py

@@ -80,6 +80,16 @@ class MigrationsTestCase(unittest.TestCase):
         self.assertEquals(self.getTableFields(Model4), [('date', 'Date'), ('f3', 'DateTime'), ('f2', 'String')])
         self.assertEquals(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])
+
 
 # Several different models with the same table name, to simulate a table that changes over time