Skip to content

Commit

Permalink
Enum8 and Enum16 supported (#263)
Browse files Browse the repository at this point in the history
* Enum8 and Enum16 supported
  • Loading branch information
akurdyukov authored Sep 28, 2023
1 parent 3ce18ad commit 1c7553d
Show file tree
Hide file tree
Showing 3 changed files with 92 additions and 1 deletion.
3 changes: 2 additions & 1 deletion clickhouse_sqlalchemy/types/common.py
Original file line number Diff line number Diff line change
Expand Up @@ -152,12 +152,13 @@ def __init__(self, precision=3, timezone=None):

class Enum(types.Enum, ClickHouseTypeEngine):
__visit_name__ = 'enum'
native_enum = True

def __init__(self, *enums, **kw):
if not enums:
enums = kw.get('_enums', ()) # passed as keyword

super(Enum, self).__init__(*enums, **kw)
super(Enum, self).__init__(*enums, **kw, convert_unicode=False)


class Enum8(Enum):
Expand Down
45 changes: 45 additions & 0 deletions tests/types/test_enum16.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,45 @@
import enum

from sqlalchemy import Column
from sqlalchemy.sql.ddl import CreateTable

from clickhouse_sqlalchemy import types, engines, Table
from tests.testcase import BaseTestCase, CompilationTestCase
from tests.util import with_native_and_http_sessions


class TestEnum(enum.IntEnum):
First = 1
Second = 2


class Enum16CompilationTestCase(CompilationTestCase):
def test_create_table(self):
table = Table(
"test",
BaseTestCase.metadata(),
Column("x", types.Enum16(TestEnum)),
engines.Memory(),
)

self.assertEqual(
self.compile(CreateTable(table)),
"CREATE TABLE test (x Enum16('First' = 1, 'Second' = 2)) "
"ENGINE = Memory"
)


@with_native_and_http_sessions
class Enum8TestCase(BaseTestCase):
table = Table(
'test', BaseTestCase.metadata(),
Column('x', types.Enum16(TestEnum)),
engines.Memory()
)

def test_select_insert(self):
value = TestEnum.First
with self.create_table(self.table):
self.session.execute(self.table.insert(), [{'x': value}])
self.assertEqual(self.session.query(self.table.c.x).scalar(),
value)
45 changes: 45 additions & 0 deletions tests/types/test_enum8.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,45 @@
import enum

from sqlalchemy import Column
from sqlalchemy.sql.ddl import CreateTable

from clickhouse_sqlalchemy import types, engines, Table
from tests.testcase import BaseTestCase, CompilationTestCase
from tests.util import with_native_and_http_sessions


class TestEnum(enum.IntEnum):
First = 1
Second = 2


class Enum8CompilationTestCase(CompilationTestCase):
def test_create_table(self):
table = Table(
"test",
BaseTestCase.metadata(),
Column("x", types.Enum8(TestEnum)),
engines.Memory(),
)

self.assertEqual(
self.compile(CreateTable(table)),
"CREATE TABLE test (x Enum8('First' = 1, 'Second' = 2)) "
"ENGINE = Memory"
)


@with_native_and_http_sessions
class Enum8TestCase(BaseTestCase):
table = Table(
'test', BaseTestCase.metadata(),
Column('x', types.Enum8(TestEnum)),
engines.Memory()
)

def test_select_insert(self):
value = TestEnum.First
with self.create_table(self.table):
self.session.execute(self.table.insert(), [{'x': value}])
self.assertEqual(self.session.query(self.table.c.x).scalar(),
value)

0 comments on commit 1c7553d

Please sign in to comment.