allow columns with quotes

Author: lhoestq

libs/libapi/src/libapi/duckdb.py

@@ -13,7 +13,6 @@
 from typing import Any, Optional, TypeVar
 
 import anyio
-import duckdb
 import filelock
 from datasets import Features
 from filelock import AsyncFileLock
@@ -27,8 +26,11 @@
     DUCKDB_DEFAULT_PARTIAL_INDEX_FILENAME,
     compute_transformed_data,
     create_index,
+    duckdb_connect,
     get_indexable_columns,
     get_monolingual_stemmer,
+    key_sql,
+    varchar_sql,
 )
 from libcommon.parquet_utils import (
     ParquetFileMetadataItem,
@@ -110,7 +112,7 @@ def build_index_file(
     split_parquet_files = split_parquet_files[:num_parquet_files_to_index]
     parquet_file_names = [parquet_file["filename"] for parquet_file in split_parquet_files]
 
-    column_names_sql = ",".join(f'"{column}"' for column in features)
+    column_names_sql = ",".join(key_sql(column) for column in features)
 
     # look for indexable columns (= possibly nested columns containing string data)
     indexable_columns = get_indexable_columns(Features.from_dict(features))
@@ -142,18 +144,18 @@ def build_index_file(
     Path(index_file_location).parent.mkdir(exist_ok=True, parents=True)
 
     try:
-        with duckdb.connect(index_file_location) as con:
+        with duckdb_connect(index_file_location=index_file_location, extensions_directory=extensions_directory) as con:
             if transformed_df is not None:
                 logging.debug(transformed_df.head())
                 # update original data with results of transformations (string lengths, audio durations, etc.):
                 logging.info(f"Updating data with {transformed_df.columns}")
                 create_command_sql = CREATE_TABLE_JOIN_WITH_TRANSFORMED_DATA_COMMAND_FROM_LIST_OF_PARQUET_FILES.format(
-                    columns=column_names_sql, source=[str(p) for p in all_split_parquets]
+                    columns=column_names_sql, source="[" + ",".join(varchar_sql(str(p)) for p in all_split_parquets) + "]"
                 )
 
             else:
                 create_command_sql = CREATE_TABLE_COMMAND_FROM_LIST_OF_PARQUET_FILES.format(
-                    columns=column_names_sql, source=[str(p) for p in all_split_parquets]
+                    columns=column_names_sql, source="[" + ",".join(varchar_sql(str(p)) for p in all_split_parquets) + "]"
                 )
 
             logging.info(create_command_sql)

libs/libcommon/src/libcommon/duckdb_utils.py

@@ -34,6 +34,14 @@
 DEFAULT_STEMMER = "none"  # Exact word matches
 DUCKDB_DEFAULT_INDEX_FILENAME = "index.duckdb"
 DUCKDB_DEFAULT_PARTIAL_INDEX_FILENAME = "partial-index.duckdb"
+
+ATTACH_DATABASE = "ATTACH '{database}' as db; USE db;"
+ATTACH_READ_ONLY_DATABASE = "ATTACH '{database}' as db (READ_ONLY); USE db;"
+LOAD_FTS_COMMAND = "INSTALL 'fts'; LOAD 'fts';"
+DISABLE_EXTERNAL_ACCESS_COMMAND = "SET enable_external_access=false;"
+LOCK_CONFIG_COMMAND = "SET lock_configuration=true;"
+SET_EXTENSIONS_DIRECTORY_COMMAND = "SET extension_directory='{directory}';"
+
 CREATE_INDEX_COMMAND = (
     f"PRAGMA create_fts_index('data', '{ROW_IDX_COLUMN}', {{columns}}, stemmer='{{stemmer}}', overwrite=1);"
 )
@@ -216,16 +224,7 @@ def _sql(con: duckdb.DuckDBPyConnection, query: str) -> duckdb.DuckDBPyRelation:
         return out
 
     with tempfile.TemporaryDirectory(suffix=".duckdb") as tmp_dir:
-        with duckdb.connect(":memory:") as con:
-            # configure duckdb extensions
-            if extensions_directory is not None:
-                con.execute(SET_EXTENSIONS_DIRECTORY_COMMAND.format(directory=extensions_directory))
-            con.execute(INSTALL_AND_LOAD_EXTENSION_COMMAND)
-
-            # init
-            _sql(con, "ATTACH '%database%' as db;")
-            _sql(con, "USE db;")
-
+        with duckdb_connect(database=database, extensions_directory=extensions_directory) as con:
             # check input_table and get number of rows
             _count = _sql(con, "SELECT count(*) FROM %input_table%;").fetchone()
             if _count and isinstance(_count[0], int):
@@ -255,7 +254,8 @@ def _sql(con: duckdb.DuckDBPyConnection, query: str) -> duckdb.DuckDBPyRelation:
             )
 
             # create fields table
-            field_values = ", ".join(f"({i}, '{field}')" for i, field in enumerate(columns))
+
+            field_values = ", ".join(f"({i}, {varchar_sql(field)})" for i, field in enumerate(columns))
             _sql(
                 con,
                 """
@@ -278,17 +278,8 @@ def _sql(con: duckdb.DuckDBPyConnection, query: str) -> duckdb.DuckDBPyRelation:
         batch_size = 1 + count // num_jobs
         commands = [
             (
-                (
-                    SET_EXTENSIONS_DIRECTORY_COMMAND.format(directory=extensions_directory)
-                    if extensions_directory is not None
-                    else ""
-                )
-                + INSTALL_AND_LOAD_EXTENSION_COMMAND
-                + (
-                    "ATTACH IF NOT EXISTS '%database%' as db (READ_ONLY);"  # nosec - tmp_dir, batch_size, rank and i are safe
-                    "USE db;"
-                    f"ATTACH '{tmp_dir}/tmp_{rank}_{i}.duckdb' as tmp_{rank}_{i};"
-                    f"""
+                f"ATTACH '{tmp_dir}/tmp_{rank}_{i}.duckdb' as tmp_{rank}_{i};"  # nosec - tmp_dir, batch_size, rank and i are safe
+                f"""
                     CREATE TABLE tmp_{rank}_{i}.tokenized AS (
                         SELECT unnest(%fts_schema%.tokenize(fts_ii."{column}")) AS w,
                             {rank * batch_size} + row_number() OVER () - 1 AS docid,
@@ -299,14 +290,13 @@ def _sql(con: duckdb.DuckDBPyConnection, query: str) -> duckdb.DuckDBPyRelation:
                     );
                     CHECKPOINT;
                     """
-                )
             )
             for rank in range(num_jobs)
             for i, column in enumerate(columns)
         ]
 
         def _parallel_sql(command: str) -> None:
