Columnstore: Adding tests and error handling for non TEXT columns

fetzerms · amotl · commit c4cedf3eea49 · 2023-06-08T11:26:46.000-07:00
diff --git a/src/crate/client/sqlalchemy/compiler.py b/src/crate/client/sqlalchemy/compiler.py
@@ -25,6 +25,7 @@
 import sqlalchemy as sa
 from sqlalchemy.dialects.postgresql.base import PGCompiler
 from sqlalchemy.sql import compiler
+from sqlalchemy.types import String
 from .types import MutableDict, _Craty, Geopoint, Geoshape
 from .sa_version import SA_VERSION, SA_1_4
 
@@ -129,6 +130,11 @@ def get_column_specification(self, column, **kwargs):
             colspec += " INDEX OFF"
 
         if column.dialect_options['crate'].get('columnstore') is False:
+            if not isinstance(column.type, (String, )):
+                raise sa.exc.CompileError(
+                    "Controlling the columnstore is only allowed for STRING columns"
+                )
+
             colspec += " STORAGE WITH (columnstore = false)"
 
         return colspec
diff --git a/src/crate/client/sqlalchemy/tests/create_table_test.py b/src/crate/client/sqlalchemy/tests/create_table_test.py
@@ -20,6 +20,7 @@
 # software solely pursuant to the terms of the relevant commercial agreement.
 
 import sqlalchemy as sa
+
 try:
     from sqlalchemy.orm import declarative_base
 except ImportError:
@@ -31,7 +32,6 @@
 from unittest import TestCase
 from unittest.mock import patch, MagicMock
 
-
 fake_cursor = MagicMock(name='fake_cursor')
 FakeCursor = MagicMock(name='FakeCursor', spec=Cursor)
 FakeCursor.return_value = fake_cursor
@@ -77,6 +77,7 @@ class DummyTable(self.Base):
             __tablename__ = 'dummy'
             pk = sa.Column(sa.String, primary_key=True)
             obj_col = sa.Column(Object)
+
         self.Base.metadata.create_all(bind=self.engine)
         fake_cursor.execute.assert_called_with(
             ('\nCREATE TABLE dummy (\n\tpk STRING NOT NULL, \n\tobj_col OBJECT, '
@@ -91,6 +92,7 @@ class DummyTable(self.Base):
             }
             pk = sa.Column(sa.String, primary_key=True)
             p = sa.Column(sa.String)
+
         self.Base.metadata.create_all(bind=self.engine)
         fake_cursor.execute.assert_called_with(
             ('\nCREATE TABLE t (\n\t'
@@ -105,6 +107,7 @@ class DummyTable(self.Base):
             __tablename__ = 't'
             ts = sa.Column(sa.BigInteger, primary_key=True)
             p = sa.Column(sa.BigInteger, sa.Computed("date_trunc('day', ts)"))
+
         self.Base.metadata.create_all(bind=self.engine)
         fake_cursor.execute.assert_called_with(
             ('\nCREATE TABLE t (\n\t'
@@ -119,6 +122,7 @@ class DummyTable(self.Base):
             __tablename__ = 't'
             ts = sa.Column(sa.BigInteger, primary_key=True)
             p = sa.Column(sa.BigInteger, sa.Computed("date_trunc('day', ts)", persisted=False))
+
         with self.assertRaises(sa.exc.CompileError):
             self.Base.metadata.create_all(bind=self.engine)
 
@@ -131,6 +135,7 @@ class DummyTable(self.Base):
             }
             pk = sa.Column(sa.String, primary_key=True)
             p = sa.Column(sa.String)
+
         self.Base.metadata.create_all(bind=self.engine)
         fake_cursor.execute.assert_called_with(
             ('\nCREATE TABLE t (\n\t'
@@ -166,6 +171,7 @@ class DummyTable(self.Base):
             }
             pk = sa.Column(sa.String, primary_key=True)
             p = sa.Column(sa.String, primary_key=True)
+
         self.Base.metadata.create_all(bind=self.engine)
         fake_cursor.execute.assert_called_with(
             ('\nCREATE TABLE t (\n\t'
@@ -207,6 +213,7 @@ def test_column_pk_nullable(self):
         class DummyTable(self.Base):
             __tablename__ = 't'
             pk = sa.Column(sa.String, primary_key=True, nullable=True)
+
         with self.assertRaises(sa.exc.CompileError):
             self.Base.metadata.create_all(bind=self.engine)
 
@@ -230,5 +237,33 @@ class DummyTable(self.Base):
             __tablename__ = 't'
             pk = sa.Column(sa.String, primary_key=True)
             a = sa.Column(Geopoint, crate_index=False)
+
+        with self.assertRaises(sa.exc.CompileError):
+            self.Base.metadata.create_all(bind=self.engine)
+
+    def test_text_column_without_columnstore(self):
+        class DummyTable(self.Base):
+            __tablename__ = 't'
+            pk = sa.Column(sa.String, primary_key=True)
+            a = sa.Column(sa.String, crate_columnstore=False)
+            b = sa.Column(sa.String, crate_columnstore=True)
+            c = sa.Column(sa.String)
+
+        self.Base.metadata.create_all(bind=self.engine)
+
+        fake_cursor.execute.assert_called_with(
+            ('\nCREATE TABLE t (\n\t'
+             'pk STRING NOT NULL, \n\t'
+             'a STRING STORAGE WITH (columnstore = false), \n\t'
+             'b STRING, \n\t'
+             'c STRING, \n\t'
+             'PRIMARY KEY (pk)\n)\n\n'), ())
+
+    def test_non_text_column_without_columnstore(self):
+        class DummyTable(self.Base):
+            __tablename__ = 't'
+            pk = sa.Column(sa.String, primary_key=True)
+            a = sa.Column(sa.Integer, crate_columnstore=False)
+
         with self.assertRaises(sa.exc.CompileError):
             self.Base.metadata.create_all(bind=self.engine)
