fix: Table column name escape

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()

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

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:

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}"'

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))