Skip to content

Commit

Permalink
Core: Improve TableAddress.fullname
Browse files Browse the repository at this point in the history
By making the canonical `DatabaseAdapter.quote_relation_name` static, 
it can be used by `TableAddress.fullname` now.

This makes table name quoting less fragmented across the code base.
  • Loading branch information
amotl committed Jul 20, 2024
1 parent 731e007 commit 155fe10
Show file tree
Hide file tree
Showing 3 changed files with 48 additions and 25 deletions.
12 changes: 5 additions & 7 deletions cratedb_toolkit/model.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@

from boltons.urlutils import URL

from cratedb_toolkit.util.database import decode_database_table
from cratedb_toolkit.util.database import DatabaseAdapter, decode_database_table


@dataclasses.dataclass
Expand Down Expand Up @@ -85,12 +85,10 @@ class TableAddress:

@property
def fullname(self):
if self.schema is None and self.table is None:
raise ValueError("Uninitialized table address can not be serialized")
if self.schema and self.table:
return f'"{self.schema}"."{self.table}"'
else:
return f'"{self.table}"'
"""
Return a full-qualified quoted table identifier.
"""
return DatabaseAdapter.quote_relation_name(f"{self.schema}.{self.table}")


@dataclasses.dataclass
Expand Down
45 changes: 28 additions & 17 deletions cratedb_toolkit/util/database.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
from cratedb_sqlparse import sqlparse as sqlparse_cratedb
from sqlalchemy.exc import ProgrammingError
from sqlalchemy.sql.elements import AsBoolean
from sqlalchemy_cratedb.dialect import CrateDialect

from cratedb_toolkit.util.data import str_contains

Expand All @@ -24,6 +25,10 @@ def run_sql(dburi: str, sql: str, records: bool = False):
return DatabaseAdapter(dburi=dburi).run_sql(sql=sql, records=records)


# Just an instance of the dialect used for quoting purposes.
dialect = CrateDialect()


class DatabaseAdapter:
"""
Wrap SQLAlchemy connection to database.
Expand All @@ -35,37 +40,43 @@ def __init__(self, dburi: str, echo: bool = False):
# TODO: Make that go away.
self.connection = self.engine.connect()

def quote_relation_name(self, ident: str) -> str:
@staticmethod
def quote_relation_name(ident: str) -> str:
"""
Quote the given, possibly full-qualified, relation name if needed.
Quote a simple or full-qualified table/relation name, when needed.
In: foo
Out: foo
Simple: <table>
Full-qualified: <schema>.<table>
In: Foo
Out: "Foo"
Happy path examples:
In: "Foo"
Out: "Foo"
foo => foo
Foo => "Foo"
"Foo" => "Foo"
foo.bar => foo.bar
foo-bar.baz_qux => "foo-bar".baz_qux
In: foo.bar
Out: "foo"."bar"
Such input strings will not be modified:
In: "foo.bar"
Out: "foo.bar"
"foo.bar" => "foo.bar"
"""
if ident[0] == '"' and ident[len(ident) - 1] == '"':

# 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 (
self.engine.dialect.identifier_preparer.quote_schema(parts[0])
+ "."
+ self.engine.dialect.identifier_preparer.quote(parts[1])
dialect.identifier_preparer.quote_schema(parts[0]) + "." + dialect.identifier_preparer.quote(parts[1])
)
return self.engine.dialect.identifier_preparer.quote(ident=ident)
return dialect.identifier_preparer.quote(ident=ident)

def run_sql(
self,
Expand Down
16 changes: 15 additions & 1 deletion tests/util/database.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,10 @@
from cratedb_toolkit.util import DatabaseAdapter


def test_quote_relation_name():
def test_quote_relation_name_single():
"""
Verify quoting a simple or full-qualified relation name.
"""
database = DatabaseAdapter(dburi="crate://localhost")
assert database.quote_relation_name("my_table") == "my_table"
assert database.quote_relation_name("my-table") == '"my-table"'
Expand All @@ -17,6 +20,17 @@ def test_quote_relation_name():
assert database.quote_relation_name("table") == '"table"'


def test_quote_relation_name_double():
"""
Verify quoting a relation name twice does not cause any harm.
"""
database = DatabaseAdapter(dburi="crate://localhost")
input_fqn = "foo-bar.baz_qux"
output_fqn = '"foo-bar".baz_qux'
assert database.quote_relation_name(input_fqn) == output_fqn
assert database.quote_relation_name(output_fqn) == output_fqn


def test_quote_relation_name_with_invalid_fqn():
database = DatabaseAdapter(dburi="crate://localhost")
with pytest.raises(ValueError):
Expand Down

0 comments on commit 155fe10

Please sign in to comment.