diff --git a/buildout.cfg b/buildout.cfg index f5bd512..f0c8519 100644 --- a/buildout.cfg +++ b/buildout.cfg @@ -13,7 +13,7 @@ install_requires = [ 'pytz', 'requests', 'setuptools' -] + ] version_file = src/infi/clickhouse_utils/__version__.py description = A Python library for working with the ClickHouse database long_description = A Python library for working with the ClickHouse database @@ -47,6 +47,8 @@ eggs = ${project:name} infi.unittest infi.traceback zc.buildout +scripts = ipython + nosetests interpreter = python [pack] diff --git a/src/infi/clickhouse_utils/database.py b/src/infi/clickhouse_utils/database.py new file mode 100644 index 0000000..b7c0548 --- /dev/null +++ b/src/infi/clickhouse_utils/database.py @@ -0,0 +1,38 @@ +import requests + + +class DatabaseException(Exception): + pass + + +class Database(object): + + def __init__(self, db_name, db_url='http://localhost:8123/', username=None, password=None): + self.db_name = db_name + self.db_url = db_url + self.username = username + self.password = password + self._send('CREATE DATABASE IF NOT EXISTS ' + db_name) + + def create_table(self, model_class): + self._send(model_class.create_table_sql(self.db_name)) + + def drop_table(self, model_class): + self._send(model_class.drop_table_sql(self.db_name)) + + def drop_database(self): + self._send('DROP DATABASE ' + self.db_name) + + def _send(self, sql, settings=None): + params = self._build_params(settings) + r = requests.post(self.db_url, params=params, data=sql) + if r.status_code != 200: + raise DatabaseException(r.text) + + def _build_params(self, settings): + params = dict(settings or {}) + if self.username: + params['username'] = username + if self.password: + params['password'] = password + return params diff --git a/src/infi/clickhouse_utils/models.py b/src/infi/clickhouse_utils/models.py index c4074f7..f9fad43 100644 --- a/src/infi/clickhouse_utils/models.py +++ b/src/infi/clickhouse_utils/models.py @@ -31,15 +31,21 @@ class Model(object): return cls.__name__.lower() @classmethod - def create_table_sql(cls, db): - parts = ['CREATE TABLE IF NOT EXISTS %s.%s (' % (db, cls.table_name())] + def create_table_sql(cls, db_name): + parts = ['CREATE TABLE IF NOT EXISTS %s.%s (' % (db_name, cls.table_name())] + cols = [] for name, field in cls._fields: default = field.get_db_prep_value(field.default) - parts.append(' %s %s DEFAULT %s,' % (name, field.db_type, escape(default))) + cols.append(' %s %s DEFAULT %s' % (name, field.db_type, escape(default))) + parts.append(', \n'.join(cols)) parts.append(')') parts.append('ENGINE = ' + cls.engine.create_table_sql()) return '\n'.join(parts) + @classmethod + def drop_table_sql(cls, db_name): + return 'DROP TABLE IF EXISTS %s.%s' % (db_name, cls.table_name()) + @classmethod def from_tsv(cls, line): ''' diff --git a/tests/test_orm.py b/tests/test_orm.py new file mode 100644 index 0000000..c4f8d74 --- /dev/null +++ b/tests/test_orm.py @@ -0,0 +1,32 @@ +import unittest + +from infi.clickhouse_utils.database import Database +from infi.clickhouse_utils.models import Model +from infi.clickhouse_utils.fields import * +from infi.clickhouse_utils.engines import * + + +class ORMTestCase(unittest.TestCase): + + def setUp(self): + self.database = Database('test_db') + + def tearDown(self): + self.database.drop_database() + + def test_create_table(self): + self.database.create_table(Person) + self.database.drop_table(Person) + + +class Person(Model): + + first_name = StringField() + last_name = StringField() + birthday = DateField() + height = Float32Field() + + engine = MergeTree('birthday', ('first_name', 'last_name', 'birthday')) + + +