mirror of
https://github.com/langflow-ai/langflow.git
synced 2026-07-23 23:13:58 +08:00
fix: ibm db2 bulk insert pr 13438 (#13449)
* feat(ibm): add IBM Db2 Vector Search component with comprehensive tests - Add Db2VectorStore component with vector search capabilities - Implement search methods (similarity, MMR, similarity with score) - Add comprehensive test suite covering all search operations - Enhance db2vs.py with improved error handling and connection management * fix(ibm): preserve DB2 bulk insert values --------- Co-authored-by: Dhruv Chaturvedi <dhruv_insights@mac.lkw-in.ibm.com>
This commit is contained in:
Eric Hare
@ -160,6 +160,17 @@ class DB2VectorStoreComponent(LCVectorStoreComponent):
|
||||
advanced=True,
|
||||
info="Distance calculation strategy",
|
||||
),
|
||||
BoolInput(
|
||||
name="use_bulk_insert",
|
||||
display_name="Use Bulk Insert",
|
||||
value=True,
|
||||
advanced=True,
|
||||
info=(
|
||||
"Enable bulk insert using executemany() for better performance. "
|
||||
"When enabled, all documents are inserted in a single batch (unlimited chunk size). "
|
||||
"When disabled, uses row-by-row insert with execute() (chunk size: 1 document per call)."
|
||||
),
|
||||
),
|
||||
]
|
||||
|
||||
outputs = [
|
||||
@ -372,6 +383,11 @@ class DB2VectorStoreComponent(LCVectorStoreComponent):
|
||||
)
|
||||
raise ValueError(msg)
|
||||
|
||||
# Log bulk insert mode
|
||||
bulk_insert_enabled = getattr(self, "use_bulk_insert", True)
|
||||
insert_mode = "bulk insert (executemany)" if bulk_insert_enabled else "row-by-row insert (execute)"
|
||||
self.log(f"Insert mode: {insert_mode}")
|
||||
|
||||
# Build vector store (will automatically generate embeddings for existing empty rows)
|
||||
self.log(f"Connecting to DB2 table: {validated_table_name}")
|
||||
vector_store = DB2VS(
|
||||
@ -379,6 +395,7 @@ class DB2VectorStoreComponent(LCVectorStoreComponent):
|
||||
embedding_function=self.embedding,
|
||||
table_name=validated_table_name,
|
||||
distance_strategy=distance_strategy_map.get(self.distance_strategy, DistanceStrategy.COSINE),
|
||||
use_bulk_insert=bulk_insert_enabled,
|
||||
)
|
||||
|
||||
self.log(f"Connected to DB2 table: {validated_table_name}")
|
||||
|
||||
@ -344,6 +344,7 @@ class DB2VS(VectorStore):
|
||||
distance_strategy: DistanceStrategy = DistanceStrategy.EUCLIDEAN_DISTANCE,
|
||||
query: str | None = "What is a Db2 database",
|
||||
params: dict[str, Any] | None = None,
|
||||
use_bulk_insert: bool = True, # noqa: FBT001, FBT002
|
||||
):
|
||||
"""Initialize DB2 vector store with security validations.
|
||||
|
||||
@ -354,6 +355,7 @@ class DB2VS(VectorStore):
|
||||
distance_strategy: Strategy for distance calculation
|
||||
query: Optional default query
|
||||
params: Optional additional parameters
|
||||
use_bulk_insert: Enable bulk insert using executemany() for better performance (default: True)
|
||||
|
||||
Raises:
|
||||
ValueError: If table name is invalid or other validation fails
|
||||
@ -371,6 +373,7 @@ class DB2VS(VectorStore):
|
||||
|
||||
self.client = client
|
||||
self.table_name = validated_table_name
|
||||
self.use_bulk_insert = use_bulk_insert
|
||||
|
||||
if not isinstance(embedding_function, Embeddings):
|
||||
logger.warning(
|
||||
@ -532,6 +535,9 @@ class DB2VS(VectorStore):
|
||||
generated_ids = [str(uuid.uuid4()) for _ in texts]
|
||||
processed_ids = [hashlib.sha256(_id.encode()).hexdigest()[:16].upper() for _id in generated_ids]
|
||||
|
||||
if not texts:
|
||||
return []
|
||||
|
||||
# Generate embeddings
|
||||
embeddings = self._embed_documents(texts)
|
||||
|
||||
@ -545,47 +551,77 @@ class DB2VS(VectorStore):
|
||||
|
||||
cursor = self.client.cursor()
|
||||
try:
|
||||
logger.debug("Inserting %s documents into %s", len(processed_ids), self.table_name)
|
||||
if self.use_bulk_insert:
|
||||
# BULK INSERT MODE: Use executemany() for better performance
|
||||
logger.debug("Using BULK INSERT mode (executemany) for %s documents", len(processed_ids))
|
||||
|
||||
# SECURITY FIX: Use parameterized queries where possible
|
||||
# For DB2 VECTOR() function, we need to use string formatting but with proper sanitization
|
||||
for idx, (id_, embedding, metadata, text) in enumerate(
|
||||
zip(processed_ids, embeddings, metadatas, texts, strict=False), 1
|
||||
):
|
||||
# Convert numpy array to Python list if needed
|
||||
embedding_list = embedding.tolist() if hasattr(embedding, "tolist") else list(embedding)
|
||||
|
||||
# Format as string for VECTOR() function
|
||||
vector_str = str(embedding_list)
|
||||
|
||||
# SECURITY: Properly sanitize text using SQL-safe escaping
|
||||
text_sanitized = sanitize_sql_string(text) if text else ""
|
||||
|
||||
# SECURITY: Sanitize metadata JSON
|
||||
metadata_json = json.dumps(metadata)
|
||||
metadata_sanitized = sanitize_sql_string(metadata_json)
|
||||
|
||||
# SECURITY: Sanitize ID (already limited to 100 chars)
|
||||
id_sanitized = sanitize_sql_string(id_)
|
||||
|
||||
# Build SQL with sanitized values
|
||||
# Note: table_name and column_names are validated during initialization
|
||||
# Build the SQL statement with placeholders for bulk insert
|
||||
sql_insert = f"""
|
||||
INSERT INTO {_quoted_table_identifier(self.table_name)}
|
||||
({self.column_names["id"]}, {self.column_names["embedding"]},
|
||||
{self.column_names["metadata"]}, {self.column_names["text"]})
|
||||
VALUES ('{id_sanitized}', VECTOR('{vector_str}', {embedding_len}, FLOAT32),
|
||||
'{metadata_sanitized}', '{text_sanitized}')
|
||||
VALUES (?, VECTOR(?, {embedding_len}, FLOAT32), ?, ?)
|
||||
""" # noqa: S608
|
||||
|
||||
# Reduced logging - only log summary
|
||||
if idx == 1:
|
||||
logger.debug("Inserting documents 1-%s into %s", len(processed_ids), self.table_name)
|
||||
# Prepare data tuples for executemany
|
||||
insert_data = []
|
||||
for id_, embedding, metadata, text in zip(processed_ids, embeddings, metadatas, texts, strict=False):
|
||||
# Convert numpy array to Python list if needed
|
||||
embedding_list = embedding.tolist() if hasattr(embedding, "tolist") else list(embedding)
|
||||
|
||||
cursor.execute(sql_insert)
|
||||
# Format as string for VECTOR() function
|
||||
vector_str = str(embedding_list)
|
||||
|
||||
metadata_json = json.dumps(metadata)
|
||||
|
||||
insert_data.append((id_, vector_str, metadata_json, text or ""))
|
||||
|
||||
# Use executemany for bulk insert - much more efficient
|
||||
cursor.executemany(sql_insert, insert_data)
|
||||
logger.debug("Successfully inserted %s documents using bulk insert", len(processed_ids))
|
||||
|
||||
else:
|
||||
# ROW-BY-ROW INSERT MODE: Use execute() for each row (for debugging/compatibility)
|
||||
logger.debug("Using ROW-BY-ROW INSERT mode (execute) for %s documents", len(processed_ids))
|
||||
|
||||
for idx, (id_, embedding, metadata, text) in enumerate(
|
||||
zip(processed_ids, embeddings, metadatas, texts, strict=False), 1
|
||||
):
|
||||
# Convert numpy array to Python list if needed
|
||||
embedding_list = embedding.tolist() if hasattr(embedding, "tolist") else list(embedding)
|
||||
|
||||
# Format as string for VECTOR() function
|
||||
vector_str = str(embedding_list)
|
||||
|
||||
# SECURITY: Properly sanitize text using SQL-safe escaping
|
||||
text_sanitized = sanitize_sql_string(text) if text else ""
|
||||
|
||||
# SECURITY: Sanitize metadata JSON
|
||||
metadata_json = json.dumps(metadata)
|
||||
metadata_sanitized = sanitize_sql_string(metadata_json)
|
||||
|
||||
# SECURITY: Sanitize ID (already limited to 100 chars)
|
||||
id_sanitized = sanitize_sql_string(id_)
|
||||
|
||||
# Build SQL with sanitized values
|
||||
sql_insert = f"""
|
||||
INSERT INTO {_quoted_table_identifier(self.table_name)}
|
||||
({self.column_names["id"]}, {self.column_names["embedding"]},
|
||||
{self.column_names["metadata"]}, {self.column_names["text"]})
|
||||
VALUES ('{id_sanitized}', VECTOR('{vector_str}', {embedding_len}, FLOAT32),
|
||||
'{metadata_sanitized}', '{text_sanitized}')
|
||||
""" # noqa: S608
|
||||
|
||||
# Log progress for first insert
|
||||
if idx == 1:
|
||||
logger.debug("Inserting documents 1-%s into %s", len(processed_ids), self.table_name)
|
||||
|
||||
# Execute the insert for this row
|
||||
cursor.execute(sql_insert)
|
||||
|
||||
logger.debug("Successfully inserted %s documents using row-by-row insert", len(processed_ids))
|
||||
|
||||
cursor.execute("COMMIT")
|
||||
logger.debug("Successfully inserted %s documents", len(processed_ids))
|
||||
except Exception as e:
|
||||
# Rollback on error
|
||||
import contextlib
|
||||
|
||||
@ -11,11 +11,13 @@ historical schema fixtures to validate against.
|
||||
"""
|
||||
|
||||
import importlib.util
|
||||
import json
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
from unittest.mock import MagicMock, Mock, patch
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from langchain_core.embeddings import Embeddings
|
||||
from lfx.schema.data import Data
|
||||
from lfx.schema.message import Message
|
||||
from lfx_ibm.components.ibm.db2_vector import DB2VectorStoreComponent
|
||||
@ -31,6 +33,19 @@ requires_ibm_db = pytest.mark.skipif(
|
||||
)
|
||||
|
||||
|
||||
class FakeEmbeddings(Embeddings):
|
||||
"""Small Embeddings test double that follows the LangChain contract."""
|
||||
|
||||
def __init__(self, vector: list[float] | None = None):
|
||||
self.vector = vector or [0.1, 0.2, 0.3]
|
||||
|
||||
def embed_documents(self, texts: list[str]) -> list[list[float]]:
|
||||
return [self.vector.copy() for _ in texts]
|
||||
|
||||
def embed_query(self, _text: str) -> list[float]:
|
||||
return self.vector.copy()
|
||||
|
||||
|
||||
@requires_ibm_db
|
||||
class TestDB2VectorStoreComponent:
|
||||
"""Test DB2 Vector Store Component."""
|
||||
@ -54,11 +69,8 @@ class TestDB2VectorStoreComponent:
|
||||
|
||||
@pytest.fixture
|
||||
def mock_embedding(self):
|
||||
"""Create a mock embedding model."""
|
||||
embedding = Mock()
|
||||
embedding.embed_documents = Mock(return_value=[[0.1, 0.2, 0.3]])
|
||||
embedding.embed_query = Mock(return_value=[0.1, 0.2, 0.3])
|
||||
return embedding
|
||||
"""Create an embedding model test double."""
|
||||
return FakeEmbeddings()
|
||||
|
||||
def test_component_metadata(self):
|
||||
"""Test component metadata is correctly set."""
|
||||
@ -821,5 +833,400 @@ class TestDB2VectorStoreComponent:
|
||||
):
|
||||
component._add_documents_to_vector_store(mock_vector_store)
|
||||
|
||||
def test_bulk_insert_enabled_by_default(self, component, mock_embedding):
|
||||
"""Test that bulk insert is enabled by default."""
|
||||
component.embedding = mock_embedding
|
||||
|
||||
with (
|
||||
patch("ibm_db_dbi.connect") as mock_connect,
|
||||
patch.object(component, "_add_documents_to_vector_store"),
|
||||
):
|
||||
mock_connection = MagicMock()
|
||||
mock_connect.return_value = mock_connection
|
||||
|
||||
from lfx_ibm.components.ibm.db2vs import DB2VS
|
||||
|
||||
with patch.object(DB2VS, "__init__", return_value=None) as mock_init:
|
||||
component.build_vector_store()
|
||||
|
||||
# Verify use_bulk_insert was passed as True (default)
|
||||
call_kwargs = mock_init.call_args[1]
|
||||
assert call_kwargs["use_bulk_insert"] is True
|
||||
|
||||
def test_bulk_insert_toggle_enabled(self, component, mock_embedding):
|
||||
"""Test bulk insert when explicitly enabled."""
|
||||
component.embedding = mock_embedding
|
||||
component.use_bulk_insert = True
|
||||
|
||||
with (
|
||||
patch("ibm_db_dbi.connect") as mock_connect,
|
||||
patch.object(component, "_add_documents_to_vector_store"),
|
||||
):
|
||||
mock_connection = MagicMock()
|
||||
mock_connect.return_value = mock_connection
|
||||
|
||||
from lfx_ibm.components.ibm.db2vs import DB2VS
|
||||
|
||||
with patch.object(DB2VS, "__init__", return_value=None) as mock_init:
|
||||
component.build_vector_store()
|
||||
|
||||
# Verify use_bulk_insert was passed as True
|
||||
call_kwargs = mock_init.call_args[1]
|
||||
assert call_kwargs["use_bulk_insert"] is True
|
||||
|
||||
def test_bulk_insert_toggle_disabled(self, component, mock_embedding):
|
||||
"""Test bulk insert when explicitly disabled."""
|
||||
component.embedding = mock_embedding
|
||||
component.use_bulk_insert = False
|
||||
|
||||
with (
|
||||
patch("ibm_db_dbi.connect") as mock_connect,
|
||||
patch.object(component, "_add_documents_to_vector_store"),
|
||||
):
|
||||
mock_connection = MagicMock()
|
||||
mock_connect.return_value = mock_connection
|
||||
|
||||
from lfx_ibm.components.ibm.db2vs import DB2VS
|
||||
|
||||
with patch.object(DB2VS, "__init__", return_value=None) as mock_init:
|
||||
component.build_vector_store()
|
||||
|
||||
# Verify use_bulk_insert was passed as False
|
||||
call_kwargs = mock_init.call_args[1]
|
||||
assert call_kwargs["use_bulk_insert"] is False
|
||||
|
||||
def test_bulk_insert_mode_uses_executemany(self, mock_embedding):
|
||||
"""Test that bulk insert mode uses executemany() for better performance."""
|
||||
from lfx_ibm.components.ibm.db2vs import DB2VS
|
||||
|
||||
mock_connection = MagicMock()
|
||||
mock_cursor = MagicMock()
|
||||
mock_connection.cursor.return_value = mock_cursor
|
||||
|
||||
# Mock table operations
|
||||
with (
|
||||
patch("lfx_ibm.components.ibm.db2vs._table_exists", return_value=False),
|
||||
patch("lfx_ibm.components.ibm.db2vs._create_table"),
|
||||
patch(
|
||||
"lfx_ibm.components.ibm.db2vs._get_column_names",
|
||||
return_value={"id": "id", "embedding": "embedding", "metadata": "metadata", "text": "text"},
|
||||
),
|
||||
):
|
||||
# Create DB2VS instance with bulk insert enabled
|
||||
vector_store = DB2VS(
|
||||
client=mock_connection,
|
||||
embedding_function=mock_embedding,
|
||||
table_name="test_table",
|
||||
use_bulk_insert=True,
|
||||
)
|
||||
|
||||
# Add some test documents
|
||||
texts = ["Document 1", "Document 2", "Document 3"]
|
||||
metadatas = [{"source": "test1"}, {"source": "test2"}, {"source": "test3"}]
|
||||
|
||||
vector_store.add_texts(texts, metadatas=metadatas)
|
||||
|
||||
# Verify executemany was called (bulk insert)
|
||||
assert mock_cursor.executemany.called, "executemany should be called in bulk insert mode"
|
||||
# Verify execute was NOT called for individual inserts
|
||||
# (execute is only called for COMMIT)
|
||||
execute_calls = [call for call in mock_cursor.execute.call_args_list if "INSERT" in str(call)]
|
||||
assert len(execute_calls) == 0, "execute should not be called for individual inserts in bulk mode"
|
||||
|
||||
def test_row_by_row_insert_mode_uses_execute(self, mock_embedding):
|
||||
"""Test that row-by-row insert mode uses execute() for each document."""
|
||||
from lfx_ibm.components.ibm.db2vs import DB2VS
|
||||
|
||||
mock_connection = MagicMock()
|
||||
mock_cursor = MagicMock()
|
||||
mock_connection.cursor.return_value = mock_cursor
|
||||
|
||||
# Mock table operations
|
||||
with (
|
||||
patch("lfx_ibm.components.ibm.db2vs._table_exists", return_value=False),
|
||||
patch("lfx_ibm.components.ibm.db2vs._create_table"),
|
||||
patch(
|
||||
"lfx_ibm.components.ibm.db2vs._get_column_names",
|
||||
return_value={"id": "id", "embedding": "embedding", "metadata": "metadata", "text": "text"},
|
||||
),
|
||||
):
|
||||
# Create DB2VS instance with bulk insert disabled
|
||||
vector_store = DB2VS(
|
||||
client=mock_connection,
|
||||
embedding_function=mock_embedding,
|
||||
table_name="test_table",
|
||||
use_bulk_insert=False,
|
||||
)
|
||||
|
||||
# Add some test documents
|
||||
texts = ["Document 1", "Document 2", "Document 3"]
|
||||
metadatas = [{"source": "test1"}, {"source": "test2"}, {"source": "test3"}]
|
||||
|
||||
vector_store.add_texts(texts, metadatas=metadatas)
|
||||
|
||||
# Verify executemany was NOT called
|
||||
assert not mock_cursor.executemany.called, "executemany should not be called in row-by-row mode"
|
||||
# Verify execute was called multiple times (once per document + COMMIT)
|
||||
assert mock_cursor.execute.call_count >= len(texts), "execute should be called for each document"
|
||||
|
||||
def test_bulk_insert_with_multiple_documents(self, component, mock_embedding):
|
||||
"""Test bulk insert with multiple documents."""
|
||||
component.embedding = mock_embedding
|
||||
component.use_bulk_insert = True
|
||||
component.ingest_data = [
|
||||
Data(text="Document 1"),
|
||||
Data(text="Document 2"),
|
||||
Data(text="Document 3"),
|
||||
Data(text="Document 4"),
|
||||
Data(text="Document 5"),
|
||||
]
|
||||
|
||||
with (
|
||||
patch("ibm_db_dbi.connect") as mock_connect,
|
||||
):
|
||||
mock_connection = MagicMock()
|
||||
mock_cursor = MagicMock()
|
||||
mock_connection.cursor.return_value = mock_cursor
|
||||
mock_connect.return_value = mock_connection
|
||||
|
||||
from lfx_ibm.components.ibm.db2vs import DB2VS
|
||||
|
||||
with (
|
||||
patch.object(DB2VS, "__init__", return_value=None),
|
||||
patch.object(DB2VS, "add_documents") as mock_add_docs,
|
||||
):
|
||||
vector_store = component.build_vector_store()
|
||||
component._add_documents_to_vector_store(vector_store)
|
||||
|
||||
# Verify documents were added
|
||||
if mock_add_docs.called:
|
||||
added_docs = mock_add_docs.call_args[0][0]
|
||||
assert len(added_docs) == 5, "All 5 documents should be added"
|
||||
|
||||
def test_bulk_insert_performance_benefit(self, mock_embedding):
|
||||
"""Test that bulk insert reduces database round-trips."""
|
||||
from lfx_ibm.components.ibm.db2vs import DB2VS
|
||||
|
||||
mock_connection = MagicMock()
|
||||
mock_cursor = MagicMock()
|
||||
mock_connection.cursor.return_value = mock_cursor
|
||||
|
||||
# Mock table operations
|
||||
with (
|
||||
patch("lfx_ibm.components.ibm.db2vs._table_exists", return_value=False),
|
||||
patch("lfx_ibm.components.ibm.db2vs._create_table"),
|
||||
patch(
|
||||
"lfx_ibm.components.ibm.db2vs._get_column_names",
|
||||
return_value={"id": "id", "embedding": "embedding", "metadata": "metadata", "text": "text"},
|
||||
),
|
||||
):
|
||||
# Test with bulk insert enabled
|
||||
vector_store_bulk = DB2VS(
|
||||
client=mock_connection,
|
||||
embedding_function=mock_embedding,
|
||||
table_name="test_table",
|
||||
use_bulk_insert=True,
|
||||
)
|
||||
|
||||
texts = ["Doc 1", "Doc 2", "Doc 3", "Doc 4", "Doc 5"]
|
||||
vector_store_bulk.add_texts(texts)
|
||||
|
||||
# Count database calls with bulk insert
|
||||
bulk_executemany_calls = mock_cursor.executemany.call_count
|
||||
bulk_execute_calls = mock_cursor.execute.call_count
|
||||
|
||||
# Reset mocks
|
||||
mock_cursor.reset_mock()
|
||||
|
||||
# Test with bulk insert disabled
|
||||
vector_store_row = DB2VS(
|
||||
client=mock_connection,
|
||||
embedding_function=mock_embedding,
|
||||
table_name="test_table",
|
||||
use_bulk_insert=False,
|
||||
)
|
||||
|
||||
vector_store_row.add_texts(texts)
|
||||
|
||||
# Count database calls with row-by-row insert
|
||||
row_executemany_calls = mock_cursor.executemany.call_count
|
||||
row_execute_calls = mock_cursor.execute.call_count
|
||||
|
||||
# Verify bulk insert uses fewer calls
|
||||
assert bulk_executemany_calls > 0, "Bulk insert should use executemany"
|
||||
assert row_executemany_calls == 0, "Row-by-row should not use executemany"
|
||||
assert row_execute_calls > bulk_execute_calls, "Row-by-row should make more execute calls"
|
||||
|
||||
def test_bulk_insert_with_empty_documents(self, mock_embedding):
|
||||
"""Test bulk insert handles empty document list gracefully."""
|
||||
from lfx_ibm.components.ibm.db2vs import DB2VS
|
||||
|
||||
mock_connection = MagicMock()
|
||||
mock_cursor = MagicMock()
|
||||
mock_connection.cursor.return_value = mock_cursor
|
||||
|
||||
# Mock table operations
|
||||
with (
|
||||
patch("lfx_ibm.components.ibm.db2vs._table_exists", return_value=False),
|
||||
patch("lfx_ibm.components.ibm.db2vs._create_table"),
|
||||
patch(
|
||||
"lfx_ibm.components.ibm.db2vs._get_column_names",
|
||||
return_value={"id": "id", "embedding": "embedding", "metadata": "metadata", "text": "text"},
|
||||
),
|
||||
):
|
||||
vector_store = DB2VS(
|
||||
client=mock_connection,
|
||||
embedding_function=mock_embedding,
|
||||
table_name="test_table",
|
||||
use_bulk_insert=True,
|
||||
)
|
||||
|
||||
# Add empty list
|
||||
result = vector_store.add_texts([])
|
||||
|
||||
# Verify no database calls were made
|
||||
mock_connection.cursor.assert_not_called()
|
||||
assert not mock_cursor.executemany.called
|
||||
assert result == []
|
||||
|
||||
def test_bulk_insert_with_special_characters(self, mock_embedding):
|
||||
"""Test bulk insert preserves special characters in bound parameters."""
|
||||
from lfx_ibm.components.ibm.db2vs import DB2VS
|
||||
|
||||
mock_connection = MagicMock()
|
||||
mock_cursor = MagicMock()
|
||||
mock_connection.cursor.return_value = mock_cursor
|
||||
|
||||
# Mock table operations
|
||||
with (
|
||||
patch("lfx_ibm.components.ibm.db2vs._table_exists", return_value=False),
|
||||
patch("lfx_ibm.components.ibm.db2vs._create_table"),
|
||||
patch(
|
||||
"lfx_ibm.components.ibm.db2vs._get_column_names",
|
||||
return_value={"id": "id", "embedding": "embedding", "metadata": "metadata", "text": "text"},
|
||||
),
|
||||
):
|
||||
vector_store = DB2VS(
|
||||
client=mock_connection,
|
||||
embedding_function=mock_embedding,
|
||||
table_name="test_table",
|
||||
use_bulk_insert=True,
|
||||
)
|
||||
|
||||
# Add documents with special characters
|
||||
texts = [
|
||||
"Document with 'single quotes'",
|
||||
'Document with "double quotes"',
|
||||
"Document with \n newlines \t tabs",
|
||||
"Document with SQL injection'; DROP TABLE users;--",
|
||||
]
|
||||
metadatas = [
|
||||
{"source": "O'Brien"},
|
||||
{"source": 'double "quote"'},
|
||||
{"source": "line\nbreak"},
|
||||
{"source": "payload'; DROP TABLE users;--"},
|
||||
]
|
||||
ids = ["id'1", 'id"2', "id\n3", "id;4"]
|
||||
|
||||
vector_store.add_texts(texts, metadatas=metadatas, ids=ids)
|
||||
|
||||
# Verify executemany was called (documents were processed)
|
||||
assert mock_cursor.executemany.called
|
||||
# Bound parameters should stay raw; the DB driver handles escaping.
|
||||
call_args = mock_cursor.executemany.call_args[0]
|
||||
assert len(call_args) == 2 # SQL statement and data tuples
|
||||
insert_data = call_args[1]
|
||||
assert insert_data[0][0] == ids[0]
|
||||
assert insert_data[0][2] == json.dumps(metadatas[0])
|
||||
assert insert_data[0][3] == texts[0]
|
||||
assert "O''Brien" not in insert_data[0][2]
|
||||
assert "''single quotes''" not in insert_data[0][3]
|
||||
|
||||
def test_bulk_insert_logging(self, component, mock_embedding):
|
||||
"""Test that bulk insert mode is logged correctly."""
|
||||
component.embedding = mock_embedding
|
||||
component.use_bulk_insert = True
|
||||
|
||||
with (
|
||||
patch("ibm_db_dbi.connect") as mock_connect,
|
||||
patch.object(component, "_add_documents_to_vector_store"),
|
||||
):
|
||||
mock_connection = MagicMock()
|
||||
mock_connect.return_value = mock_connection
|
||||
|
||||
from lfx_ibm.components.ibm.db2vs import DB2VS
|
||||
|
||||
with patch.object(DB2VS, "__init__", return_value=None), patch.object(component, "log") as mock_log:
|
||||
component.build_vector_store()
|
||||
|
||||
messages = [str(call.args[0]).lower() for call in mock_log.call_args_list]
|
||||
assert any("bulk insert" in message or "executemany" in message for message in messages)
|
||||
|
||||
def test_row_by_row_insert_logging(self, component, mock_embedding):
|
||||
"""Test that row-by-row insert mode is logged correctly."""
|
||||
component.embedding = mock_embedding
|
||||
component.use_bulk_insert = False
|
||||
|
||||
with (
|
||||
patch("ibm_db_dbi.connect") as mock_connect,
|
||||
patch.object(component, "_add_documents_to_vector_store"),
|
||||
):
|
||||
mock_connection = MagicMock()
|
||||
mock_connect.return_value = mock_connection
|
||||
|
||||
from lfx_ibm.components.ibm.db2vs import DB2VS
|
||||
|
||||
with patch.object(DB2VS, "__init__", return_value=None), patch.object(component, "log") as mock_log:
|
||||
component.build_vector_store()
|
||||
|
||||
messages = [str(call.args[0]).lower() for call in mock_log.call_args_list]
|
||||
assert any("row-by-row" in message or "execute" in message for message in messages)
|
||||
|
||||
def test_bulk_insert_error_handling(self, mock_embedding):
|
||||
"""Test that bulk insert handles errors gracefully with rollback."""
|
||||
from lfx_ibm.components.ibm.db2vs import DB2VS
|
||||
|
||||
mock_connection = MagicMock()
|
||||
mock_cursor = MagicMock()
|
||||
mock_connection.cursor.return_value = mock_cursor
|
||||
|
||||
# Make executemany raise an error
|
||||
mock_cursor.executemany.side_effect = Exception("Database error")
|
||||
|
||||
# Mock table operations
|
||||
with (
|
||||
patch("lfx_ibm.components.ibm.db2vs._table_exists", return_value=False),
|
||||
patch("lfx_ibm.components.ibm.db2vs._create_table"),
|
||||
patch(
|
||||
"lfx_ibm.components.ibm.db2vs._get_column_names",
|
||||
return_value={"id": "id", "embedding": "embedding", "metadata": "metadata", "text": "text"},
|
||||
),
|
||||
):
|
||||
vector_store = DB2VS(
|
||||
client=mock_connection,
|
||||
embedding_function=mock_embedding,
|
||||
table_name="test_table",
|
||||
use_bulk_insert=True,
|
||||
)
|
||||
|
||||
# Attempt to add documents
|
||||
with pytest.raises(RuntimeError, match="during document insertion"):
|
||||
vector_store.add_texts(["Document 1", "Document 2"])
|
||||
|
||||
# Verify rollback was attempted
|
||||
rollback_calls = [call for call in mock_cursor.execute.call_args_list if "ROLLBACK" in str(call)]
|
||||
assert len(rollback_calls) > 0, "ROLLBACK should be called on error"
|
||||
|
||||
def test_bulk_insert_input_exists(self, component):
|
||||
"""Test that use_bulk_insert input is defined in component."""
|
||||
# Check that the input exists in the component's inputs
|
||||
input_names = [inp.name for inp in component.inputs]
|
||||
assert "use_bulk_insert" in input_names, "use_bulk_insert input should be defined"
|
||||
|
||||
# Find the input and verify its properties
|
||||
bulk_insert_input = next(inp for inp in component.inputs if inp.name == "use_bulk_insert")
|
||||
assert bulk_insert_input.value is True, "Default value should be True"
|
||||
assert bulk_insert_input.advanced is True, "Should be an advanced setting"
|
||||
|
||||
|
||||
# Made with Bob
|
||||
|
||||
Reference in New Issue
Block a user