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 b1e4e37
Show file tree
Hide file tree
Showing 4 changed files with 91 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
44 changes: 44 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,41 @@ 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

# Heuristically consider if a dot is included, that it's a full-qualified
# identifier like <schema>.<table>. It needs to be split, in order to apply
# identifier quoting properly.
if "." in ident:
parts = ident.split(".")
if len(parts) > 2:
raise ValueError(f"Invalid relation name {ident}")
return (
identifier_preparer.quote_schema(parts[0]) + "." + identifier_preparer.quote(parts[1])
)
return identifier_preparer.quote(ident=ident)
44 changes: 44 additions & 0 deletions tests/test_support_util.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,44 @@
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.
"""
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"'
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"'


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("my-db.my-schema.my-table")

0 comments on commit b1e4e37

Please sign in to comment.