From f4b4ab3781bdf766909a8b237a0fc2eb0876b447 Mon Sep 17 00:00:00 2001
From: Andreas Motl <andreas.motl@crate.io>
Date: Mon, 15 Jan 2024 23:53:50 +0100
Subject: [PATCH] Compiler: Add CrateIdentifierPreparer

By using this component of the SQLAlchemy dialect compiler, it can
define CrateDB's reserved words to be quoted properly when building
SQL statements.

This allows to quote reserved words like `index` or `object` properly,
for example when used as column names.
---
 CHANGES.md                         |  2 ++
 src/sqlalchemy_cratedb/compiler.py | 37 ++++++++++++++++++++++++++++++
 src/sqlalchemy_cratedb/dialect.py  |  4 +++-
 tests/compiler_test.py             | 28 ++++++++++++++++++++++
 4 files changed, 70 insertions(+), 1 deletion(-)

diff --git a/CHANGES.md b/CHANGES.md
index 3b65ced0..e731e52d 100644
--- a/CHANGES.md
+++ b/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
diff --git a/src/sqlalchemy_cratedb/compiler.py b/src/sqlalchemy_cratedb/compiler.py
index d3a66188..d9e3bb7d 100644
--- a/src/sqlalchemy_cratedb/compiler.py
+++ b/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 @@ def for_update_clause(self, select, **kw):
         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
diff --git a/src/sqlalchemy_cratedb/dialect.py b/src/sqlalchemy_cratedb/dialect.py
index 43af2fc4..c5cc2deb 100644
--- a/src/sqlalchemy_cratedb/dialect.py
+++ b/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
diff --git a/tests/compiler_test.py b/tests/compiler_test.py
index a40ebb0f..91f7d8d0 100644
--- a/tests/compiler_test.py
+++ b/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