unstructured-ingest 0.3.6__tar.gz → 0.3.8__tar.gz

{unstructured_ingest-0.3.6 → unstructured_ingest-0.3.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: unstructured-ingest
-Version: 0.3.6
+Version: 0.3.8
 Summary: A library that prepares raw documents for downstream ML tasks.
 Home-page: https://github.com/Unstructured-IO/unstructured-ingest
 Author: Unstructured Technologies
@@ -22,13 +22,13 @@ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Requires-Python: >=3.9.0,<3.13
 Description-Content-Type: text/markdown
 License-File: LICENSE.md
-Requires-Dist: python-dateutil
-Requires-Dist: pandas
-Requires-Dist: pydantic>=2.7
 Requires-Dist: opentelemetry-sdk
+Requires-Dist: pandas
+Requires-Dist: python-dateutil
 Requires-Dist: dataclasses_json
 Requires-Dist: tqdm
 Requires-Dist: click
+Requires-Dist: pydantic>=2.7
 Provides-Extra: remote
 Requires-Dist: unstructured-client>=0.26.1; extra == "remote"
 Provides-Extra: csv
@@ -71,18 +71,18 @@ Requires-Dist: fsspec; extra == "azure"
 Provides-Extra: azure-ai-search
 Requires-Dist: azure-search-documents; extra == "azure-ai-search"
 Provides-Extra: biomed
-Requires-Dist: bs4; extra == "biomed"
 Requires-Dist: requests; extra == "biomed"
+Requires-Dist: bs4; extra == "biomed"
 Provides-Extra: box
-Requires-Dist: boxfs; extra == "box"
 Requires-Dist: fsspec; extra == "box"
+Requires-Dist: boxfs; extra == "box"
 Provides-Extra: chroma
 Requires-Dist: chromadb; extra == "chroma"
 Provides-Extra: clarifai
 Requires-Dist: clarifai; extra == "clarifai"
 Provides-Extra: confluence
-Requires-Dist: atlassian-python-api; extra == "confluence"
 Requires-Dist: requests; extra == "confluence"
+Requires-Dist: atlassian-python-api; extra == "confluence"
 Provides-Extra: couchbase
 Requires-Dist: couchbase; extra == "couchbase"
 Provides-Extra: delta-table
@@ -91,8 +91,10 @@ Requires-Dist: boto3; extra == "delta-table"
 Provides-Extra: discord
 Requires-Dist: discord-py; extra == "discord"
 Provides-Extra: dropbox
-Requires-Dist: fsspec; extra == "dropbox"
 Requires-Dist: dropboxdrivefs; extra == "dropbox"
+Requires-Dist: fsspec; extra == "dropbox"
+Provides-Extra: duckdb
+Requires-Dist: duckdb; extra == "duckdb"
 Provides-Extra: elasticsearch
 Requires-Dist: elasticsearch[async]; extra == "elasticsearch"
 Provides-Extra: gcs
@@ -107,8 +109,8 @@ Requires-Dist: python-gitlab; extra == "gitlab"
 Provides-Extra: google-drive
 Requires-Dist: google-api-python-client; extra == "google-drive"
 Provides-Extra: hubspot
-Requires-Dist: urllib3; extra == "hubspot"
 Requires-Dist: hubspot-api-client; extra == "hubspot"
+Requires-Dist: urllib3; extra == "hubspot"
 Provides-Extra: jira
 Requires-Dist: atlassian-python-api; extra == "jira"
 Provides-Extra: kafka
@@ -172,8 +174,8 @@ Requires-Dist: requests; extra == "vectara"
 Provides-Extra: embed-huggingface
 Requires-Dist: sentence-transformers; extra == "embed-huggingface"
 Provides-Extra: embed-octoai
-Requires-Dist: tiktoken; extra == "embed-octoai"
 Requires-Dist: openai; extra == "embed-octoai"
+Requires-Dist: tiktoken; extra == "embed-octoai"
 Provides-Extra: embed-vertexai
 Requires-Dist: vertexai; extra == "embed-vertexai"
 Provides-Extra: embed-voyageai
@@ -181,8 +183,8 @@ Requires-Dist: voyageai; extra == "embed-voyageai"
 Provides-Extra: embed-mixedbreadai
 Requires-Dist: mixedbread-ai; extra == "embed-mixedbreadai"
 Provides-Extra: openai
-Requires-Dist: tiktoken; extra == "openai"
 Requires-Dist: openai; extra == "openai"
+Requires-Dist: tiktoken; extra == "openai"
 Provides-Extra: bedrock
 Requires-Dist: boto3; extra == "bedrock"
 Provides-Extra: togetherai

{unstructured_ingest-0.3.6 → unstructured_ingest-0.3.8}/setup.py

@@ -95,6 +95,7 @@ connectors_reqs = {
     "delta-table": load_requirements("requirements/connectors/delta-table.in"),
     "discord": load_requirements("requirements/connectors/discord.in"),
     "dropbox": load_requirements("requirements/connectors/dropbox.in"),
+    "duckdb": load_requirements("requirements/connectors/duckdb.in"),
     "elasticsearch": load_requirements("requirements/connectors/elasticsearch.in"),
     "gcs": load_requirements("requirements/connectors/gcs.in"),
     "github": load_requirements("requirements/connectors/github.in"),

unstructured_ingest-0.3.8/test/integration/connectors/duckdb/test_duckdb.py

@@ -0,0 +1,82 @@
+import tempfile
+from contextlib import contextmanager
+from pathlib import Path
+from typing import Generator
+
+import duckdb
+import pandas as pd
+import pytest
+
+from test.integration.connectors.utils.constants import DESTINATION_TAG
+from unstructured_ingest.v2.interfaces.file_data import FileData, SourceIdentifiers
+from unstructured_ingest.v2.processes.connectors.duckdb.duckdb import (
+    CONNECTOR_TYPE,
+    DuckDBConnectionConfig,
+    DuckDBUploader,
+    DuckDBUploaderConfig,
+    DuckDBUploadStager,
+)
+
+
+@contextmanager
+def duckdbd_setup() -> Generator[Path, None, None]:
+    with tempfile.TemporaryDirectory() as temp_dir:
+        db_path = Path(temp_dir) / "temp_duck.db"
+        db_init_path = Path(__file__).parent / "duckdb-schema.sql"
+        assert db_init_path.exists()
+        assert db_init_path.is_file()
+        with duckdb.connect(database=db_path) as duckdb_connection:
+            with db_init_path.open("r") as f:
+                query = f.read()
+            duckdb_connection.execute(query)
+            duckdb_connection.close()
+        yield db_path
+
+
+def validate_duckdb_destination(db_path: Path, expected_num_elements: int):
+    conn = None
+    try:
+        conn = duckdb.connect(db_path)
+        _results = conn.sql("select count(*) from elements").fetchall()
+        _count = _results[0][0]
+        assert (
+            _count == expected_num_elements
+        ), f"dest check failed: got {_count}, expected {expected_num_elements}"
+        conn.close()
+    finally:
+        if conn:
+            conn.close()
+
+
+@pytest.mark.tags(CONNECTOR_TYPE, DESTINATION_TAG, "duckdb")
+def test_duckdb_destination(upload_file: Path):
+    with duckdbd_setup() as test_db_path:
+        with tempfile.TemporaryDirectory() as temp_dir:
+            file_data = FileData(
+                source_identifiers=SourceIdentifiers(
+                    fullpath=upload_file.name, filename=upload_file.name
+                ),
+                connector_type=CONNECTOR_TYPE,
+                identifier="mock-file-data",
+            )
+
+            # deafults to default stager config
+            stager = DuckDBUploadStager()
+            stager_params = {
+                "elements_filepath": upload_file,
+                "file_data": file_data,
+                "output_dir": temp_dir,
+                "output_filename": "test_db",
+            }
+            staged_path = stager.run(**stager_params)
+
+            connection_config = DuckDBConnectionConfig(database=str(test_db_path))
+            upload_config = DuckDBUploaderConfig()
+            uploader = DuckDBUploader(
+                connection_config=connection_config, upload_config=upload_config
+            )
+
+            uploader.run(path=staged_path, file_data=file_data)
+
+            staged_df = pd.read_json(staged_path, orient="records", lines=True)
+            validate_duckdb_destination(db_path=test_db_path, expected_num_elements=len(staged_df))

unstructured_ingest-0.3.8/test/integration/connectors/duckdb/test_motherduck.py

@@ -0,0 +1,106 @@
+import os
+import tempfile
+import uuid
+from contextlib import contextmanager
+from pathlib import Path
+from typing import Generator
+
+import duckdb
+import pandas as pd
+import pytest
+
+from test.integration.connectors.utils.constants import DESTINATION_TAG
+from test.integration.utils import requires_env
+from unstructured_ingest.v2.interfaces.file_data import FileData, SourceIdentifiers
+from unstructured_ingest.v2.processes.connectors.duckdb.motherduck import (
+    CONNECTOR_TYPE,
+    MotherDuckAccessConfig,
+    MotherDuckConnectionConfig,
+    MotherDuckUploader,
+    MotherDuckUploaderConfig,
+    MotherDuckUploadStager,
+)
+
+
+@contextmanager
+def motherduck_setup(md_token: str) -> Generator[Path, None, None]:
+    database_name = f"test_{str(uuid.uuid4()).replace('-', '_')}"
+    try:
+        db_init_path = Path(__file__).parent / "duckdb-schema.sql"
+        assert db_init_path.exists()
+        assert db_init_path.is_file()
+        with duckdb.connect(f"md:?motherduck_token={md_token}") as md_conn:
+            with db_init_path.open("r") as f:
+                query = f.read()
+            md_conn.execute(f"CREATE DATABASE {database_name}")
+            md_conn.execute(f"USE {database_name}")
+            md_conn.execute(query)
+            md_conn.close()
+        yield database_name
+    finally:
+        with duckdb.connect(f"md:?motherduck_token={md_token}") as md_conn:
+            md_conn.execute(f"DROP DATABASE {database_name}")
+            md_conn.close()
+
+
+def validate_motherduck_destination(database: str, expected_num_elements: int, md_token: str):
+    conn = None
+    try:
+        conn = duckdb.connect(f"md:?motherduck_token={md_token}")
+        conn.execute(f"USE {database}")
+        _results = conn.sql("select count(*) from elements").fetchall()
+        _count = _results[0][0]
+        assert (
+            _count == expected_num_elements
+        ), f"dest check failed: got {_count}, expected {expected_num_elements}"
+        conn.close()
+    finally:
+        if conn:
+            conn.close()
+
+
+def get_motherduck_token() -> dict:
+    motherduck_token = os.getenv("MOTHERDUCK_TOKEN", None)
+    assert motherduck_token
+    return motherduck_token
+
+
+@pytest.mark.tags(CONNECTOR_TYPE, DESTINATION_TAG, "motherduck")
+@requires_env("MOTHERDUCK_TOKEN")
+def test_motherduck_destination(upload_file: Path):
+    md_token = get_motherduck_token()
+    with motherduck_setup(md_token) as test_database:
+        with tempfile.TemporaryDirectory() as temp_dir:
+            file_data = FileData(
+                source_identifiers=SourceIdentifiers(
+                    fullpath=upload_file.name, filename=upload_file.name
+                ),
+                connector_type=CONNECTOR_TYPE,
+                identifier="mock-file-data",
+            )
+
+            # deafults to default stager config
+            stager = MotherDuckUploadStager()
+            stager_params = {
+                "elements_filepath": upload_file,
+                "file_data": file_data,
+                "output_dir": temp_dir,
+                "output_filename": "test_db",
+            }
+            staged_path = stager.run(**stager_params)
+
+            access_config = MotherDuckAccessConfig(md_token=md_token)
+            connection_config = MotherDuckConnectionConfig(
+                database=test_database, access_config=access_config
+            )
+            upload_config = MotherDuckUploaderConfig()
+            uploader = MotherDuckUploader(
+                connection_config=connection_config, upload_config=upload_config
+            )
+
+            uploader.run(path=staged_path, file_data=file_data)
+
+            staged_df = pd.read_json(staged_path, orient="records", lines=True)
+            validate_motherduck_destination(
+                database=test_database, expected_num_elements=len(staged_df), md_token=md_token
+            )

{unstructured_ingest-0.3.6 → unstructured_ingest-0.3.8}/test/integration/connectors/test_kafka.py

@@ -1,4 +1,5 @@
 import json
+import os
 import tempfile
 import time
 from pathlib import Path
@@ -17,8 +18,17 @@ from test.integration.connectors.utils.validation import (
     ValidationConfigs,
     source_connector_validation,
 )
+from test.integration.utils import requires_env
 from unstructured_ingest.error import DestinationConnectionError, SourceConnectionError
 from unstructured_ingest.v2.interfaces import FileData, SourceIdentifiers
+from unstructured_ingest.v2.processes.connectors.kafka.cloud import (
+    CloudKafkaAccessConfig,
+    CloudKafkaConnectionConfig,
+    CloudKafkaDownloader,
+    CloudKafkaDownloaderConfig,
+    CloudKafkaIndexer,
+    CloudKafkaIndexerConfig,
+)
 from unstructured_ingest.v2.processes.connectors.kafka.local import (
     CONNECTOR_TYPE,
     LocalKafkaConnectionConfig,
@@ -47,20 +57,27 @@ def docker_compose_ctx():
         yield ctx
 
 
-def wait_for_topic(topic: str, retries: int = 10, interval: int = 1):
-    admin_client = get_admin_client()
+def wait_for_topic(
+    topic: str,
+    retries: int = 10,
+    interval: int = 1,
+    exists: bool = True,
+    admin_client=None,
+):
+    if admin_client is None:
+        admin_client = get_admin_client()
     current_topics = admin_client.list_topics().topics
     attempts = 0
-    while topic not in current_topics and attempts < retries:
+    while (topic not in current_topics) == exists and attempts < retries:
         attempts += 1
         print(
-            "Attempt {}: Waiting for topic {} to exist in {}".format(
-                attempts, topic, ", ".join(current_topics)
+            "Attempt {}: Waiting for topic {} to {} exist. Current topics: [{}]".format(
+                attempts, topic, "" if exists else "not", ", ".join(current_topics)
             )
         )
         time.sleep(interval)
         current_topics = admin_client.list_topics().topics
-    if topic not in current_topics:
+    if (topic not in current_topics) == exists:
         raise TimeoutError(f"Timeout out waiting for topic {topic} to exist")
 
 
@@ -110,6 +127,92 @@ async def test_kafka_source_local(kafka_seed_topic: str):
         )
 
 
+@pytest.fixture
+def kafka_seed_topic_cloud(expected_messages: int = 5) -> int:
+    conf = {
+        "bootstrap.servers": os.environ["KAFKA_BOOTSTRAP_SERVER"],
+        "sasl.username": os.environ["KAFKA_API_KEY"],
+        "sasl.password": os.environ["KAFKA_SECRET"],
+        "sasl.mechanism": "PLAIN",
+        "security.protocol": "SASL_SSL",
+    }
+    admin_client = AdminClient(conf)
+    try:
+        res = admin_client.delete_topics([TOPIC], operation_timeout=10)
+        for topic, f in res.items():
+            f.result()
+            print(f"Topic {topic} removed")
+            wait_for_topic(TOPIC, 5, 1, False, admin_client)
+    except Exception:
+        pass
+
+    cluster_meta = admin_client.list_topics()
+    current_topics = [topic for topic in cluster_meta.topics if topic != "__consumer_offsets"]
+
+    assert TOPIC not in current_topics, f"Topic {TOPIC} shouldn't exist"
+
+    # Kafka Cloud allows to use replication_factor=1 only for Dedicated clusters.
+    topic_obj = NewTopic(TOPIC, num_partitions=1, replication_factor=3)
+
+    res = admin_client.create_topics([topic_obj], operation_timeout=10, validate_only=False)
+    for topic, f in res.items():
+        f.result()
+
+    producer = Producer(conf)
+    for i in range(expected_messages):
+        message = f"This is some text for message {i}"
+        producer.produce(topic=TOPIC, value=message)
+    producer.flush(timeout=10)
+    return expected_messages
+
+
+@pytest.mark.asyncio
+@pytest.mark.tags(CONNECTOR_TYPE, SOURCE_TAG)
+@requires_env("KAFKA_API_KEY", "KAFKA_SECRET", "KAFKA_BOOTSTRAP_SERVER")
+async def test_kafka_source_cloud(kafka_seed_topic_cloud: int):
+    """
+    In order to have this test succeed, you need to create cluster on Confluent Cloud,
+    and create the API key with admin privileges. By default, user account keys have it.
+    """
+
+    expected_messages = kafka_seed_topic_cloud
+
+    connection_config = CloudKafkaConnectionConfig(
+        bootstrap_server=os.environ["KAFKA_BOOTSTRAP_SERVER"],
+        port=9092,
+        access_config=CloudKafkaAccessConfig(
+            kafka_api_key=os.environ["KAFKA_API_KEY"],
+            secret=os.environ["KAFKA_SECRET"],
+        ),
+    )
+
+    with tempfile.TemporaryDirectory() as tempdir:
+        tempdir_path = Path(tempdir)
+        download_config = CloudKafkaDownloaderConfig(download_dir=tempdir_path)
+        indexer = CloudKafkaIndexer(
+            connection_config=connection_config,
+            index_config=CloudKafkaIndexerConfig(
+                topic=TOPIC,
+                num_messages_to_consume=expected_messages,
+            ),
+        )
+        downloader = CloudKafkaDownloader(
+            connection_config=connection_config, download_config=download_config
+        )
+        indexer.precheck()
+        await source_connector_validation(
+            indexer=indexer,
+            downloader=downloader,
+            configs=ValidationConfigs(
+                test_id="kafka",
+                exclude_fields_extend=["connector_type"],
+                expected_num_files=expected_messages,
+                validate_downloaded_files=True,
+                validate_file_data=True,
+            ),
+        )
+
+
 @pytest.mark.tags(CONNECTOR_TYPE, SOURCE_TAG)
 def test_kafka_source_local_precheck_fail_no_cluster():
     connection_config = LocalKafkaConnectionConfig(bootstrap_server="localhost", port=29092)

{unstructured_ingest-0.3.6 → unstructured_ingest-0.3.8}/test/integration/connectors/test_qdrant.py

@@ -1,4 +1,5 @@
 import json
+import os
 import uuid
 from contextlib import asynccontextmanager
 from pathlib import Path
@@ -9,7 +10,16 @@ from qdrant_client import AsyncQdrantClient
 
 from test.integration.connectors.utils.constants import DESTINATION_TAG
 from test.integration.connectors.utils.docker import container_context
+from test.integration.utils import requires_env
 from unstructured_ingest.v2.interfaces.file_data import FileData, SourceIdentifiers
+from unstructured_ingest.v2.processes.connectors.qdrant.cloud import (
+    CloudQdrantAccessConfig,
+    CloudQdrantConnectionConfig,
+    CloudQdrantUploader,
+    CloudQdrantUploaderConfig,
+    CloudQdrantUploadStager,
+    CloudQdrantUploadStagerConfig,
+)
 from unstructured_ingest.v2.processes.connectors.qdrant.local import (
     CONNECTOR_TYPE as LOCAL_CONNECTOR_TYPE,
 )
@@ -135,3 +145,48 @@ async def test_qdrant_destination_server(upload_file: Path, tmp_path: Path, dock
         uploader.run(path=upload_file, file_data=file_data)
     async with qdrant_client(connection_kwargs) as client:
         await validate_upload(client=client, upload_file=upload_file)
+
+
+@pytest.mark.asyncio
+@pytest.mark.tags(SERVER_CONNECTOR_TYPE, DESTINATION_TAG, "qdrant")
+@requires_env("QDRANT_API_KEY", "QDRANT_SERVER_URL")
+async def test_qdrant_destination_cloud(upload_file: Path, tmp_path: Path):
+    server_url = os.environ["QDRANT_SERVER_URL"]
+    api_key = os.environ["QDRANT_API_KEY"]
+    connection_kwargs = {"location": server_url, "api_key": api_key}
+    async with qdrant_client(connection_kwargs) as client:
+        await client.create_collection(COLLECTION_NAME, vectors_config=VECTORS_CONFIG)
+    AsyncQdrantClient(**connection_kwargs)
+
+    stager = CloudQdrantUploadStager(
+        upload_stager_config=CloudQdrantUploadStagerConfig(),
+    )
+    uploader = CloudQdrantUploader(
+        connection_config=CloudQdrantConnectionConfig(
+            url=server_url,
+            access_config=CloudQdrantAccessConfig(
+                api_key=api_key,
+            ),
+        ),
+        upload_config=CloudQdrantUploaderConfig(collection_name=COLLECTION_NAME),
+    )
+
+    file_data = FileData(
+        source_identifiers=SourceIdentifiers(fullpath=upload_file.name, filename=upload_file.name),
+        connector_type=SERVER_CONNECTOR_TYPE,
+        identifier="mock-file-data",
+    )
+
+    staged_upload_file = stager.run(
+        elements_filepath=upload_file,
+        file_data=file_data,
+        output_dir=tmp_path,
+        output_filename=upload_file.name,
+    )
+
+    if uploader.is_async():
+        await uploader.run_async(path=staged_upload_file, file_data=file_data)
+    else:
+        uploader.run(path=staged_upload_file, file_data=file_data)
+    async with qdrant_client(connection_kwargs) as client:
+        await validate_upload(client=client, upload_file=upload_file)

unstructured_ingest-0.3.8/test/unit/v2/connectors/test_confluence.py

@@ -0,0 +1,39 @@
+import pytest
+from pydantic import ValidationError
+
+from unstructured_ingest.v2.processes.connectors.confluence import (
+    ConfluenceAccessConfig,
+    ConfluenceConnectionConfig,
+)
+
+
+def test_connection_config_multiple_auth():
+    with pytest.raises(ValidationError):
+        ConfluenceConnectionConfig(
+            access_config=ConfluenceAccessConfig(
+                api_token="api_token",
+                access_token="access_token",
+            ),
+            user_email="user_email",
+            url="url",
+        )
+
+
+def test_connection_config_no_auth():
+    with pytest.raises(ValidationError):
+        ConfluenceConnectionConfig(access_config=ConfluenceAccessConfig(), url="url")
+
+
+def test_connection_config_basic_auth():
+    ConfluenceConnectionConfig(
+        access_config=ConfluenceAccessConfig(api_token="api_token"),
+        url="url",
+        user_email="user_email",
+    )
+
+
+def test_connection_config_pat_auth():
+    ConfluenceConnectionConfig(
+        access_config=ConfluenceAccessConfig(access_token="access_token"),
+        url="url",
+    )

unstructured_ingest-0.3.8/unstructured_ingest/__version__.py

@@ -0,0 +1 @@
+__version__ = "0.3.8"  # pragma: no cover

{unstructured_ingest-0.3.6 → unstructured_ingest-0.3.8}/unstructured_ingest/v2/processes/connectors/__init__.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import unstructured_ingest.v2.processes.connectors.databricks  # noqa: F401
+import unstructured_ingest.v2.processes.connectors.duckdb  # noqa: F401
 import unstructured_ingest.v2.processes.connectors.elasticsearch  # noqa: F401
 import unstructured_ingest.v2.processes.connectors.fsspec  # noqa: F401
 import unstructured_ingest.v2.processes.connectors.kafka  # noqa: F401

{unstructured_ingest-0.3.6 → unstructured_ingest-0.3.8}/unstructured_ingest/v2/processes/connectors/azure_ai_search.py

@@ -1,5 +1,4 @@
 import json
-import uuid
 from dataclasses import dataclass, field
 from pathlib import Path
 from typing import TYPE_CHECKING, Any
@@ -24,6 +23,7 @@ from unstructured_ingest.v2.processes.connector_registry import (
     DestinationRegistryEntry,
 )
 from unstructured_ingest.v2.processes.connectors.utils import parse_datetime
+from unstructured_ingest.v2.utils import get_enhanced_element_id
 
 if TYPE_CHECKING:
     from azure.search.documents import SearchClient
@@ -100,7 +100,7 @@ class AzureAISearchUploadStager(UploadStager):
         Azure Cognitive Search index
         """
 
-        data["id"] = str(uuid.uuid4())
+        data["id"] = get_enhanced_element_id(element_dict=data, file_data=file_data)
         data[RECORD_ID_LABEL] = file_data.identifier
 
         if points := data.get("metadata", {}).get("coordinates", {}).get("points"):
@@ -173,8 +173,10 @@ class AzureAISearchUploader(Uploader):
     connector_type: str = CONNECTOR_TYPE
 
     def query_docs(self, record_id: str, index_key: str) -> list[str]:
-        client = self.connection_config.get_search_client()
-        results = list(client.search(filter=f"record_id eq '{record_id}'", select=[index_key]))
+        with self.connection_config.get_search_client() as search_client:
+            results = list(
+                search_client.search(filter=f"record_id eq '{record_id}'", select=[index_key])
+            )
         return [result[index_key] for result in results]
 
     def delete_by_record_id(self, file_data: FileData, index_key: str) -> None:
@@ -186,10 +188,10 @@ class AzureAISearchUploader(Uploader):
         doc_ids_to_delete = self.query_docs(record_id=file_data.identifier, index_key=index_key)
         if not doc_ids_to_delete:
             return
-        client: SearchClient = self.connection_config.get_search_client()
-        results = client.delete_documents(
-            documents=[{index_key: doc_id} for doc_id in doc_ids_to_delete]
-        )
+        with self.connection_config.get_search_client() as search_client:
+            results = search_client.delete_documents(
+                documents=[{index_key: doc_id} for doc_id in doc_ids_to_delete]
+            )
         errors = []
         success = []
         for result in results:
@@ -207,7 +209,9 @@ class AzureAISearchUploader(Uploader):
 
     @DestinationConnectionError.wrap
     @requires_dependencies(["azure"], extras="azure-ai-search")
-    def write_dict(self, elements_dict: list[dict[str, Any]]) -> None:
+    def write_dict(
+        self, elements_dict: list[dict[str, Any]], search_client: "SearchClient"
+    ) -> None:
         import azure.core.exceptions
 
         logger.info(
@@ -215,12 +219,10 @@ class AzureAISearchUploader(Uploader):
             f"index at {self.connection_config.index}",
         )
         try:
-            results = self.connection_config.get_search_client().upload_documents(
-                documents=elements_dict
-            )
-
+            results = search_client.upload_documents(documents=elements_dict)
         except azure.core.exceptions.HttpResponseError as http_error:
             raise WriteError(f"http error: {http_error}") from http_error
+
         errors = []
         success = []
         for result in results:
@@ -240,8 +242,8 @@ class AzureAISearchUploader(Uploader):
             )
 
     def can_delete(self) -> bool:
-        search_index_client = self.connection_config.get_search_index_client()
-        index = search_index_client.get_index(name=self.connection_config.index)
+        with self.connection_config.get_search_index_client() as search_index_client:
+            index = search_index_client.get_index(name=self.connection_config.index)
         index_fields = index.fields
         record_id_fields = [
             field for field in index_fields if field.name == self.upload_config.record_id_key
@@ -252,8 +254,8 @@ class AzureAISearchUploader(Uploader):
         return record_id_field.filterable
 
     def get_index_key(self) -> str:
-        search_index_client = self.connection_config.get_search_index_client()
-        index = search_index_client.get_index(name=self.connection_config.index)
+        with self.connection_config.get_search_index_client() as search_index_client:
+            index = search_index_client.get_index(name=self.connection_config.index)
         index_fields = index.fields
         key_fields = [field for field in index_fields if field.key]
         if not key_fields:
@@ -262,8 +264,8 @@ class AzureAISearchUploader(Uploader):
 
     def precheck(self) -> None:
         try:
-            client = self.connection_config.get_search_client()
-            client.get_document_count()
+            with self.connection_config.get_search_client() as search_client:
+                search_client.get_document_count()
         except Exception as e:
             logger.error(f"failed to validate connection: {e}", exc_info=True)
             raise DestinationConnectionError(f"failed to validate connection: {e}")
@@ -284,8 +286,9 @@ class AzureAISearchUploader(Uploader):
             logger.warning("criteria for deleting previous content not met, skipping")
 
         batch_size = self.upload_config.batch_size
-        for chunk in batch_generator(elements_dict, batch_size):
-            self.write_dict(elements_dict=chunk)  # noqa: E203
+        with self.connection_config.get_search_client() as search_client:
+            for chunk in batch_generator(elements_dict, batch_size):
+                self.write_dict(elements_dict=chunk, search_client=search_client)  # noqa: E203
 
 
 azure_ai_search_destination_entry = DestinationRegistryEntry(