Support disabling indexing in SQLAlchemy ORM column definitions

Co-authored-by: Andreas Motl <[email protected]>

Author: hammerhead

CHANGES.txt

@@ -5,6 +5,8 @@ Changes for crate
 Unreleased
 ==========
 
+- SQLAlchemy: Added support for ``crate_index`` and ``nullable`` attributes in
+  ORM column definitions.
 
 2022/12/02 0.28.0
 =================

docs/sqlalchemy.rst

@@ -181,9 +181,9 @@ system`_::
     ...     }
     ...
     ...     id = sa.Column(sa.String, primary_key=True, default=gen_key)
-    ...     name = sa.Column(sa.String)
+    ...     name = sa.Column(sa.String, crate_index=False)
     ...     name_normalized = sa.Column(sa.String, sa.Computed("lower(name)"))
-    ...     quote = sa.Column(sa.String)
+    ...     quote = sa.Column(sa.String, nullable=False)
     ...     details = sa.Column(types.Object)
     ...     more_details = sa.Column(types.ObjectArray)
     ...     name_ft = sa.Column(sa.String)
@@ -201,6 +201,8 @@ In this example, we:
 - Use the ``gen_key`` function to provide a default value for the ``id`` column
   (which is also the primary key)
 - Use standard SQLAlchemy types for the ``id``, ``name``, and ``quote`` columns
+- Use ``nullable=False`` to define a ``NOT NULL`` constraint
+- Disable indexing of the ``name`` column using ``crate_index=False``
 - Define a computed column ``name_normalized`` (based on ``name``) that
   translates into a generated column
 - Use the `Object`_ extension type for the ``details`` column
@@ -250,7 +252,7 @@ A table schema like this
 .. code-block:: sql
 
    CREATE TABLE "doc"."logs" (
-     "ts" TIMESTAMP WITH TIME ZONE,
+     "ts" TIMESTAMP WITH TIME ZONE NOT NULL,
      "level" TEXT,
      "message" TEXT
    )

src/crate/client/sqlalchemy/compiler.py

@@ -25,7 +25,7 @@
 import sqlalchemy as sa
 from sqlalchemy.dialects.postgresql.base import PGCompiler
 from sqlalchemy.sql import compiler, crud, selectable
-from .types import MutableDict
+from .types import MutableDict, _Craty, Geopoint, Geoshape
 from .sa_version import SA_VERSION, SA_1_4
 
 
@@ -119,6 +119,15 @@ def get_column_specification(self, column, **kwargs):
                 "Primary key columns cannot be nullable"
             )
 
+        if column.dialect_options['crate'].get('index') is False:
+            if isinstance(column.type, (Geopoint, Geoshape, _Craty)):
+                raise sa.exc.CompileError(
+                    "Disabling indexing is not supported for column "
+                    "types OBJECT, GEO_POINT, and GEO_SHAPE"
+                )
+
+            colspec += " INDEX OFF"
+
         return colspec
 
     def visit_computed_column(self, generated):

src/crate/client/sqlalchemy/tests/create_table_test.py

@@ -22,7 +22,7 @@
 import sqlalchemy as sa
 from sqlalchemy.ext.declarative import declarative_base
 
-from crate.client.sqlalchemy.types import Object, ObjectArray
+from crate.client.sqlalchemy.types import Object, ObjectArray, Geopoint
 from crate.client.cursor import Cursor
 
 from unittest import TestCase
@@ -41,7 +41,7 @@ def setUp(self):
         self.engine = sa.create_engine('crate://')
         self.Base = declarative_base(bind=self.engine)
 
-    def test_create_table_with_basic_types(self):
+    def test_table_basic_types(self):
         class User(self.Base):
             __tablename__ = 'users'
             string_col = sa.Column(sa.String, primary_key=True)
@@ -69,7 +69,7 @@ class User(self.Base):
              '\n\tPRIMARY KEY (string_col)\n)\n\n'),
             ())
 
-    def test_with_obj_column(self):
+    def test_column_obj(self):
         class DummyTable(self.Base):
             __tablename__ = 'dummy'
             pk = sa.Column(sa.String, primary_key=True)
@@ -80,7 +80,7 @@ class DummyTable(self.Base):
              '\n\tPRIMARY KEY (pk)\n)\n\n'),
             ())
 
-    def test_with_clustered_by(self):
+    def test_table_clustered_by(self):
         class DummyTable(self.Base):
             __tablename__ = 't'
             __table_args__ = {
@@ -97,7 +97,7 @@ class DummyTable(self.Base):
              ') CLUSTERED BY (p)\n\n'),
             ())
 
-    def test_with_computed_column(self):
+    def test_column_computed(self):
         class DummyTable(self.Base):
             __tablename__ = 't'
             ts = sa.Column(sa.BigInteger, primary_key=True)
@@ -111,15 +111,15 @@ class DummyTable(self.Base):
              ')\n\n'),
             ())
 
-    def test_with_virtual_computed_column(self):
+    def test_column_computed_virtual(self):
         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()
 
-    def test_with_partitioned_by(self):
+    def test_table_partitioned_by(self):
         class DummyTable(self.Base):
             __tablename__ = 't'
             __table_args__ = {
@@ -137,7 +137,7 @@ class DummyTable(self.Base):
              ') PARTITIONED BY (p)\n\n'),
             ())
 
-    def test_with_number_of_shards_and_replicas(self):
+    def test_table_number_of_shards_and_replicas(self):
         class DummyTable(self.Base):
             __tablename__ = 't'
             __table_args__ = {
@@ -154,7 +154,7 @@ class DummyTable(self.Base):
              ') CLUSTERED INTO 3 SHARDS WITH (NUMBER_OF_REPLICAS = 2)\n\n'),
             ())
 
-    def test_with_clustered_by_and_number_of_shards(self):
+    def test_table_clustered_by_and_number_of_shards(self):
         class DummyTable(self.Base):
             __tablename__ = 't'
             __table_args__ = {
@@ -172,7 +172,7 @@ class DummyTable(self.Base):
              ') CLUSTERED BY (p) INTO 3 SHARDS\n\n'),
             ())
 
-    def test_table_with_object_array(self):
+    def test_column_object_array(self):
         class DummyTable(self.Base):
             __tablename__ = 't'
             pk = sa.Column(sa.String, primary_key=True)
@@ -185,7 +185,7 @@ class DummyTable(self.Base):
              'tags ARRAY(OBJECT), \n\t'
              'PRIMARY KEY (pk)\n)\n\n'), ())
 
-    def test_table_with_nullable(self):
+    def test_column_nullable(self):
         class DummyTable(self.Base):
             __tablename__ = 't'
             pk = sa.Column(sa.String, primary_key=True)
@@ -200,9 +200,32 @@ class DummyTable(self.Base):
              'b INT NOT NULL, \n\t'
              'PRIMARY KEY (pk)\n)\n\n'), ())
 
-    def test_with_pk_nullable(self):
+    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()
+
+    def test_column_crate_index(self):
+        class DummyTable(self.Base):
+            __tablename__ = 't'
+            pk = sa.Column(sa.String, primary_key=True)
+            a = sa.Column(sa.Integer, crate_index=False)
+            b = sa.Column(sa.Integer, crate_index=True)
+
+        self.Base.metadata.create_all()
+        fake_cursor.execute.assert_called_with(
+            ('\nCREATE TABLE t (\n\t'
+             'pk STRING NOT NULL, \n\t'
+             'a INT INDEX OFF, \n\t'
+             'b INT, \n\t'
+             'PRIMARY KEY (pk)\n)\n\n'), ())
+
+    def test_column_geopoint_without_index(self):
+        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()