From 8488f9a9ae6453976b4724e8971f2e638bdf36c2 Mon Sep 17 00:00:00 2001 From: Alik Kurdyukov Date: Tue, 26 Sep 2023 17:09:59 +0400 Subject: [PATCH 1/3] Enum8 and Enum16 supported --- clickhouse_sqlalchemy/types/common.py | 3 +- tests/types/test_enum16.py | 43 +++++++++++++++++++++++++++ tests/types/test_enum8.py | 43 +++++++++++++++++++++++++++ 3 files changed, 88 insertions(+), 1 deletion(-) create mode 100644 tests/types/test_enum16.py create mode 100644 tests/types/test_enum8.py diff --git a/clickhouse_sqlalchemy/types/common.py b/clickhouse_sqlalchemy/types/common.py index 84cd78ff..947a8708 100644 --- a/clickhouse_sqlalchemy/types/common.py +++ b/clickhouse_sqlalchemy/types/common.py @@ -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): diff --git a/tests/types/test_enum16.py b/tests/types/test_enum16.py new file mode 100644 index 00000000..9d88cfaf --- /dev/null +++ b/tests/types/test_enum16.py @@ -0,0 +1,43 @@ +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) diff --git a/tests/types/test_enum8.py b/tests/types/test_enum8.py new file mode 100644 index 00000000..3a710aa4 --- /dev/null +++ b/tests/types/test_enum8.py @@ -0,0 +1,43 @@ +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) From aa3a31f5089542c40b060f02aa56f150c76dcd02 Mon Sep 17 00:00:00 2001 From: Alik Kurdyukov Date: Thu, 28 Sep 2023 22:45:29 +0400 Subject: [PATCH 2/3] flake8 fixed --- tests/types/test_enum16.py | 3 ++- tests/types/test_enum8.py | 3 ++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/types/test_enum16.py b/tests/types/test_enum16.py index 9d88cfaf..a86bdbcb 100644 --- a/tests/types/test_enum16.py +++ b/tests/types/test_enum16.py @@ -24,7 +24,8 @@ def test_create_table(self): self.assertEqual( self.compile(CreateTable(table)), - "CREATE TABLE test (x Enum16('First' = 1, 'Second' = 2)) ENGINE = Memory" + "CREATE TABLE test (x Enum16('First' = 1, 'Second' = 2)) " + "ENGINE = Memory" ) diff --git a/tests/types/test_enum8.py b/tests/types/test_enum8.py index 3a710aa4..ae4c4bb1 100644 --- a/tests/types/test_enum8.py +++ b/tests/types/test_enum8.py @@ -24,7 +24,8 @@ def test_create_table(self): self.assertEqual( self.compile(CreateTable(table)), - "CREATE TABLE test (x Enum8('First' = 1, 'Second' = 2)) ENGINE = Memory" + "CREATE TABLE test (x Enum8('First' = 1, 'Second' = 2)) " + "ENGINE = Memory" ) From 0c3349b027566c6f15801d3d3d2a2744d22f0bd0 Mon Sep 17 00:00:00 2001 From: Alik Kurdyukov Date: Thu, 28 Sep 2023 22:50:11 +0400 Subject: [PATCH 3/3] more flake8 fixes --- tests/types/test_enum16.py | 3 ++- tests/types/test_enum8.py | 3 ++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/types/test_enum16.py b/tests/types/test_enum16.py index a86bdbcb..a3dd61af 100644 --- a/tests/types/test_enum16.py +++ b/tests/types/test_enum16.py @@ -41,4 +41,5 @@ 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) + self.assertEqual(self.session.query(self.table.c.x).scalar(), + value) diff --git a/tests/types/test_enum8.py b/tests/types/test_enum8.py index ae4c4bb1..9f2afc02 100644 --- a/tests/types/test_enum8.py +++ b/tests/types/test_enum8.py @@ -41,4 +41,5 @@ 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) + self.assertEqual(self.session.query(self.table.c.x).scalar(), + value)