diff --git a/arbalister/adbc.py b/arbalister/adbc.py index 2e0b546..596e4f7 100644 --- a/arbalister/adbc.py +++ b/arbalister/adbc.py @@ -5,6 +5,8 @@ import adbc_driver_sqlite.dbapi as adbc_sqlite import pyarrow as pa +from . import utils as utils + def write_sqlite( table: pa.Table, @@ -57,7 +59,7 @@ def read_sqlite(cls, context: Any, path: pathlib.Path | str, table_name: str | N with adbc_sqlite.connect(str(path)) as connection: with connection.cursor() as cursor: - cursor.execute(f'SELECT COUNT(*) FROM "{table_name}"') + cursor.execute(f"SELECT COUNT(*) FROM {utils.escape(table_name)}") num_rows = cursor.fetchone()[0] # type: ignore[index] schema = connection.adbc_get_table_schema(table_name) @@ -88,13 +90,9 @@ def to_arrow_table(self) -> pa.Table: else "" ) - # Escape column names or default to wildcard - if self._select is not None: - columns = ",".join(f'"{c}"' for c in self._select) - else: - columns = "*" + columns = ",".join(self._select) if self._select is not None else "*" with adbc_sqlite.connect(self._path) as connection: with connection.cursor() as cursor: - cursor.execute(f'SELECT {columns} FROM "{self._table_name}" {limit}') + cursor.execute(f"SELECT {columns} FROM {utils.escape(self._table_name)} {limit}") return cursor.fetch_arrow_table() diff --git a/arbalister/routes.py b/arbalister/routes.py index b867d82..7071ab0 100644 --- a/arbalister/routes.py +++ b/arbalister/routes.py @@ -14,6 +14,7 @@ from . import arrow as abw from . import file_format as ff from . import params as params +from . import utils as utils @dataclasses.dataclass(frozen=True, slots=True) @@ -104,8 +105,8 @@ async def get(self, path: str) -> None: df = df.limit(count=count, offset=offset) if params.start_col is not None and params.end_col is not None: - col_names = df.schema().names - df = df.select(*col_names[params.start_col : params.end_col]) + col_names = df.schema().names[params.start_col : params.end_col] + df = df.select(*(utils.escape(c) for c in col_names)) table: pa.Table = df.to_arrow_table() @@ -158,7 +159,7 @@ async def get(self, path: str) -> None: # No dedicated exception type coming from DataFusion if str(e).startswith("DataFusion"): first_col: str = schema.names[0] - batches = df.aggregate([], [dnf.count(dn.col(first_col))]).collect() + batches = df.aggregate([], [dnf.count(dn.col(utils.escape(first_col)))]).collect() num_rows = batches[0].column(0)[0].as_py() # Create a zero-row IPC stream with the table schema diff --git a/arbalister/tests/test_routes.py b/arbalister/tests/test_routes.py index 5d96a12..ae584e1 100644 --- a/arbalister/tests/test_routes.py +++ b/arbalister/tests/test_routes.py @@ -59,9 +59,25 @@ def dummy_table_1(num_rows: int = 10) -> pa.Table: """Generate a table with fake data.""" data = { "lower": random.choices(string.ascii_lowercase, k=num_rows), - "sequence": list(range(num_rows)), - "upper": random.choices(string.ascii_uppercase, k=num_rows), - "number": [random.random() for _ in range(num_rows)], + "some sequence": list(range(num_rows)), + "UPPER": random.choices(string.ascii_uppercase, k=num_rows), + "number 💯": [random.random() for _ in range(num_rows)], + } + table = pa.table(data) + return table + + +@pytest.fixture(scope="module") +def avro_table(num_rows: int = 10) -> pa.Table: + """Generate a table with fake data and Avro compatible column names. + + Avro field names are restricted to ``[A-Za-z_][A-Za-z0-9_]*``. + """ + data = { + "lower": random.choices(string.ascii_lowercase, k=num_rows), + "some_sequence": list(range(num_rows)), + "UPPER": random.choices(string.ascii_uppercase, k=num_rows), + "number_": [random.random() for _ in range(num_rows)], } table = pa.table(data) return table @@ -82,9 +98,17 @@ def dummy_table_2(num_rows: int = 13) -> pa.Table: @pytest.fixture(scope="module") -def full_table(file_params: ff.FileFormat, dummy_table_1: pa.Table, dummy_table_2: pa.Table) -> pa.Table: +def full_table( + file_format: ff.FileFormat, + file_params: arb.routes.FileReadOptions, + dummy_table_1: pa.Table, + dummy_table_2: pa.Table, + avro_table: pa.Table, +) -> pa.Table: """Return the full table on which we are executed queries.""" - if isinstance(file_params, arb.routes.SqliteReadOptions): + if file_format == ff.FileFormat.Avro: + return avro_table + if isinstance(file_params, arb.routes.SqliteReadOptions) and file_params.table_name: return { "dummy_table_1": dummy_table_1, "dummy_table_2": dummy_table_2, @@ -97,6 +121,7 @@ def table_file( jp_root_dir: pathlib.Path, dummy_table_1: pa.Table, dummy_table_2: pa.Table, + avro_table: pa.Table, file_format: ff.FileFormat, file_params: arb.routes.FileReadOptions, ) -> pathlib.Path: @@ -105,6 +130,8 @@ def table_file( table_path = jp_root_dir / f"test.{str(file_format).lower()}" match file_format: + case ff.FileFormat.Avro: + write_table(avro_table, table_path) case ff.FileFormat.Csv: write_table(dummy_table_1, table_path, delimiter=getattr(file_params, "delimiter", ",")) case ff.FileFormat.Sqlite: diff --git a/arbalister/utils.py b/arbalister/utils.py new file mode 100644 index 0000000..2a7e658 --- /dev/null +++ b/arbalister/utils.py @@ -0,0 +1,7 @@ +def escape(name: str) -> str: + """Quote a SQL identifier so that it is taken verbatim. + + Unquoted identifiers are lowercased and dots are read as qualifiers. + """ + escaped = name.replace('"', '""') + return f'"{escaped}"' diff --git a/data/generate.py b/data/generate.py index dea46cd..163a5c9 100644 --- a/data/generate.py +++ b/data/generate.py @@ -1,6 +1,7 @@ import argparse import pathlib import random +import re import datafusion as dn import datafusion.functions as dnf @@ -33,10 +34,10 @@ def generate_table(num_rows: int) -> pa.Table: gen = faker.Faker() data = { - "name": [gen.name() for _ in range(num_rows)], - "address": [gen.address().replace("\n", ", ") for _ in range(num_rows)], - "age": [gen.random_number(digits=2) for _ in range(num_rows)], - "id": [gen.uuid4() for _ in range(num_rows)], + "Name": [gen.name() for _ in range(num_rows)], + "Address 🏠": [gen.address().replace("\n", ", ") for _ in range(num_rows)], + '"Age"': [gen.random_number(digits=2) for _ in range(num_rows)], + "__id": [gen.uuid4() for _ in range(num_rows)], } return pa.table(data) @@ -150,8 +151,16 @@ def shuffle_table(table: pa.Table, seed: int | None = None) -> pa.Table: return table.select(col_order).take(row_indices) +def avro_compatible(table: pa.Table) -> pa.Table: + """Rename the columns to comply with Avro field names ``[A-Za-z_][A-Za-z0-9_]*``.""" + names = [re.sub(r"[^A-Za-z0-9_]", "_", name) for name in table.column_names] + return table.rename_columns([n if re.match(r"[A-Za-z_]", n) else f"_{n}" for n in names]) + + def save_table(table: pa.Table, path: pathlib.Path, file_type: ff.FileFormat) -> None: """Save a table to file with the given file type.""" + if file_type == ff.FileFormat.Avro: + table = avro_compatible(table) path.parent.mkdir(exist_ok=True, parents=True) write_table = aa.get_table_writer(file_type) write_table(table, str(path))