From fb046c02dfa8322956c2fe754e63100d71f8bed8 Mon Sep 17 00:00:00 2001 From: Eric Hare Date: Mon, 1 Jun 2026 13:38:34 -0700 Subject: [PATCH] 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 --- .../src/lfx_ibm/components/ibm/db2_vector.py | 17 + .../ibm/src/lfx_ibm/components/ibm/db2vs.py | 98 ++-- src/bundles/ibm/tests/test_db2_vector.py | 419 +++++++++++++++++- 3 files changed, 497 insertions(+), 37 deletions(-) diff --git a/src/bundles/ibm/src/lfx_ibm/components/ibm/db2_vector.py b/src/bundles/ibm/src/lfx_ibm/components/ibm/db2_vector.py index 5ef3a62cf4..7c3d099ef7 100644 --- a/src/bundles/ibm/src/lfx_ibm/components/ibm/db2_vector.py +++ b/src/bundles/ibm/src/lfx_ibm/components/ibm/db2_vector.py @@ -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}") diff --git a/src/bundles/ibm/src/lfx_ibm/components/ibm/db2vs.py b/src/bundles/ibm/src/lfx_ibm/components/ibm/db2vs.py index 9874ffa1c5..403f349eec 100644 --- a/src/bundles/ibm/src/lfx_ibm/components/ibm/db2vs.py +++ b/src/bundles/ibm/src/lfx_ibm/components/ibm/db2vs.py @@ -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 diff --git a/src/bundles/ibm/tests/test_db2_vector.py b/src/bundles/ibm/tests/test_db2_vector.py index 20ff6f7bfc..ec2653e5a5 100644 --- a/src/bundles/ibm/tests/test_db2_vector.py +++ b/src/bundles/ibm/tests/test_db2_vector.py @@ -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