Add nullable and crate_index in ORM column definition

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
 
 
@@ -107,11 +107,27 @@ class CrateDDLCompiler(compiler.DDLCompiler):
     def get_column_specification(self, column, **kwargs):
         colspec = self.preparer.format_column(column) + " " + \
             self.dialect.type_compiler.process(column.type)
-        # TODO: once supported add default / NOT NULL here
+        # TODO: once supported add default here
 
         if column.computed is not None:
             colspec += " " + self.process(column.computed)
 
+        if column.nullable is False:
+            colspec += " NOT NULL"
+        elif column.nullable and column.primary_key:
+            raise sa.exc.CompileError(
+                "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)
@@ -59,7 +59,7 @@ class User(self.Base):
 
         self.Base.metadata.create_all()
         fake_cursor.execute.assert_called_with(
-            ('\nCREATE TABLE users (\n\tstring_col STRING, '
+            ('\nCREATE TABLE users (\n\tstring_col STRING NOT NULL, '
              '\n\tunicode_col STRING, \n\ttext_col STRING, \n\tint_col INT, '
              '\n\tlong_col1 LONG, \n\tlong_col2 LONG, '
              '\n\tbool_col BOOLEAN, '
@@ -69,18 +69,18 @@ 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)
             obj_col = sa.Column(Object)
         self.Base.metadata.create_all()
         fake_cursor.execute.assert_called_with(
-            ('\nCREATE TABLE dummy (\n\tpk STRING, \n\tobj_col OBJECT, '
+            ('\nCREATE TABLE dummy (\n\tpk STRING NOT NULL, \n\tobj_col OBJECT, '
              '\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__ = {
@@ -91,35 +91,35 @@ class DummyTable(self.Base):
         self.Base.metadata.create_all()
         fake_cursor.execute.assert_called_with(
             ('\nCREATE TABLE t (\n\t'
-             'pk STRING, \n\t'
+             'pk STRING NOT NULL, \n\t'
              'p STRING, \n\t'
              'PRIMARY KEY (pk)\n'
              ') 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)
             p = sa.Column(sa.BigInteger, sa.Computed("date_trunc('day', ts)"))
         self.Base.metadata.create_all()
         fake_cursor.execute.assert_called_with(
             ('\nCREATE TABLE t (\n\t'
-             'ts LONG, \n\t'
+             'ts LONG NOT NULL, \n\t'
              'p LONG GENERATED ALWAYS AS (date_trunc(\'day\', ts)), \n\t'
              'PRIMARY KEY (ts)\n'
              ')\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__ = {
@@ -131,13 +131,13 @@ class DummyTable(self.Base):
         self.Base.metadata.create_all()
         fake_cursor.execute.assert_called_with(
             ('\nCREATE TABLE t (\n\t'
-             'pk STRING, \n\t'
+             'pk STRING NOT NULL, \n\t'
              'p STRING, \n\t'
              'PRIMARY KEY (pk)\n'
              ') 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__ = {
@@ -149,12 +149,12 @@ class DummyTable(self.Base):
         self.Base.metadata.create_all()
         fake_cursor.execute.assert_called_with(
             ('\nCREATE TABLE t (\n\t'
-             'pk STRING, \n\t'
+             'pk STRING NOT NULL, \n\t'
              'PRIMARY KEY (pk)\n'
              ') 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__ = {
@@ -166,13 +166,13 @@ class DummyTable(self.Base):
         self.Base.metadata.create_all()
         fake_cursor.execute.assert_called_with(
             ('\nCREATE TABLE t (\n\t'
-             'pk STRING, \n\t'
-             'p STRING, \n\t'
+             'pk STRING NOT NULL, \n\t'
+             'p STRING NOT NULL, \n\t'
              'PRIMARY KEY (pk, p)\n'
              ') 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)
@@ -181,6 +181,51 @@ class DummyTable(self.Base):
         self.Base.metadata.create_all()
         fake_cursor.execute.assert_called_with(
             ('\nCREATE TABLE t (\n\t'
-             'pk STRING, \n\t'
+             'pk STRING NOT NULL, \n\t'
              'tags ARRAY(OBJECT), \n\t'
              'PRIMARY KEY (pk)\n)\n\n'), ())
+
+    def test_column_nullable(self):
+        class DummyTable(self.Base):
+            __tablename__ = 't'
+            pk = sa.Column(sa.String, primary_key=True)
+            a = sa.Column(sa.Integer, nullable=True)
+            b = sa.Column(sa.Integer, nullable=False)
+
+        self.Base.metadata.create_all()
+        fake_cursor.execute.assert_called_with(
+            ('\nCREATE TABLE t (\n\t'
+             'pk STRING NOT NULL, \n\t'
+             'a INT, \n\t'
+             'b INT NOT NULL, \n\t'
+             'PRIMARY KEY (pk)\n)\n\n'), ())
+
+    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()