-            with duckdb.connect(":memory:") as rank_con:
+            with duckdb_connect(database=database, extensions_directory=extensions_directory) as rank_con:
                 _sql(rank_con, command)
 
         thread_map(_parallel_sql, commands, desc="Tokenize")
@@ -324,16 +314,9 @@ def _parallel_sql(command: str) -> None:
         #     """)
         # union_fields_query = " UNION ALL ".join(f"SELECT * FROM tmp.tokenized_{i}" for i in range(len(columns)))
 
-        with duckdb.connect(":memory:") as con:
-            # configure duckdb extensions
-            if extensions_directory is not None:
-                con.execute(SET_EXTENSIONS_DIRECTORY_COMMAND.format(directory=extensions_directory))
-            con.execute(INSTALL_AND_LOAD_EXTENSION_COMMAND)
-
+        with duckdb_connect(database=database, extensions_directory=extensions_directory) as con:
             # init
             _sql(con, f"ATTACH '{tmp_dir}/tmp.duckdb' as tmp;")  # nosec - tmp_dir is safe
-            _sql(con, "ATTACH '%database%' as db;")
-            _sql(con, "USE db;")
             _sql(
                 con,
                 ";".join(
@@ -526,3 +509,32 @@ def _parallel_sql(command: str) -> None:
             """,
             )
             _sql(con, "CHECKPOINT;")
+
+
+def varchar_sql(value: str) -> str:
+    """escape the value and return the varchar `'{value}'`"""
+    return "'" + value.replace("'", "''") + "'"
+
+
+def key_sql(value: str) -> str:
+    """escape the value and return the key `"{value}"`"""
+    return '"' + value.replace('"', '""') + '"'
+
+
+def duckdb_connect(
+    index_file_location: Optional[str] = None,
+    database: Optional[str] = None,
+    extensions_directory: Optional[str] = None,
+    read_only: bool = False,
+    **kwargs: Any,
+) -> duckdb.DuckDBPyConnection:
+    """In-memory session with the current database attached with read-only and fts extension"""
+    con = duckdb.connect(":memory:" if index_file_location is None else index_file_location, **kwargs)
+    if database is not None:
+        con.execute((ATTACH_READ_ONLY_DATABASE if read_only else ATTACH_DATABASE).format(database=database))
+    if extensions_directory is not None:
+        con.execute(SET_EXTENSIONS_DIRECTORY_COMMAND.format(directory=extensions_directory))
+    con.sql(LOAD_FTS_COMMAND)
+    con.sql(DISABLE_EXTERNAL_ACCESS_COMMAND)
+    con.sql(LOCK_CONFIG_COMMAND)
+    return con

services/search/src/search/duckdb_connection.py

@@ -30,14 +30,13 @@
     get_json_ok_response,
 )
 from libcommon.constants import ROW_IDX_COLUMN
+from libcommon.duckdb_utils import duckdb_connect, key_sql
 from libcommon.prometheus import StepProfiler
 from libcommon.storage import StrPath, clean_dir
 from libcommon.storage_client import StorageClient
 from starlette.requests import Request
 from starlette.responses import Response
 
-from search.duckdb_connection import duckdb_connect_readonly
-
 FILTER_QUERY = """\
     SELECT {columns}
     FROM data
@@ -224,9 +223,11 @@ def execute_filter_query(
     offset: int,
     extensions_directory: Optional[str] = None,
 ) -> tuple[int, pa.Table]:
-    with duckdb_connect_readonly(extensions_directory=extensions_directory, database=index_file_location) as con:
+    with duckdb_connect(
+        index_file_location=index_file_location, extensions_directory=extensions_directory, read_only=True
+    ) as con:
         filter_query = FILTER_QUERY.format(
-            columns=",".join([f'"{column}"' for column in columns]),
+            columns=",".join([key_sql(column) for column in columns]),
             where=f"WHERE {where}" if where else "",
             orderby=f"ORDER BY {orderby}" if orderby else "",
             limit=limit,

services/search/src/search/routes/filter.py

@@ -32,15 +32,14 @@
 )
 from libcommon.constants import HF_FTS_SCORE, MAX_NUM_ROWS_PER_PAGE, ROW_IDX_COLUMN
 from libcommon.dtos import PaginatedResponse
+from libcommon.duckdb_utils import duckdb_connect, key_sql
 from libcommon.prometheus import StepProfiler
 from libcommon.storage import StrPath, clean_dir
 from libcommon.storage_client import StorageClient
 from libcommon.viewer_utils.features import to_features_list
 from starlette.requests import Request
 from starlette.responses import Response
 
-from search.duckdb_connection import duckdb_connect_readonly
-
 logger = logging.getLogger(__name__)
 
 FTS_STAGE_TABLE_COMMAND = f"SELECT * FROM (SELECT {ROW_IDX_COLUMN}, fts_main_data.match_bm25({ROW_IDX_COLUMN}, ?) AS {HF_FTS_SCORE} FROM data) A WHERE {HF_FTS_SCORE} IS NOT NULL;"  # nosec
@@ -57,13 +56,15 @@ def full_text_search(
     length: int,
     extensions_directory: Optional[str] = None,
 ) -> tuple[int, pa.Table]:
-    with duckdb_connect_readonly(extensions_directory=extensions_directory, database=index_file_location) as con:
+    with duckdb_connect(
+        extensions_directory=extensions_directory, database=index_file_location, read_only=True
+    ) as con:
         fts_stage_table = con.execute(query=FTS_STAGE_TABLE_COMMAND, parameters=[query]).arrow().read_all()
         num_rows_total = fts_stage_table.num_rows
         logging.info(f"got {num_rows_total=} results for {query=} using {offset=} {length=}")
         fts_stage_table = fts_stage_table.sort_by([(HF_FTS_SCORE, "descending")]).slice(offset, length)
         join_stage_and_data_query = JOIN_STAGE_AND_DATA_COMMAND.format(
-            columns=",".join([f'"{column}"' for column in columns]),
+            columns=",".join([key_sql(column) for column in columns]),
             row_idx_column=ROW_IDX_COLUMN,
             hf_fts_score=HF_FTS_SCORE,
         )

services/search/src/search/routes/search.py

@@ -48,8 +48,6 @@ def ds() -> Dataset:
 def index_file_location(ds: Dataset) -> Generator[str, None, None]:
     index_file_location = "index.duckdb"
     con = duckdb.connect(index_file_location)
-    con.execute("INSTALL 'httpfs';")
-    con.execute("LOAD 'httpfs';")
     con.execute("INSTALL 'fts';")
     con.execute("LOAD 'fts';")
     con.sql("CREATE OR REPLACE SEQUENCE serial START 0 MINVALUE 0;")

services/search/tests/routes/test_filter.py

@@ -63,8 +63,6 @@ def test_full_text_search(
     # simulate index file
     index_file_location = "index.duckdb"
     con = duckdb.connect(index_file_location)
-    con.execute("INSTALL 'httpfs';")
-    con.execute("LOAD 'httpfs';")
     con.execute("INSTALL 'fts';")
     con.execute("LOAD 'fts';")
     con.sql("CREATE OR REPLACE SEQUENCE serial START 0 MINVALUE 0;")