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 7970219
Show file tree
Hide file tree
Showing 4 changed files with 75 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
42 changes: 42 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,44 @@ 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 = \
"cross, current_date, intersect, else, end, except, using, case, and, current_schema, any, " \
"all, set, limit, input, natural, cast, directory, is, when, if, table, right, outer, full, " \
"order, select, join, add, session_user, current_time, grant, true, left, into, try_cast, " \
"current_role, insert, some, exists, update, false, create, reset, offset, object, " \
"transient, current_user, in, or, for, alter, asc, function, null, from, default, not, " \
"like, union, distinct, nulls, having, inner, by, persistent, stratify, array, revoke, " \
"match, drop, escape, where, costs, with, group, index, delete, column, on, unbounded, " \
"returns, then, last, user, called, recursive, between, describe, as, extract, " \
"current_timestamp, deny, first, constraint, desc".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(

Check warning on line 349 in src/sqlalchemy_cratedb/compiler.py

View check run for this annotation

Codecov / codecov/patch

src/sqlalchemy_cratedb/compiler.py#L348-L349

Added lines #L348 - L349 were not covered by tests
self.escape_to_quote, self.escape_quote
)
return value

Check warning on line 352 in src/sqlalchemy_cratedb/compiler.py

View check run for this annotation

Codecov / codecov/patch

src/sqlalchemy_cratedb/compiler.py#L352

Added line #L352 was not covered by tests

def format_type(self, type_, use_schema=True):
if not type_.name:
raise sa.exc.CompileError("Type requires a name.")

Check warning on line 356 in src/sqlalchemy_cratedb/compiler.py

View check run for this annotation

Codecov / codecov/patch

src/sqlalchemy_cratedb/compiler.py#L355-L356

Added lines #L355 - L356 were not covered by tests

name = self.quote(type_.name)
effective_schema = self.schema_for_object(type_)

Check warning on line 359 in src/sqlalchemy_cratedb/compiler.py

View check run for this annotation

Codecov / codecov/patch

src/sqlalchemy_cratedb/compiler.py#L358-L359

Added lines #L358 - L359 were not covered by tests

if (

Check warning on line 361 in src/sqlalchemy_cratedb/compiler.py

View check run for this annotation

Codecov / codecov/patch

src/sqlalchemy_cratedb/compiler.py#L361

Added line #L361 was not covered by tests
not self.omit_schema
and use_schema
and effective_schema is not None
):
name = self.quote_schema(effective_schema) + "." + name
return name

Check warning on line 367 in src/sqlalchemy_cratedb/compiler.py

View check run for this annotation

Codecov / codecov/patch

src/sqlalchemy_cratedb/compiler.py#L366-L367

Added lines #L366 - L367 were not covered by tests
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 7970219

Please sign in to comment.