Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Core: Improve quoting for TableAddress.fullname #202

Merged
merged 1 commit into from
Jul 22, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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
20 changes: 17 additions & 3 deletions tests/util/database.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,8 +3,11 @@
from cratedb_toolkit.util import DatabaseAdapter


def test_quote_relation_name():
database = DatabaseAdapter(dburi="crate://localhost")
def test_quote_relation_name_single():
"""
Verify quoting a simple or full-qualified relation name.
"""
database = DatabaseAdapter
assert database.quote_relation_name("my_table") == "my_table"
assert database.quote_relation_name("my-table") == '"my-table"'
assert database.quote_relation_name("MyTable") == '"MyTable"'
Expand All @@ -17,7 +20,18 @@ 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
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")
database = DatabaseAdapter
with pytest.raises(ValueError):
database.quote_relation_name("my-db.my-schema.my-table")