feat(ibm_db): Add async APIs support using ibm_db 3.3.0 by GeetikaChughIBM · Pull Request #4082 · deepset-ai/haystack-core-integrations · GitHub
Skip to content

feat(ibm_db): Add async APIs support using ibm_db 3.3.0 - #4082

Open
GeetikaChughIBM wants to merge 10 commits into
deepset-ai:mainfrom
GeetikaChughIBM:feature/ibm-db-async-support
Open

GeetikaChughIBM wants to merge 10 commits into
deepset-ai:mainfrom
GeetikaChughIBM:feature/ibm-db-async-support

Conversation

@GeetikaChughIBM

Copy link
Copy Markdown
Contributor

Summary

Adds full async support to the IBM Db2 Haystack integration, built on the
AsyncConnection and AsyncCursor APIs introduced in ibm_db 3.3.0.

All 13 public document-store operations now have async counterparts, and the
embedding retriever gains run_async() and close_async(). The implementation
is 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.py

New async connection infrastructure

  • _get_connection_async() — uses ibm_db_dbi.pconnect_async() (persistent pool)
    wrapped in AsyncConnection rather than bare AsyncConnection.connect(). The
    non-persistent connect path hangs indefinitely on VECTOR_DISTANCE queries due
    to large CLOB parameter binding; the persistent pool avoids this entirely.
  • asyncio.Lock is created lazily on first use inside the running event loop,
    avoiding the DeprecationWarning raised in Python 3.10+ when locks are
    instantiated at __init__ time.
  • The already-connected fast-path check is performed inside the lock to
    eliminate the TOCTOU race condition in concurrent coroutine scenarios.
  • _transaction_async() — asynccontextmanager that commits on success,
    rolls back on failure, and re-raises DuplicateDocumentError /
    DocumentStoreError unwrapped; only truly unexpected exceptions are wrapped.
  • _ensure_table_exists_async() — async DDL setup, called outside
    _transaction_async in delete_all_documents_async to avoid nested-cursor
    conflicts 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 inside threading.Lock (same race
    condition that existed in the async path).
  • delete_all_documents: _ensure_table_exists now runs outside _transaction
    for the same nested-cursor reason.
  • _transaction / _transaction_async: both now re-raise known store exceptions
    unwrapped.
  • count_unique_metadata_by_filter (sync + async): replaced fetchall() +
    Python set dedup with a single COUNT(DISTINCT …) SQL aggregate — eliminates
    unbounded memory use on large tables.
  • _to_row: json.dumps(doc.embedding) replaces the f-string form, which
    relied on Python's list __repr__ (implementation-defined, not JSON spec).
  • _embedding_retrieval (sync): BaseException narrowed to Exception to
    avoid silencing KeyboardInterrupt / SystemExit on the zero-vector error path.

embedding_retriever.py

  • run_async() — async counterpart to run(), decorated with
    @component.output_types, delegates to embedding_retrieval_async and applies
    the same filter-merge policy.
  • close_async() — releases only the async connection without touching the sync
    connection.

pyproject.toml

  • ibm-db lower bound bumped to >=3.3.0.
  • pytest-asyncio added to test dependencies.
  • asyncio_mode = "auto" and asyncio_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 (single
    connection under race), DuplicateDocumentError indicator mapping (4
    parametrised SQL fragments), _transaction_async rollback, no-rewrap of known
    store exceptions, and the numeric→lex fallback in get_metadata_field_min_max_async.
  • tests/test_embedding_retriever_async.py — unit tests for run_async
    filter merging, top-k override, output shape, and close_async isolation from
    sync connection; plus integration tests against a live Db2 instance.
  • tests/conftest.py — document_store fixture converted to
    @pytest_asyncio.fixture (async) so close_async() is properly awaited in
    teardown, releasing the persistent-pool lock before DROP TABLE.

Test results

Suite Count Notes
Unit tests (no DB) 183 passed hatch run test:unit-cov-retry
Integration tests 158 passed live Db2 via docker compose up
mypy no issues 5 source files
ruff all checks passed 14 files

Related

- 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)
@GeetikaChughIBM
GeetikaChughIBM requested a review from a team as a code owner October 10, 2026 22:31
@GeetikaChughIBM
GeetikaChughIBM requested review from julian-risch and removed request for a team October 10, 2026 22:31
@github-actions

Copy link
Copy Markdown
Contributor

Hi @GeetikaChughIBM, thanks for your interest in contributing to Haystack! 🙏

⚠️ Issue #3971 is already being addressed by open pull request(s) #3980. Before opening a PR for an issue, please check whether a PR is already linked to it, and consider contributing to the existing PR instead. We may close duplicate PRs to keep the review queue manageable.

This is an automated message to help us keep the review queue healthy.

@github-actions github-actions Bot added integration:ibm-db type:documentation Improvements or additions to documentation labels Oct 10, 2026
@github-actions

github-actions Bot commented Oct 10, 2026 •

Copy link
Copy Markdown
Contributor

Coverage report (ibm_db)

Click to see where and how coverage changed

FileStatementsMissingCoverageCoverage
(new stmts)
Lines missing
  integrations/ibm_db/src/haystack_integrations/components/retrievers/ibm_db
  embedding_retriever.py
  integrations/ibm_db/src/haystack_integrations/document_stores/ibm_db
  document_store.py 260-262, 266, 280-288, 359-360, 372-400, 409-410, 511-513, 523-524, 552-553, 563, 591, 594, 621, 630-631, 635, 660-661, 692, 696-697, 712-713, 737-738, 752-753, 791-792, 804-805, 838-839, 868-869, 899-909, 977, 1143-1144, 1153-1154, 1172-1177, 1188-1218, 1249, 1294, 1310-1327, 1346-1364, 1388, 1475, 1509
Project Total  

This report was generated by python-coverage-comment-action

@julian-risch

Copy link
Copy Markdown
Member

@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 bogdankostic left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for picking this up @GeetikaChughIBM! There are a few things left before we can merge this, I left inline comments with details.

Comment on lines +268 to +277
# 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)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 MERGEs

In 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.

Suggested change
# 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,
)

Comment on lines +257 to +266
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)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment on lines +160 to +166
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()

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment on lines +663 to +693
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)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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)

Comment on lines +894 to +909
# 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

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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:

Suggested change
# 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

Comment on lines +115 to +120
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

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is the same as test_count_not_empty_async in CountDocumentsAsyncTest, so we can remove it.

Comment on lines +151 to +152
# Sync connection is still alive
assert document_store._connection is sync_conn

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This only checks that the attribute wasn't reset. Let's also check that the sync connection still works.

Suggested change
# 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

Comment on lines +103 to +106
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

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This passes as long as any filters are passed. Let's check the merged filters instead.

Suggested change
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,
)

Comment on lines +109 to +123
@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)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment on lines +126 to +138
@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]

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Minor: These two tests can be merged into one.

Suggested change

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

integration:ibm-db type:documentation Improvements or additions to documentation

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Add async API support for ibm_db Integration

3 participants