mirror of
https://github.com/Infinidat/infi.clickhouse_orm.git
synced 2024-09-21 03:08:53 +03:00
120 lines
4.2 KiB
Python
120 lines
4.2 KiB
Python
from __future__ import unicode_literals
|
|
import unittest
|
|
import six
|
|
from uuid import UUID
|
|
from infi.clickhouse_orm.database import Database
|
|
from infi.clickhouse_orm.fields import Field, Int16Field
|
|
from infi.clickhouse_orm.models import Model
|
|
from infi.clickhouse_orm.engines import Memory
|
|
from infi.clickhouse_orm.utils import escape
|
|
|
|
|
|
class CustomFieldsTest(unittest.TestCase):
|
|
|
|
def setUp(self):
|
|
self.database = Database('test-db')
|
|
|
|
def tearDown(self):
|
|
self.database.drop_database()
|
|
|
|
def test_boolean_field(self):
|
|
# Create a model
|
|
class TestModel(Model):
|
|
i = Int16Field()
|
|
f = BooleanField()
|
|
engine = Memory()
|
|
self.database.create_table(TestModel)
|
|
# Check valid values
|
|
for index, value in enumerate([1, '1', True, 0, '0', False]):
|
|
rec = TestModel(i=index, f=value)
|
|
self.database.insert([rec])
|
|
self.assertEquals([rec.f for rec in TestModel.objects_in(self.database).order_by('i')],
|
|
[True, True, True, False, False, False])
|
|
# Check invalid values
|
|
for value in [None, 'zzz', -5, 7]:
|
|
with self.assertRaises(ValueError):
|
|
TestModel(i=1, f=value)
|
|
|
|
def test_uuid_field(self):
|
|
# Create a model
|
|
class TestModel(Model):
|
|
i = Int16Field()
|
|
f = UUIDField()
|
|
engine = Memory()
|
|
self.database.create_table(TestModel)
|
|
# Check valid values (all values are the same UUID)
|
|
values = [
|
|
'{12345678-1234-5678-1234-567812345678}',
|
|
'12345678123456781234567812345678',
|
|
'urn:uuid:12345678-1234-5678-1234-567812345678',
|
|
'\x12\x34\x56\x78'*4,
|
|
(0x12345678, 0x1234, 0x5678, 0x12, 0x34, 0x567812345678),
|
|
0x12345678123456781234567812345678,
|
|
]
|
|
for index, value in enumerate(values):
|
|
rec = TestModel(i=index, f=value)
|
|
self.database.insert([rec])
|
|
for rec in TestModel.objects_in(self.database):
|
|
self.assertEquals(rec.f, UUID(values[0]))
|
|
# Check that ClickHouse encoding functions are supported
|
|
for rec in self.database.select("SELECT i, UUIDNumToString(f) AS f FROM testmodel", TestModel):
|
|
self.assertEquals(rec.f, UUID(values[0]))
|
|
for rec in self.database.select("SELECT 1 as i, UUIDStringToNum('12345678-1234-5678-1234-567812345678') AS f", TestModel):
|
|
self.assertEquals(rec.f, UUID(values[0]))
|
|
# Check invalid values
|
|
for value in [None, 'zzz', -1, '123']:
|
|
with self.assertRaises(ValueError):
|
|
TestModel(i=1, f=value)
|
|
|
|
|
|
class BooleanField(Field):
|
|
|
|
# The ClickHouse column type to use
|
|
db_type = 'UInt8'
|
|
|
|
# The default value if empty
|
|
class_default = False
|
|
|
|
def to_python(self, value, timezone_in_use):
|
|
# Convert valid values to bool
|
|
if value in (1, '1', True):
|
|
return True
|
|
elif value in (0, '0', False):
|
|
return False
|
|
else:
|
|
raise ValueError('Invalid value for BooleanField: %r' % value)
|
|
|
|
def to_db_string(self, value, quote=True):
|
|
# The value was already converted by to_python, so it's a bool
|
|
return '1' if value else '0'
|
|
|
|
|
|
class UUIDField(Field):
|
|
|
|
# The ClickHouse column type to use
|
|
db_type = 'FixedString(16)'
|
|
|
|
# The default value if empty
|
|
class_default = UUID(int=0)
|
|
|
|
def to_python(self, value, timezone_in_use):
|
|
# Convert valid values to UUID instance
|
|
if isinstance(value, UUID):
|
|
return value
|
|
elif isinstance(value, six.string_types):
|
|
return UUID(bytes=value.encode('latin1')) if len(value) == 16 else UUID(value)
|
|
elif isinstance(value, six.integer_types):
|
|
return UUID(int=value)
|
|
elif isinstance(value, tuple):
|
|
return UUID(fields=value)
|
|
else:
|
|
raise ValueError('Invalid value for UUIDField: %r' % value)
|
|
|
|
def to_db_string(self, value, quote=True):
|
|
# The value was already converted by to_python, so it's a UUID instance
|
|
val = value.bytes
|
|
if six.PY3:
|
|
val = str(val, 'latin1')
|
|
return escape(val, quote)
|
|
|