Skip to content

Commit

Permalink
Compiler: Add CrateIdentifierPreparer
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
amotl committed Jun 24, 2024
1 parent 877ebaa commit 9a95be0
Show file tree
Hide file tree
Showing 4 changed files with 70 additions and 1 deletion.
2 changes: 2 additions & 0 deletions CHANGES.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
37 changes: 37 additions & 0 deletions src/sqlalchemy_cratedb/compiler.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
4 changes: 3 additions & 1 deletion src/sqlalchemy_cratedb/dialect.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
28 changes: 28 additions & 0 deletions tests/compiler_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

0 comments on commit 9a95be0

Please sign in to comment.