Compiler: Add CrateIdentifierPreparer for properly quoting reserved words

CHANGES.md

@@ -3,6 +3,8 @@
 
 ## Unreleased
 - Added/reactivated documentation as `sqlalchemy-cratedb`
+- Added `CrateIdentifierPreparer`, in order to quote reserved words
+  like `object` properly, for example when used as column names.
 
 ## 2024/06/13 0.37.0
 - Added support for CrateDB's [FLOAT_VECTOR] data type and its accompanying

src/sqlalchemy_cratedb/compiler.py

@@ -25,6 +25,7 @@
 
 import sqlalchemy as sa
 from sqlalchemy.dialects.postgresql.base import PGCompiler
+from sqlalchemy.dialects.postgresql.base import RESERVED_WORDS as POSTGRESQL_RESERVED_WORDS
 from sqlalchemy.sql import compiler
 from sqlalchemy.types import String
 from .type.geo import Geopoint, Geoshape
@@ -323,3 +324,39 @@
         warnings.warn("CrateDB does not support the 'INSERT ... FOR UPDATE' clause, "
                       "it will be omitted when generating SQL statements.")
         return ''
+
+
+CRATEDB_RESERVED_WORDS = \
+    "add, alter, between, by, called, costs, delete, deny, directory, drop, escape, exists, " \
+    "extract, first, function, if, index, input, insert, last, match, nulls, object, " \
+    "persistent, recursive, reset, returns, revoke, set, stratify, transient, try_cast, " \
+    "unbounded, update".split(", ")
+
+
+class CrateIdentifierPreparer(sa.sql.compiler.IdentifierPreparer):
+    """
+    Define CrateDB's reserved words to be quoted properly.
+    """
+    reserved_words = set(list(POSTGRESQL_RESERVED_WORDS) + CRATEDB_RESERVED_WORDS)
+
+    def _unquote_identifier(self, value):
+        if value[0] == self.initial_quote:
+            value = value[1:-1].replace(
+                self.escape_to_quote, self.escape_quote
+            )
+        return value
+
+    def format_type(self, type_, use_schema=True):
+        if not type_.name:
+            raise sa.exc.CompileError("Type requires a name.")
+
+        name = self.quote(type_.name)
+        effective_schema = self.schema_for_object(type_)
+
+        if (
+            not self.omit_schema
+            and use_schema
+            and effective_schema is not None
+        ):
+            name = self.quote_schema(effective_schema) + "." + name
+        return name

src/sqlalchemy_cratedb/dialect.py

@@ -29,7 +29,8 @@
 
 from .compiler import (
     CrateTypeCompiler,
-    CrateDDLCompiler
+    CrateDDLCompiler,
+    CrateIdentifierPreparer,
 )
 from crate.client.exceptions import TimezoneUnawareException
 from .sa_version import SA_VERSION, SA_1_4, SA_2_0
@@ -174,6 +175,7 @@ class CrateDialect(default.DefaultDialect):
     statement_compiler = statement_compiler
     ddl_compiler = CrateDDLCompiler
     type_compiler = CrateTypeCompiler
+    preparer = CrateIdentifierPreparer
     use_insertmanyvalues = True
     use_insertmanyvalues_wo_returning = True
     supports_multivalues_insert = True

tests/compiler_test.py

@@ -432,3 +432,31 @@ class FooBar(Base):
         self.assertIsSubclass(w[-1].category, UserWarning)
         self.assertIn("CrateDB does not support unique constraints, "
                       "they will be omitted when generating DDL statements.", str(w[-1].message))
+
+    def test_ddl_with_reserved_words(self):
+        """
+        Verify CrateDB's reserved words like `object` are quoted properly.
+        """
+
+        Base = declarative_base(metadata=self.metadata)
+
+        class FooBar(Base):
+            """The entity."""
+
+            __tablename__ = "foobar"
+
+            index = sa.Column(sa.Integer, primary_key=True)
+            array = sa.Column(sa.String)
+            object = sa.Column(sa.String)
+
+        # Verify SQL DDL statement.
+        self.metadata.create_all(self.engine, tables=[FooBar.__table__], checkfirst=False)
+        self.assertEqual(self.executed_statement, dedent("""
+            CREATE TABLE testdrive.foobar (
+            \t"index" INT NOT NULL, 
+            \t"array" STRING, 
+            \t"object" STRING, 
+            \tPRIMARY KEY ("index")
+            )
+
+        """))  # noqa: W291, W293