Skip to content

Commit

Permalink
Add quote_relation_name support utility function
Browse files Browse the repository at this point in the history
  • Loading branch information
amotl committed Aug 27, 2024
1 parent 5e4368f commit fab6f4b
Show file tree
Hide file tree
Showing 4 changed files with 97 additions and 2 deletions.
2 changes: 1 addition & 1 deletion CHANGES.md
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
# Changelog


## Unreleased
- Added `quote_relation_name` support utility function

## 2024/06/25 0.38.0
- Added/reactivated documentation as `sqlalchemy-cratedb`
Expand Down
3 changes: 2 additions & 1 deletion src/sqlalchemy_cratedb/support/__init__.py
Original file line number Diff line number Diff line change
@@ -1,12 +1,13 @@
from sqlalchemy_cratedb.support.pandas import insert_bulk, table_kwargs
from sqlalchemy_cratedb.support.polyfill import check_uniqueness_factory, refresh_after_dml, \
patch_autoincrement_timestamp
from sqlalchemy_cratedb.support.util import refresh_table, refresh_dirty
from sqlalchemy_cratedb.support.util import quote_relation_name, refresh_table, refresh_dirty

__all__ = [
check_uniqueness_factory,
insert_bulk,
patch_autoincrement_timestamp,
quote_relation_name,
refresh_after_dml,
refresh_dirty,
refresh_table,
Expand Down
43 changes: 43 additions & 0 deletions src/sqlalchemy_cratedb/support/util.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,13 +3,19 @@

import sqlalchemy as sa

from sqlalchemy_cratedb.dialect import CrateDialect

if t.TYPE_CHECKING:
try:
from sqlalchemy.orm import DeclarativeBase
except ImportError:
pass


# An instance of the dialect used for quoting purposes.
identifier_preparer = CrateDialect().identifier_preparer


def refresh_table(connection, target: t.Union[str, "DeclarativeBase", "sa.sql.selectable.TableClause"]):
"""
Invoke a `REFRESH TABLE` statement.
Expand Down Expand Up @@ -39,3 +45,40 @@ def refresh_dirty(session, flush_context=None):
dirty_classes = {entity.__class__ for entity in dirty_entities}
for class_ in dirty_classes:
refresh_table(session, class_)


def quote_relation_name(ident: str) -> str:
"""
Quote a simple or full-qualified table/relation name, when needed.
Simple: <table>
Full-qualified: <schema>.<table>
Happy path examples:
foo => foo
Foo => "Foo"
"Foo" => "Foo"
foo.bar => foo.bar
foo-bar.baz_qux => "foo-bar".baz_qux
Such input strings will not be modified:
"foo.bar" => "foo.bar"
"""

# Heuristically consider that if a quote exists at the beginning or the end
# of the input string, that the relation name has been quoted already.
if ident.startswith('"') or ident.endswith('"'):
return ident

# When no dot is included, it's not a full-qualified identifier.
if "." not in ident:
return identifier_preparer.quote(ident=ident)

# If a dot is included, it's a full-qualified identifier like <schema>.<table>.
# It needs to be split, in order to apply identifier quoting properly.
parts = ident.split(".")
if len(parts) > 3:
raise ValueError(f"Invalid relation name, too many parts: {ident}")
return ".".join(map(identifier_preparer.quote, parts))
51 changes: 51 additions & 0 deletions tests/test_support_util.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,51 @@
import pytest

from sqlalchemy_cratedb.support import quote_relation_name


def test_quote_relation_name_once():
"""
Verify quoting a simple or full-qualified relation name.
"""

# Table name only.
assert quote_relation_name("my_table") == "my_table"
assert quote_relation_name("my-table") == '"my-table"'
assert quote_relation_name("MyTable") == '"MyTable"'
assert quote_relation_name('"MyTable"') == '"MyTable"'

# Schema and table name.
assert quote_relation_name("my_schema.my_table") == "my_schema.my_table"
assert quote_relation_name("my-schema.my_table") == '"my-schema".my_table'
assert quote_relation_name('"wrong-quoted-fqn.my_table"') == '"wrong-quoted-fqn.my_table"'
assert quote_relation_name('"my_schema"."my_table"') == '"my_schema"."my_table"'

# Catalog, schema, and table name.
assert quote_relation_name("crate.doc.t01") == "crate.doc.t01"


def test_quote_relation_name_twice():
"""
Verify quoting a relation name twice does not cause any harm.
"""
input_fqn = "foo-bar.baz_qux"
output_fqn = '"foo-bar".baz_qux'
assert quote_relation_name(input_fqn) == output_fqn
assert quote_relation_name(output_fqn) == output_fqn


def test_quote_relation_name_reserved_keywords():
"""
Verify quoting a simple relation name that is a reserved keyword.
"""
assert quote_relation_name("table") == '"table"'
assert quote_relation_name("true") == '"true"'
assert quote_relation_name("select") == '"select"'


def test_quote_relation_name_with_invalid_fqn():
"""
Verify quoting a relation name with an invalid fqn raises an error.
"""
with pytest.raises(ValueError):
quote_relation_name("too-many.my-db.my-schema.my-table")

0 comments on commit fab6f4b

Please sign in to comment.