Repository navigation
feat(ibm_db): Add async APIs support using ibm_db 3.3.0 - #4082
GeetikaChughIBM wants to merge 10 commits into
Conversation
- Bump ibm-db dependency from >=3.2.8 to >=3.3.0 - Add pytest-asyncio to test deps; set asyncio_mode=auto in pyproject.toml IBMDb2DocumentStore: - Add _async_connection, _async_connection_lock, _async_table_initialized state separate from the synchronous connection path - Add _get_connection_async(): uses ibm_db_dbi.pconnect_async() (persistent pool) wrapped in AsyncConnection; required for VECTOR_DISTANCE queries which hang on non-persistent connect()-based connections - Add _transaction_async(): asynccontextmanager yielding AsyncCursor, commit on success, rollback + DocumentStoreError on failure - Add _ensure_table_exists_async() - Add close_async(): closes only the async connection, sync path unaffected - Add async counterparts for all public operations: count_documents_async, count_documents_by_filter_async, write_documents_async, filter_documents_async, delete_documents_async, delete_by_filter_async, delete_all_documents_async, update_by_filter_async, get_metadata_field_unique_values_async, get_metadata_field_min_max_async, get_metadata_fields_info_async, count_unique_metadata_by_filter_async, _embedding_retrieval_async - Private insert/skip/upsert helpers: _insert_documents_async, _skip_duplicate_documents_async, _upsert_documents_async IBMDb2EmbeddingRetriever: - Add run_async(): decorated with @component.output_types, awaits _embedding_retrieval_async, applies same filter policy as run() - Add close_async(): delegates to document_store.close_async() only; does not touch the synchronous connection Tests: - tests/conftest.py: switch document_store fixture to pytest_asyncio.fixture (async) so close_async() is properly awaited in teardown, releasing pconnect pool locks before table drop - tests/test_document_store_async.py: Haystack async mixin suite (CountDocumentsAsyncTest, WriteDocumentsAsyncTest, DeleteDocumentsAsyncTest, DeleteAllAsyncTest, DeleteByFilterAsyncTest, FilterDocumentsAsyncTest, UpdateByFilterAsyncTest, CountDocumentsByFilterAsyncTest, CountUniqueMetadataByFilterAsyncTest, GetMetadataFieldsInfoAsyncTest, GetMetadataFieldMinMaxAsyncTest, GetMetadataFieldUniqueValuesAsyncTest) plus custom integration and unit tests - tests/test_embedding_retriever_async.py: unit tests with AsyncMock and live integration tests for run_async() and close_async() All 46 integration tests pass against a live Db2 instance (SSL/port 50001). All 17 unit tests pass without a live DB2 instance.
Add _make_mock_cursor/_make_mock_conn helpers and mock_connected_store fixture that patches _get_connection_async with a mock AsyncConnection. Covers all async execution paths (count, filter, write, delete, update, metadata, embedding retrieval, transaction rollback) without requiring a live DB2 instance.
- Replace assert guards with explicit RuntimeError in _ensure_table_exists{,_async}
- Move fast-path return inside asyncio.Lock in _get_connection_async to
eliminate check-and-use race condition with concurrent coroutines / close_async()
- Fix delete_all_documents_async: run _ensure_table_exists_async outside
_transaction_async to avoid nested cursors on same connection during DROP/CREATE
- Extract _create_table_sql() and _validate_documents() helpers to eliminate
~30 lines of duplicated logic across sync/async write and table-creation paths
- Fix import order: ibm_db_dbi imports before haystack imports (ruff I001) - Replace BaseException with Exception in _embedding_retrieval_async to avoid silencing KeyboardInterrupt / SystemExit on zero-vector error path - Add asyncio_default_fixture_loop_scope = "function" to pyproject.toml to silence pytest-asyncio 0.23+ deprecation warning - Add unit test for delete_all_documents_async(recreate_index=True) to cover the previously untested DDL branch
Sync parity fixes (same class of bugs already fixed in async): - _get_connection: move fast-path check inside threading.Lock (race condition) - delete_all_documents: run _ensure_table_exists outside _transaction to avoid nested cursors on same connection during DROP/CREATE - _embedding_retrieval: BaseException -> Exception (sync path missed earlier) Connection lifecycle: - asyncio.Lock created lazily inside event loop to avoid DeprecationWarning in Python 3.10+ when creating asyncio primitives outside a running loop _transaction / _transaction_async: - Re-raise DocumentStoreError and DuplicateDocumentError unchanged; only wrap truly unexpected exceptions as DocumentStoreError get_metadata_field_min_max (sync + async): - Use _transaction instead of bare cursor for consistent rollback semantics - Extract SQL strings before the with-block to reduce nesting depth count_unique_metadata_by_filter (sync + async): - Replace fetchall + Python set-dedup with COUNT(DISTINCT ...) SQL aggregate Old approach fetched every row into memory; new approach returns 1 row _to_row: - Use json.dumps(doc.embedding) instead of f-string — f-string relies on Python list __repr__ which is implementation-defined New unit tests: - Concurrent _get_connection_async: proves only one connection created - test_insert_documents_async_maps_duplicate_error[4 parametrised cases] - test_get_metadata_field_min_max_async_fallback_to_lex - test_transaction_async_does_not_rewrap_duplicate_document_error All 327 tests passed (110 unit + 217 integration)
|
Hi @GeetikaChughIBM, thanks for your interest in contributing to Haystack! 🙏 This is an automated message to help us keep the review queue healthy. |
Coverage report (ibm_db)Click to see where and how coverage changed
This report was generated by python-coverage-comment-action |
||||||||||||||||||||||||||||||||||||
|
@bogdankostic I believe some comments you left in your review of apply here too #3980 |
All async methods share a single AsyncConnection and Db2 has one transaction per connection, so commit()/rollback() in one coroutine could commit or discard work done by another coroutine in the meantime. - Hold an asyncio.Lock for the whole unit of work in _transaction_async. _insert_documents_async, _embedding_retrieval_async and the DDL step in delete_all_documents_async take the same lock. - close_async() waits for any in-flight transaction before closing. - Share a single _table_initialized flag between the sync and async paths so recreate_table=True only drops the table once per instance. - Create asyncio.Lock instances in __init__ (Python >= 3.10 binds them to a loop on first use) and drop the lazy None checks. - Use ibm_db_dbi.AsyncConnection / AsyncCursor instead of a second from-import with its own type-ignore. - Use single backticks in docstrings instead of Sphinx roles. Tests: add regression tests for transaction serialization, lock use in insert / retrieval / close, and the shared table-initialized flag.
bogdankostic
left a comment
There was a problem hiding this comment.
Thanks for picking this up @GeetikaChughIBM! There are a few things left before we can merge this, I left inline comments with details.
| # Use pconnect_async (persistent pool) rather than AsyncConnection.connect | ||
| # (non-persistent). Non-persistent connections hang on VECTOR_DISTANCE | ||
| # queries with large CLOB parameter binding; the persistent pool avoids this. | ||
| sync_conn = await ibm_db_dbi.pconnect_async( | ||
| dsn=dsn, | ||
| user=self.username.resolve_value() or "", | ||
| password=self.password.resolve_value() or "", | ||
| conn_options=conn_options, | ||
| ) | ||
| conn = ibm_db_dbi.AsyncConnection(sync_conn) |
There was a problem hiding this comment.
pconnect caches connections per user, DSN, password, process and thread, and close() is a no-op for persistent connections (see ibm_db.c). Since pconnect_async just runs pconnect via asyncio.to_thread, and the default executor reuses idle threads, different store instances with the same credentials usually end up on the same physical connection. Each instance has its own _async_transaction_lock though, just across instances. I checked this against the Db2 container from docker-compose.yml:
a = IBMDb2DocumentStore(**conn, table_name="REPRO_A")
b = IBMDb2DocumentStore(**conn, table_name="REPRO_B")
await a.count_documents_async()
await b.count_documents_async()
# both return the same APPLICATION_ID() -> same physical connection
await b.write_documents_async([Document(id="dup")])
await asyncio.gather(
a.write_documents_async(docs, policy=DuplicatePolicy.OVERWRITE), # 150 docs -> 150
b.write_documents_async([Document(id="dup")], policy=DuplicatePolicy.FAIL), # DuplicateDocumentError -> rollback()
)
# only 129 rows are committed in REPRO_A, B's rollback discarded A's first 21 MERGEsIn the same way, SET SCHEMA of one store applies to the other one: with schema="REPRO_SA" and schema="REPRO_SB", a document written through the first store ends up in REPRO_SB. And close_async() doesn't release anything, the handle stays active afterwards.
With AsyncConnection.connect, each store gets its own connection and all of the above works as expected. I couldn't reproduce the hang on VECTOR_DISTANCE queries either, the full integration suite passes with non-persistent connections. Could you share how to reproduce it in case you still see it? Otherwise, let's switch to AsyncConnection.connect, update the docstring above accordingly, and add an integration test with two store instances like the one above.
| # Use pconnect_async (persistent pool) rather than AsyncConnection.connect | |
| # (non-persistent). Non-persistent connections hang on VECTOR_DISTANCE | |
| # queries with large CLOB parameter binding; the persistent pool avoids this. | |
| sync_conn = await ibm_db_dbi.pconnect_async( | |
| dsn=dsn, | |
| user=self.username.resolve_value() or "", | |
| password=self.password.resolve_value() or "", | |
| conn_options=conn_options, | |
| ) | |
| conn = ibm_db_dbi.AsyncConnection(sync_conn) | |
| conn = await ibm_db_dbi.AsyncConnection.connect( | |
| dsn=dsn, | |
| user=self.username.resolve_value() or "", | |
| password=self.password.resolve_value() or "", | |
| conn_options=conn_options, | |
| ) |
| dsn = f"DATABASE={self.database};HOSTNAME={self.hostname};PORT={self.port};PROTOCOL={self.protocol}" | ||
|
|
||
| if self.use_ssl: | ||
| dsn += ";SECURITY=SSL" | ||
| if self.ssl_certificate: | ||
| dsn += f";SSLServerCertificate={self.ssl_certificate}" | ||
|
|
||
| conn_options = {ibm_db_dbi.SQL_ATTR_AUTOCOMMIT: ibm_db_dbi.SQL_AUTOCOMMIT_OFF} | ||
| if self.connection_options: | ||
| conn_options.update(self.connection_options) |
There was a problem hiding this comment.
Minor: The DSN and conn_options building is a copy of the one in _get_connection. Let's move it to a separate method used by both.
| self._async_connection: ibm_db_dbi.AsyncConnection | None = None | ||
| # Guards creation/closing of the async connection. | ||
| self._async_connection_lock = asyncio.Lock() | ||
| # All coroutines share one AsyncConnection and Db2 has one transaction per connection, so | ||
| # commit()/rollback() would otherwise affect work done by other coroutines in the meantime. | ||
| # This lock serializes units of work. Lock order is always transaction lock -> connection lock. | ||
| self._async_transaction_lock = asyncio.Lock() |
There was a problem hiding this comment.
For my understanding: Do we still need _async_connection_lock? _get_connection_async is only called while _async_transaction_lock is held, and close_async takes both locks, so I think it never blocks anything. Having a single lock would make this easier to follow.
| async def _insert_documents_async(self, documents: list[Document]) -> int: | ||
| """Async version: insert documents and fail on duplicates via database integrity errors.""" | ||
| rows = [self._to_row(doc) for doc in documents] | ||
| sql = ( | ||
| f"INSERT INTO {self.table_name} (id, content, meta, embedding) " | ||
| f"VALUES (?, ?, SYSTOOLS.JSON2BSON(?), " | ||
| f"VECTOR(CAST(? AS CLOB(100000)), {self.embedding_dim}, FLOAT32))" | ||
| ) | ||
| async with self._async_transaction_lock: | ||
| conn = await self._get_connection_async() | ||
| cur = await conn.cursor() | ||
| async with cur: | ||
| try: | ||
| await cur.executemany(sql, rows) | ||
| await conn.commit() | ||
| except Exception as e: | ||
| await conn.rollback() | ||
| error_msg = str(e).lower() | ||
| duplicate_indicators = ( | ||
| "duplicate", | ||
| "unique", | ||
| "sql0803n", | ||
| "primary key", | ||
| "sqlcode=-803", | ||
| "sqlstate=23505", | ||
| ) | ||
| if any(indicator in error_msg for indicator in duplicate_indicators): | ||
| msg = f"Document already exists. Use DuplicatePolicy.OVERWRITE or SKIP. Original error: {e}" | ||
| raise DuplicateDocumentError(msg) from e | ||
| raise | ||
| return len(documents) |
There was a problem hiding this comment.
The INSERT statement and the duplicate detection are copies of _insert_documents, same for the MERGE statements in _skip_duplicate_documents_async and _upsert_documents_async. Let's move them to separate methods used by both.
Since _transaction_async now re-raises DuplicateDocumentError unchanged, this can also go through it instead of handling the lock, cursor, commit and rollback itself. Same for the sync version with _transaction:
async def _insert_documents_async(self, documents: list[Document]) -> int:
rows = [self._to_row(doc) for doc in documents]
async with self._transaction_async("Failed to insert documents") as cur:
try:
await cur.executemany(self._insert_sql(), rows)
except Exception as e:
if self._is_duplicate_error(e):
msg = f"Document already exists. Use DuplicatePolicy.OVERWRITE or SKIP. Original error: {e}"
raise DuplicateDocumentError(msg) from e
raise
return len(documents)| # Count and plain-delete share a single cursor via _transaction. | ||
| # recreate_index is handled *outside* that transaction because | ||
| # _ensure_table_exists opens its own cursor on the same connection; | ||
| # having two cursors open simultaneously on a Db2 connection during a | ||
| # DROP/CREATE causes undefined behaviour. | ||
| with self._transaction("Failed to count documents") as cur: | ||
| cur.execute(f"SELECT COUNT(*) FROM {self.table_name}") | ||
| deleted_count = cur.fetchone()[0] | ||
|
|
||
| if recreate_index: | ||
| # Drop and recreate the table | ||
| self._ensure_table_exists(recreate=True) | ||
| else: | ||
| # Just delete all rows | ||
| if recreate_index: | ||
| self._ensure_table_exists(recreate=True) | ||
| else: | ||
| with self._transaction("Failed to delete all documents") as cur: | ||
| cur.execute(f"DELETE FROM {self.table_name}") | ||
|
|
||
| return deleted_count | ||
| return deleted_count |
There was a problem hiding this comment.
On main, the count and the DELETE run in the same transaction. Now they're two separate ones, even though the comment says they share a cursor. Only the recreate_index path needs to be outside of the transaction, so let's keep the count and the DELETE together:
| # Count and plain-delete share a single cursor via _transaction. | |
| # recreate_index is handled *outside* that transaction because | |
| # _ensure_table_exists opens its own cursor on the same connection; | |
| # having two cursors open simultaneously on a Db2 connection during a | |
| # DROP/CREATE causes undefined behaviour. | |
| with self._transaction("Failed to count documents") as cur: | |
| cur.execute(f"SELECT COUNT(*) FROM {self.table_name}") | |
| deleted_count = cur.fetchone()[0] | |
| if recreate_index: | |
| # Drop and recreate the table | |
| self._ensure_table_exists(recreate=True) | |
| else: | |
| # Just delete all rows | |
| if recreate_index: | |
| self._ensure_table_exists(recreate=True) | |
| else: | |
| with self._transaction("Failed to delete all documents") as cur: | |
| cur.execute(f"DELETE FROM {self.table_name}") | |
| return deleted_count | |
| return deleted_count | |
| # recreate_index is handled outside of the transaction, because _ensure_table_exists | |
| # opens its own cursor on the same connection. | |
| with self._transaction("Failed to delete all documents") as cur: | |
| cur.execute(f"SELECT COUNT(*) FROM {self.table_name}") | |
| deleted_count = cur.fetchone()[0] | |
| if not recreate_index: | |
| cur.execute(f"DELETE FROM {self.table_name}") | |
| if recreate_index: | |
| self._ensure_table_exists(recreate=True) | |
| return deleted_count |
| async def test_count_not_empty_async(self, document_store: IBMDb2DocumentStore) -> None: | ||
| """Explicit count test with known documents.""" | ||
| await document_store.write_documents_async( | ||
| [Document(content="doc 1"), Document(content="doc 2"), Document(content="doc 3")] | ||
| ) | ||
| assert await document_store.count_documents_async() == 3 |
There was a problem hiding this comment.
This is the same as test_count_not_empty_async in CountDocumentsAsyncTest, so we can remove it.
| # Sync connection is still alive | ||
| assert document_store._connection is sync_conn |
There was a problem hiding this comment.
This only checks that the attribute wasn't reset. Let's also check that the sync connection still works.
| # Sync connection is still alive | |
| assert document_store._connection is sync_conn | |
| # Sync connection is still alive and usable | |
| assert document_store._connection is sync_conn | |
| assert document_store.count_documents() == 0 |
| call_kwargs = mock_document_store._embedding_retrieval_async.call_args # type: ignore[attr-defined] | ||
| merged = call_kwargs.kwargs["filters"] | ||
| # Merged result must reference both conditions | ||
| assert merged is not None |
There was a problem hiding this comment.
This passes as long as any filters are passed. Let's check the merged filters instead.
| call_kwargs = mock_document_store._embedding_retrieval_async.call_args # type: ignore[attr-defined] | |
| merged = call_kwargs.kwargs["filters"] | |
| # Merged result must reference both conditions | |
| assert merged is not None | |
| mock_document_store._embedding_retrieval_async.assert_awaited_once_with( # type: ignore[attr-defined] | |
| [0.0] * 768, | |
| filters={"operator": "AND", "conditions": [constructor_filters, runtime_filters]}, | |
| top_k=10, | |
| ) |
| @pytest.mark.asyncio | ||
| async def test_run_async_no_filters(retriever: IBMDb2EmbeddingRetriever) -> None: | ||
| """run_async() without runtime filters passes the constructor filters.""" | ||
| await retriever.run_async(query_embedding=[0.0] * 768) | ||
| retriever.document_store._embedding_retrieval_async.assert_awaited_once_with( # type: ignore[attr-defined] | ||
| [0.0] * 768, filters={}, top_k=5 | ||
| ) | ||
|
|
||
|
|
||
| @pytest.mark.asyncio | ||
| async def test_run_async_returns_correct_output_shape(retriever: IBMDb2EmbeddingRetriever) -> None: | ||
| """run_async() output always has the 'documents' key.""" | ||
| result = await retriever.run_async(query_embedding=[0.0] * 768) | ||
| assert "documents" in result | ||
| assert isinstance(result["documents"], list) |
There was a problem hiding this comment.
I don't think we need these two tests, test_run_async_calls_embedding_retrieval already checks that the constructor filters are used and what the output looks like.
| @pytest.mark.asyncio | ||
| async def test_close_async_only_closes_async_connection(retriever: IBMDb2EmbeddingRetriever) -> None: | ||
| """close_async() on the retriever delegates to document_store.close_async() only.""" | ||
| await retriever.close_async() | ||
| retriever.document_store.close_async.assert_awaited_once() # type: ignore[attr-defined] | ||
|
|
||
|
|
||
| @pytest.mark.asyncio | ||
| async def test_close_async_does_not_call_sync_close(retriever: IBMDb2EmbeddingRetriever) -> None: | ||
| """Calling close_async() must not touch the synchronous close().""" | ||
| retriever.document_store.close = MagicMock() # type: ignore[method-assign] | ||
| await retriever.close_async() | ||
| retriever.document_store.close.assert_not_called() # type: ignore[attr-defined] |
There was a problem hiding this comment.
Minor: These two tests can be merged into one.

