Security advisory GHSA-339r-cjv9-x78g:  LlamaIndex Retrievers Integration: DuckDBRetriever SQL Injection
Package: llama-index-retrievers-duckdb-retriever (pypi), vulnerable version: 0.3.0, fixed in: 0.4.0
CWE: CWE-89
Files changed by the fix:
llama-index-integrations/retrievers/llama-index-retrievers-duckdb-retriever/llama_index/retrievers/duckdb_retriever/base.py

Fix diff:
diff --git a/llama-index-integrations/retrievers/llama-index-retrievers-duckdb-retriever/llama_index/retrievers/duckdb_retriever/base.py b/llama-index-integrations/retrievers/llama-index-retrievers-duckdb-retriever/llama_index/retrievers/duckdb_retriever/base.py
index d4b50785e19..fa4d2e9b9b5 100644
--- a/llama-index-integrations/retrievers/llama-index-retrievers-duckdb-retriever/llama_index/retrievers/duckdb_retriever/base.py
+++ b/llama-index-integrations/retrievers/llama-index-retrievers-duckdb-retriever/llama_index/retrievers/duckdb_retriever/base.py
@@ -2,13 +2,13 @@
 import os
 from typing import List, Optional
 
+import duckdb
 from llama_index.core.base.base_retriever import BaseRetriever
 from llama_index.core.callbacks.base import CallbackManager
 from llama_index.core.constants import DEFAULT_SIMILARITY_TOP_K
-from llama_index.core.schema import TextNode, NodeWithScore, QueryBundle
+from llama_index.core.schema import NodeWithScore, QueryBundle, TextNode
 
 logger = logging.getLogger(__name__)
-import_err_msg = "`duckdb` package not found, please run `pip install duckdb`"
 
 
 class DuckDBLocalContext:
@@ -18,19 +18,13 @@ def __init__(self, database_path: str):
         self._home_dir = os.path.expanduser("~")
 
     def __enter__(self) -> "duckdb.DuckDBPyConnection":
-        try:
-            import duckdb
-        except ImportError:
-            raise ImportError(import_err_msg)
-
-        if not os.path.exists(os.path.dirname(self.database_path)):
+        if self.database_path != ":memory:" and not os.path.exists(
+            os.path.dirname(self.database_path)
+        ):
             raise ValueError(
                 f"Directory {os.path.dirname(self.database_path)} does not exist."
             )
 
-        # if not os.path.isfile(self.database_path):
-        #     raise ValueError(f"Database path {self.database_path} is not a valid file.")
-
         self._conn = duckdb.connect(self.database_path)
         self._conn.execute(f"SET home_directory='{self._home_dir}';")
 
@@ -40,8 +34,6 @@ def __enter__(self) -> "duckdb.DuckDBPyConnection":
         return self._conn
 
     def __exit__(self, exc_type, exc_val, exc_tb) -> None:
-        self._conn.close()
-
         if self._conn:
             self._conn.close()
 
@@ -49,9 +41,9 @@ def __exit__(self, exc_type, exc_val, exc_tb) -> None:
 class DuckDBRetriever(BaseRetriever):
     def __init__(
         self,
-        database_name: Optional[str] = ":memory:",
-        table_name: Optional[str] = "documents",
-        text_search_config: Optional[dict] = {
+        database_name: str = ":memory:",
+        table_name: str = "documents",
+        text_search_config: dict = {
             "stemmer": "english",
             "stopwords": "english",
             "ignore": r"(\\.|[^a-z])+",
@@ -59,9 +51,9 @@ def __init__(
             "lower": True,
             "overwrite": True,
         },
-        persist_dir: Optional[str] = "./storage",
-        node_id_column: Optional[str] = "node_id",
-        text_column: Optional[str] = "text",
+        persist_dir: str = "./storage",
+        node_id_column: str = "node_id",
+        text_column: str = "text",
         # TODO: Add more options for FTS index creation
         similarity_top_k: int = DEFAULT_SIMILARITY_TOP_K,
         callback_manager: Optional[CallbackManager] = None,
@@ -102,7 +94,7 @@ def _retrieve(self, query_bundle: QueryBundle) -> List[NodeWithScore]:
         query = query_bundle.query_str
         sql = f"""
                 SELECT
-                    fts_main_{self._table_name}.match_bm25({self._node_id_column}, '{query}') AS score,
+                    fts_main_{self._table_name}.match_bm25({self._node_id_column}, ?) AS score,
                     {self._node_id_column}, {self._text_column}
                 FROM {self._table_name}
                 WHERE score IS NOT NULL
@@ -110,7 +102,7 @@ def _retrieve(self, query_bundle: QueryBundle) -> List[NodeWithScore]:
                 LIMIT {self._similarity_top_k};
             """
         with DuckDBLocalContext(self._database_path) as conn:
-            query_result = conn.execute(sql).fetchall()
+            query_result = conn.execute(sql, [query]).fetchall()
         # Convert query result to NodeWithScore objects
         retrieve_nodes = []
         for row in query_result:
diff --git a/llama-index-integrations/retrievers/llama-index-retrievers-duckdb-retriever/pyproject.toml b/llama-index-integrations/retrievers/llama-index-retrievers-duckdb-retriever/pyproject.toml
index 299caf35148..87839bbe4a9 100644
--- a/llama-index-integrations/retrievers/llama-index-retrievers-duckdb-retriever/pyproject.toml
+++ b/llama-index-integrations/retrievers/llama-index-retrievers-duckdb-retriever/pyproject.toml
@@ -27,12 +27,12 @@ exclude = ["**/BUILD"]
 license = "MIT"
 name = "llama-index-retrievers-duckdb-retriever"
 readme = "README.md"
-version = "0.3.0"
+version = "0.4.0"
 
 [tool.poetry.dependencies]
 python = ">=3.9,<4.0"
-pymongo = "^4.6.1"
 llama-index-core = "^0.12.0"
+duckdb = "^1.1.3"
 
 [tool.poetry.group.dev.dependencies]
 ipython = "8.10.0"
diff --git a/llama-index-integrations/retrievers/llama-index-retrievers-duckdb-retriever/tests/test_retrievers_bm25_retriever.py b/llama-index-integrations/retrievers/llama-index-retrievers-duckdb-retriever/tests/test_retrievers_bm25_retriever.py
index 01cc5ff47e0..a795dfb4f00 100644
--- a/llama-index-integrations/retrievers/llama-index-retrievers-duckdb-retriever/tests/test_retrievers_bm25_retriever.py
+++ b/llama-index-integrations/retrievers/llama-index-retrievers-duckdb-retriever/tests/test_retrievers_bm25_retriever.py
@@ -1,3 +1,4 @@
+import duckdb
 from llama_index.core.base.base_retriever import BaseRetriever
 from llama_index.retrievers.duckdb_retriever.base import (
     DuckDBRetriever,
@@ -7,3 +8,16 @@
 def test_class():
     names_of_base_classes = [b.__name__ for b in DuckDBRetriever.__mro__]
     assert BaseRetriever.__name__ in names_of_base_classes
+
+
+def test_injection(tmp_path):
+    db = tmp_path / "test.db"
+    with duckdb.connect(db) as conn:
+        conn.sql(
+            "CREATE TABLE documents (node_id VARCHAR, text VARCHAR, author VARCHAR, doc_version INTEGER);"
+        )
+    r = DuckDBRetriever(database_name=db)
+    nodes = r.retrieve(
+        "life') AS score, node_id, text FROM documents UNION SELECT '1500', '!', concat('life', version()) UNION SELECT concat('0"
+    )
+    assert not nodes