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 `object` properly, for example
when used as column names.
  • Loading branch information
amotl committed Jun 22, 2024
1 parent 877ebaa commit 27f3e35
Show file tree
Hide file tree
Showing 4 changed files with 67 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
34 changes: 34 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,36 @@ 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 ''


# TODO: There are certainly more to add here than just `object`?
CRATEDB_RESERVED_WORDS = ["object"]


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 341 in src/sqlalchemy_cratedb/compiler.py

View check run for this annotation

Codecov / codecov/patch

src/sqlalchemy_cratedb/compiler.py#L340-L341

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

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

View check run for this annotation

Codecov / codecov/patch

src/sqlalchemy_cratedb/compiler.py#L344

Added line #L344 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 348 in src/sqlalchemy_cratedb/compiler.py

View check run for this annotation

Codecov / codecov/patch

src/sqlalchemy_cratedb/compiler.py#L347-L348

Added lines #L347 - L348 were not covered by tests

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

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

View check run for this annotation

Codecov / codecov/patch

src/sqlalchemy_cratedb/compiler.py#L350-L351

Added lines #L350 - L351 were not covered by tests

if (

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

View check run for this annotation

Codecov / codecov/patch

src/sqlalchemy_cratedb/compiler.py#L353

Added line #L353 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 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
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"

id = 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 (
\tid INT NOT NULL,
\t"array" STRING,
\t"object" STRING,
\tPRIMARY KEY (id)
)
""")) # noqa: W291, W293

0 comments on commit 27f3e35

Please sign in to comment.