Summary
Adds full async support to the IBM Db2 Haystack integration, built on the
AsyncConnectionandAsyncCursorAPIs introduced inibm_db 3.3.0.All 13 public document-store operations now have async counterparts, and the
embedding retriever gains
run_async()andclose_async(). The implementationis lint-clean (ruff + mypy) and ships 183 unit tests that require no live
database, plus a full integration test suite verified against a Db2 instance.
What changed
document_store.pyNew async connection infrastructure
_get_connection_async()— usesibm_db_dbi.pconnect_async()(persistent pool)wrapped in
AsyncConnectionrather than bareAsyncConnection.connect(). Thenon-persistent connect path hangs indefinitely on
VECTOR_DISTANCEqueries dueto large CLOB parameter binding; the persistent pool avoids this entirely.
asyncio.Lockis created lazily on first use inside the running event loop,avoiding the
DeprecationWarningraised in Python 3.10+ when locks areinstantiated at
__init__time.eliminate the TOCTOU race condition in concurrent coroutine scenarios.
_transaction_async()—asynccontextmanagerthat commits on success,rolls back on failure, and re-raises
DuplicateDocumentError/DocumentStoreErrorunwrapped; only truly unexpected exceptions are wrapped._ensure_table_exists_async()— async DDL setup, called outside_transaction_asyncindelete_all_documents_asyncto avoid nested-cursorconflicts during DROP/CREATE on the same connection.
New public async methods
write_documents_async,count_documents_async,count_documents_by_filter_async,filter_documents_async,delete_documents_async,delete_by_filter_async,delete_all_documents_async,update_by_filter_async,get_metadata_fields_info_async,get_metadata_field_min_max_async,get_metadata_field_unique_values_async,count_unique_metadata_by_filter_async,embedding_retrieval_async,close_async.Correctness fixes applied to sync path for parity
_get_connection: fast-path check moved insidethreading.Lock(same racecondition that existed in the async path).
delete_all_documents:_ensure_table_existsnow runs outside_transactionfor the same nested-cursor reason.
_transaction/_transaction_async: both now re-raise known store exceptionsunwrapped.
count_unique_metadata_by_filter(sync + async): replacedfetchall()+Python set dedup with a single
COUNT(DISTINCT …)SQL aggregate — eliminatesunbounded memory use on large tables.
_to_row:json.dumps(doc.embedding)replaces the f-string form, whichrelied on Python's list
__repr__(implementation-defined, not JSON spec)._embedding_retrieval(sync):BaseExceptionnarrowed toExceptiontoavoid silencing
KeyboardInterrupt/SystemExiton the zero-vector error path.embedding_retriever.pyrun_async()— async counterpart torun(), decorated with@component.output_types, delegates toembedding_retrieval_asyncand appliesthe same filter-merge policy.
close_async()— releases only the async connection without touching the syncconnection.
pyproject.tomlibm-dblower bound bumped to>=3.3.0.pytest-asyncioadded to test dependencies.asyncio_mode = "auto"andasyncio_default_fixture_loop_scope = "function"set to suppress pytest-asyncio 0.23+ deprecation warnings.
Tests
tests/test_document_store_async.py— full Haystack async mixin suite(
CountDocumentsAsyncTest,WriteDocumentsAsyncTest,FilterDocumentsAsyncTest,DeleteDocumentsAsyncTest,DeleteAllAsyncTest,DeleteByFilterAsyncTest,UpdateByFilterAsyncTest,CountDocumentsByFilterAsyncTest,CountUniqueMetadataByFilterAsyncTest,GetMetadataFieldsInfoAsyncTest,GetMetadataFieldMinMaxAsyncTest,GetMetadataFieldUniqueValuesAsyncTest)plus targeted unit tests for: concurrent
_get_connection_async(singleconnection under race),
DuplicateDocumentErrorindicator mapping (4parametrised SQL fragments),
_transaction_asyncrollback, no-rewrap of knownstore exceptions, and the numeric→lex fallback in
get_metadata_field_min_max_async.tests/test_embedding_retriever_async.py— unit tests forrun_asyncfilter merging, top-k override, output shape, and
close_asyncisolation fromsync connection; plus integration tests against a live Db2 instance.
tests/conftest.py—document_storefixture converted to@pytest_asyncio.fixture(async) soclose_async()is properly awaited inteardown, releasing the persistent-pool lock before
DROP TABLE.Test results
hatch run test:unit-cov-retrydocker compose upRelated