diff --git a/docs/schema_migrations.md b/docs/schema_migrations.md index e4647bf..4ca0400 100644 --- a/docs/schema_migrations.md +++ b/docs/schema_migrations.md @@ -50,6 +50,35 @@ A migration operation that compares the table of a given model class to the mode Default values are not altered by this 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 ------------------ diff --git a/src/infi/clickhouse_orm/migrations.py b/src/infi/clickhouse_orm/migrations.py index 1a82a5b..e347c31 100644 --- a/src/infi/clickhouse_orm/migrations.py +++ b/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 @@ -95,6 +97,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. diff --git a/tests/sample_migrations/0010.py b/tests/sample_migrations/0010.py new file mode 100644 index 0000000..7a688cf --- /dev/null +++ b/tests/sample_migrations/0010.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') ", + ]) +] \ No newline at end of file diff --git a/tests/sample_migrations/0011.py b/tests/sample_migrations/0011.py new file mode 100644 index 0000000..5ee6498 --- /dev/null +++ b/tests/sample_migrations/0011.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) +] \ No newline at end of file diff --git a/tests/test_migrations.py b/tests/test_migrations.py index 3478f9f..18174e8 100644 --- a/tests/test_migrations.py +++ b/tests/test_migrations.py @@ -70,6 +70,16 @@ class MigrationsTestCase(unittest.TestCase): self.assertEquals(self.getTableFields(AliasModel), [('date', 'Date'), ('date_alias', "Date")]) + self.database.migrate('tests.sample_migrations', 10) + 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', 11) + 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