unstructured-ingest 0.2.2__py3-none-any.whl → 0.3.0__py3-none-any.whl

unstructured_ingest/v2/processes/connectors/kafka/cloud.py

@@ -4,7 +4,10 @@ from typing import TYPE_CHECKING, Optional
 
 from pydantic import Field, Secret, SecretStr
 
-from unstructured_ingest.v2.processes.connector_registry import SourceRegistryEntry
+from unstructured_ingest.v2.processes.connector_registry import (
+    DestinationRegistryEntry,
+    SourceRegistryEntry,
+)
 from unstructured_ingest.v2.processes.connectors.kafka.kafka import (
     KafkaAccessConfig,
     KafkaConnectionConfig,
@@ -12,6 +15,8 @@ from unstructured_ingest.v2.processes.connectors.kafka.kafka import (
     KafkaDownloaderConfig,
     KafkaIndexer,
     KafkaIndexerConfig,
+    KafkaUploader,
+    KafkaUploaderConfig,
 )
 
 if TYPE_CHECKING:
@@ -41,7 +46,21 @@ class CloudKafkaConnectionConfig(KafkaConnectionConfig):
             "group.id": "default_group_id",
             "enable.auto.commit": "false",
             "auto.offset.reset": "earliest",
-            "message.max.bytes": 10485760,
+            "sasl.username": access_config.api_key,
+            "sasl.password": access_config.secret,
+            "sasl.mechanism": "PLAIN",
+            "security.protocol": "SASL_SSL",
+        }
+
+        return conf
+
+    def get_producer_configuration(self) -> dict:
+        bootstrap = self.bootstrap_server
+        port = self.port
+        access_config = self.access_config.get_secret_value()
+
+        conf = {
+            "bootstrap.servers": f"{bootstrap}:{port}",
             "sasl.username": access_config.api_key,
             "sasl.password": access_config.secret,
             "sasl.mechanism": "PLAIN",
@@ -73,6 +92,17 @@ class CloudKafkaDownloader(KafkaDownloader):
     connector_type: str = CONNECTOR_TYPE
 
 
+class CloudKafkaUploaderConfig(KafkaUploaderConfig):
+    pass
+
+
+@dataclass
+class CloudKafkaUploader(KafkaUploader):
+    connection_config: CloudKafkaConnectionConfig
+    upload_config: CloudKafkaUploaderConfig
+    connector_type: str = CONNECTOR_TYPE
+
+
 kafka_cloud_source_entry = SourceRegistryEntry(
     connection_config=CloudKafkaConnectionConfig,
     indexer=CloudKafkaIndexer,
@@ -80,3 +110,9 @@ kafka_cloud_source_entry = SourceRegistryEntry(
     downloader=CloudKafkaDownloader,
     downloader_config=CloudKafkaDownloaderConfig,
 )
+
+kafka_cloud_destination_entry = DestinationRegistryEntry(
+    connection_config=CloudKafkaConnectionConfig,
+    uploader=CloudKafkaUploader,
+    uploader_config=CloudKafkaUploaderConfig,
+)

unstructured_ingest/v2/processes/connectors/kafka/kafka.py

@@ -1,3 +1,4 @@
+import json
 from abc import ABC, abstractmethod
 from contextlib import contextmanager
 from dataclasses import dataclass, field
@@ -5,32 +6,33 @@ from pathlib import Path
 from time import time
 from typing import TYPE_CHECKING, Any, ContextManager, Generator, Optional
 
-from pydantic import Secret
+from pydantic import Field, Secret
 
 from unstructured_ingest.error import (
+    DestinationConnectionError,
     SourceConnectionError,
     SourceConnectionNetworkError,
 )
+from unstructured_ingest.utils.data_prep import batch_generator
 from unstructured_ingest.utils.dep_check import requires_dependencies
 from unstructured_ingest.v2.interfaces import (
     AccessConfig,
     ConnectionConfig,
     Downloader,
     DownloaderConfig,
+    DownloadResponse,
     FileData,
     FileDataSourceMetadata,
     Indexer,
     IndexerConfig,
     SourceIdentifiers,
-    download_responses,
+    Uploader,
+    UploaderConfig,
 )
 from unstructured_ingest.v2.logger import logger
-from unstructured_ingest.v2.processes.connector_registry import SourceRegistryEntry
 
 if TYPE_CHECKING:
-    from confluent_kafka import Consumer
-
-CONNECTOR_TYPE = "kafka"
+    from confluent_kafka import Consumer, Producer
 
 
 class KafkaAccessConfig(AccessConfig, ABC):
@@ -39,7 +41,6 @@ class KafkaAccessConfig(AccessConfig, ABC):
 
 class KafkaConnectionConfig(ConnectionConfig, ABC):
     access_config: Secret[KafkaAccessConfig]
-    timeout: Optional[float] = 1.0
     bootstrap_server: str
     port: int
 
@@ -47,6 +48,10 @@ class KafkaConnectionConfig(ConnectionConfig, ABC):
     def get_consumer_configuration(self) -> dict:
         pass
 
+    @abstractmethod
+    def get_producer_configuration(self) -> dict:
+        pass
+
     @contextmanager
     @requires_dependencies(["confluent_kafka"], extras="kafka")
     def get_consumer(self) -> ContextManager["Consumer"]:
@@ -59,20 +64,27 @@ class KafkaConnectionConfig(ConnectionConfig, ABC):
         finally:
             consumer.close()
 
+    @requires_dependencies(["confluent_kafka"], extras="kafka")
+    def get_producer(self) -> "Producer":
+        from confluent_kafka import Producer
+
+        producer = Producer(self.get_producer_configuration())
+        return producer
+
 
 class KafkaIndexerConfig(IndexerConfig):
-    topic: str
+    topic: str = Field(description="which topic to consume from")
     num_messages_to_consume: Optional[int] = 100
+    timeout: Optional[float] = Field(default=1.0, description="polling timeout")
 
     def update_consumer(self, consumer: "Consumer") -> None:
         consumer.subscribe([self.topic])
 
 
 @dataclass
-class KafkaIndexer(Indexer):
+class KafkaIndexer(Indexer, ABC):
     connection_config: KafkaConnectionConfig
     index_config: KafkaIndexerConfig
-    connector_type: str = CONNECTOR_TYPE
 
     @contextmanager
     def get_consumer(self) -> ContextManager["Consumer"]:
@@ -90,7 +102,7 @@ class KafkaIndexer(Indexer):
         num_messages_to_consume = self.index_config.num_messages_to_consume
         with self.get_consumer() as consumer:
             while messages_consumed < num_messages_to_consume and empty_polls < max_empty_polls:
-                msg = consumer.poll(timeout=self.connection_config.timeout)
+                msg = consumer.poll(timeout=self.index_config.timeout)
                 if msg is None:
                     logger.debug("No Kafka messages found")
                     empty_polls += 1
@@ -139,13 +151,13 @@ class KafkaIndexer(Indexer):
         for message in self.generate_messages():
             yield self.generate_file_data(message)
 
-    async def run_async(self, file_data: FileData, **kwargs: Any) -> download_responses:
+    async def run_async(self, file_data: FileData, **kwargs: Any) -> DownloadResponse:
         raise NotImplementedError()
 
     def precheck(self):
         try:
             with self.get_consumer() as consumer:
-                cluster_meta = consumer.list_topics(timeout=self.connection_config.timeout)
+                cluster_meta = consumer.list_topics(timeout=self.index_config.timeout)
                 current_topics = [
                     topic for topic in cluster_meta.topics if topic != "__consumer_offsets"
                 ]
@@ -160,14 +172,13 @@ class KafkaDownloaderConfig(DownloaderConfig):
 
 
 @dataclass
-class KafkaDownloader(Downloader):
+class KafkaDownloader(Downloader, ABC):
     connection_config: KafkaConnectionConfig
     download_config: KafkaDownloaderConfig = field(default_factory=KafkaDownloaderConfig)
-    connector_type: str = CONNECTOR_TYPE
     version: Optional[str] = None
     source_url: Optional[str] = None
 
-    def run(self, file_data: FileData, **kwargs: Any) -> download_responses:
+    def run(self, file_data: FileData, **kwargs: Any) -> DownloadResponse:
         source_identifiers = file_data.source_identifiers
         if source_identifiers is None:
             raise ValueError("FileData is missing source_identifiers")
@@ -187,10 +198,54 @@ class KafkaDownloader(Downloader):
         return self.generate_download_response(file_data=file_data, download_path=download_path)
 
 
-kafka_source_entry = SourceRegistryEntry(
-    connection_config=KafkaConnectionConfig,
-    indexer=KafkaIndexer,
-    indexer_config=KafkaIndexerConfig,
-    downloader=KafkaDownloader,
-    downloader_config=KafkaDownloaderConfig,
-)
+class KafkaUploaderConfig(UploaderConfig):
+    batch_size: int = Field(default=100, description="Batch size")
+    topic: str = Field(description="which topic to write to")
+    timeout: Optional[float] = Field(
+        default=10.0, description="Timeout in seconds to flush batch of messages"
+    )
+
+
+@dataclass
+class KafkaUploader(Uploader, ABC):
+    connection_config: KafkaConnectionConfig
+    upload_config: KafkaUploaderConfig
+
+    def precheck(self):
+        try:
+            with self.connection_config.get_consumer() as consumer:
+                cluster_meta = consumer.list_topics(timeout=self.upload_config.timeout)
+                current_topics = [
+                    topic for topic in cluster_meta.topics if topic != "__consumer_offsets"
+                ]
+                logger.info(f"successfully checked available topics: {current_topics}")
+        except Exception as e:
+            logger.error(f"failed to validate connection: {e}", exc_info=True)
+            raise DestinationConnectionError(f"failed to validate connection: {e}")
+
+    def produce_batch(self, elements: list[dict]) -> None:
+        from confluent_kafka.error import KafkaException
+
+        producer = self.connection_config.get_producer()
+        failed_producer = False
+
+        def acked(err, msg):
+            if err is not None:
+                logger.error("Failed to deliver message: %s: %s" % (str(msg), str(err)))
+
+        for element in elements:
+            producer.produce(
+                topic=self.upload_config.topic,
+                value=json.dumps(element),
+                callback=acked,
+            )
+
+        producer.flush(timeout=self.upload_config.timeout)
+        if failed_producer:
+            raise KafkaException("failed to produce all messages in batch")
+
+    def run(self, path: Path, file_data: FileData, **kwargs: Any) -> None:
+        with path.open("r") as elements_file:
+            elements = json.load(elements_file)
+        for element_batch in batch_generator(elements, batch_size=self.upload_config.batch_size):
+            self.produce_batch(elements=element_batch)

unstructured_ingest/v2/processes/connectors/kafka/local.py

@@ -1,10 +1,12 @@
-import socket
 from dataclasses import dataclass
 from typing import TYPE_CHECKING
 
 from pydantic import Field, Secret
 
-from unstructured_ingest.v2.processes.connector_registry import SourceRegistryEntry
+from unstructured_ingest.v2.processes.connector_registry import (
+    DestinationRegistryEntry,
+    SourceRegistryEntry,
+)
 from unstructured_ingest.v2.processes.connectors.kafka.kafka import (
     KafkaAccessConfig,
     KafkaConnectionConfig,
@@ -12,6 +14,8 @@ from unstructured_ingest.v2.processes.connectors.kafka.kafka import (
     KafkaDownloaderConfig,
     KafkaIndexer,
     KafkaIndexerConfig,
+    KafkaUploader,
+    KafkaUploaderConfig,
 )
 
 if TYPE_CHECKING:
@@ -35,11 +39,18 @@ class LocalKafkaConnectionConfig(KafkaConnectionConfig):
 
         conf = {
             "bootstrap.servers": f"{bootstrap}:{port}",
-            "client.id": socket.gethostname(),
             "group.id": "default_group_id",
             "enable.auto.commit": "false",
             "auto.offset.reset": "earliest",
-            "message.max.bytes": 10485760,
+        }
+        return conf
+
+    def get_producer_configuration(self) -> dict:
+        bootstrap = self.bootstrap_server
+        port = self.port
+
+        conf = {
+            "bootstrap.servers": f"{bootstrap}:{port}",
         }
         return conf
 
@@ -66,6 +77,17 @@ class LocalKafkaDownloader(KafkaDownloader):
     connector_type: str = CONNECTOR_TYPE
 
 
+class LocalKafkaUploaderConfig(KafkaUploaderConfig):
+    pass
+
+
+@dataclass
+class LocalKafkaUploader(KafkaUploader):
+    connection_config: LocalKafkaConnectionConfig
+    upload_config: LocalKafkaUploaderConfig
+    connector_type: str = CONNECTOR_TYPE
+
+
 kafka_local_source_entry = SourceRegistryEntry(
     connection_config=LocalKafkaConnectionConfig,
     indexer=LocalKafkaIndexer,
@@ -73,3 +95,9 @@ kafka_local_source_entry = SourceRegistryEntry(
     downloader=LocalKafkaDownloader,
     downloader_config=LocalKafkaDownloaderConfig,
 )
+
+kafka_local_destination_entry = DestinationRegistryEntry(
+    connection_config=LocalKafkaConnectionConfig,
+    uploader=LocalKafkaUploader,
+    uploader_config=LocalKafkaUploaderConfig,
+)

unstructured_ingest/v2/processes/connectors/onedrive.py

@@ -28,7 +28,6 @@ from unstructured_ingest.v2.interfaces import (
     SourceIdentifiers,
     Uploader,
     UploaderConfig,
-    download_responses,
 )
 from unstructured_ingest.v2.logger import logger
 from unstructured_ingest.v2.processes.connector_registry import (
@@ -220,7 +219,7 @@ class OnedriveDownloader(Downloader):
         return self.download_dir / Path(rel_path)
 
     @SourceConnectionError.wrap
-    def run(self, file_data: FileData, **kwargs: Any) -> download_responses:
+    def run(self, file_data: FileData, **kwargs: Any) -> DownloadResponse:
         file = self._fetch_file(file_data=file_data)
         fsize = file.get_property("size", 0)
         download_path = self.get_download_path(file_data=file_data)
@@ -233,7 +232,7 @@ class OnedriveDownloader(Downloader):
         else:
             with download_path.open(mode="wb") as f:
                 file.download(f).execute_query()
-        return DownloadResponse(file_data=file_data, path=download_path)
+        return self.generate_download_response(file_data=file_data, download_path=download_path)
 
 
 class OnedriveUploaderConfig(UploaderConfig):

unstructured_ingest/v2/processes/connectors/outlook.py

@@ -15,10 +15,10 @@ from unstructured_ingest.v2.interfaces import (
     ConnectionConfig,
     Downloader,
     DownloaderConfig,
+    DownloadResponse,
     FileData,
     Indexer,
     IndexerConfig,
-    download_responses,
 )
 from unstructured_ingest.v2.interfaces.file_data import FileDataSourceMetadata, SourceIdentifiers
 from unstructured_ingest.v2.processes.connector_registry import SourceRegistryEntry
@@ -191,7 +191,7 @@ class OutlookDownloader(Downloader):
     connection_config: OutlookConnectionConfig
     download_config: OutlookDownloaderConfig = field(default_factory=OutlookDownloaderConfig)
 
-    def run(self, file_data: FileData, **kwargs: Any) -> download_responses:
+    def run(self, file_data: FileData, **kwargs: Any) -> DownloadResponse:
         # NOTE: Indexer should provide source identifiers required to generate the download path
         download_path = self.get_download_path(file_data)
         if download_path is None:

unstructured_ingest/v2/processes/connectors/pinecone.py

@@ -9,6 +9,7 @@ from pydantic import Field, Secret
 from unstructured_ingest.error import DestinationConnectionError
 from unstructured_ingest.utils.data_prep import flatten_dict, generator_batching_wbytes
 from unstructured_ingest.utils.dep_check import requires_dependencies
+from unstructured_ingest.v2.constants import RECORD_ID_LABEL
 from unstructured_ingest.v2.interfaces import (
     AccessConfig,
     ConnectionConfig,
@@ -23,6 +24,7 @@ from unstructured_ingest.v2.processes.connector_registry import DestinationRegis
 
 if TYPE_CHECKING:
     from pinecone import Index as PineconeIndex
+    from pinecone import Pinecone
 
 
 CONNECTOR_TYPE = "pinecone"
@@ -43,16 +45,19 @@ class PineconeConnectionConfig(ConnectionConfig):
     )
 
     @requires_dependencies(["pinecone"], extras="pinecone")
-    def get_index(self, **index_kwargs) -> "PineconeIndex":
+    def get_client(self, **index_kwargs) -> "Pinecone":
         from pinecone import Pinecone
 
         from unstructured_ingest import __version__ as unstructured_version
 
-        pc = Pinecone(
+        return Pinecone(
             api_key=self.access_config.get_secret_value().pinecone_api_key,
             source_tag=f"unstructured_ingest=={unstructured_version}",
         )
 
+    def get_index(self, **index_kwargs) -> "PineconeIndex":
+        pc = self.get_client()
+
         index = pc.Index(name=self.index_name, **index_kwargs)
         logger.debug(f"connected to index: {pc.describe_index(self.index_name)}")
         return index
@@ -106,7 +111,7 @@ class PineconeUploadStager(UploadStager):
         default_factory=lambda: PineconeUploadStagerConfig()
     )
 
-    def conform_dict(self, element_dict: dict) -> dict:
+    def conform_dict(self, element_dict: dict, file_data: FileData) -> dict:
         embeddings = element_dict.pop("embeddings", None)
         metadata: dict[str, Any] = element_dict.pop("metadata", {})
         data_source = metadata.pop("data_source", {})
@@ -121,19 +126,23 @@ class PineconeUploadStager(UploadStager):
                 }
             )
 
+        metadata = flatten_dict(
+            pinecone_metadata,
+            separator="-",
+            flatten_lists=True,
+            remove_none=True,
+        )
+        metadata[RECORD_ID_LABEL] = file_data.identifier
+
         return {
             "id": str(uuid.uuid4()),
             "values": embeddings,
-            "metadata": flatten_dict(
-                pinecone_metadata,
-                separator="-",
-                flatten_lists=True,
-                remove_none=True,
-            ),
+            "metadata": metadata,
         }
 
     def run(
         self,
+        file_data: FileData,
         elements_filepath: Path,
         output_dir: Path,
         output_filename: str,
@@ -143,10 +152,15 @@ class PineconeUploadStager(UploadStager):
             elements_contents = json.load(elements_file)
 
         conformed_elements = [
-            self.conform_dict(element_dict=element) for element in elements_contents
+            self.conform_dict(element_dict=element, file_data=file_data)
+            for element in elements_contents
         ]
 
-        output_path = Path(output_dir) / Path(f"{output_filename}.json")
+        if Path(output_filename).suffix != ".json":
+            output_filename = f"{output_filename}.json"
+        else:
+            output_filename = f"{Path(output_filename).stem}.json"
+        output_path = Path(output_dir) / Path(f"{output_filename}")
         output_path.parent.mkdir(parents=True, exist_ok=True)
 
         with open(output_path, "w") as output_file:
@@ -167,6 +181,55 @@ class PineconeUploader(Uploader):
             logger.error(f"failed to validate connection: {e}", exc_info=True)
             raise DestinationConnectionError(f"failed to validate connection: {e}")
 
+    def pod_delete_by_record_id(self, file_data: FileData) -> None:
+        logger.debug(
+            f"deleting any content with metadata {RECORD_ID_LABEL}={file_data.identifier} "
+            f"from pinecone pod index"
+        )
+        index = self.connection_config.get_index(pool_threads=MAX_POOL_THREADS)
+        delete_kwargs = {"filter": {RECORD_ID_LABEL: {"$eq": file_data.identifier}}}
+        if namespace := self.upload_config.namespace:
+            delete_kwargs["namespace"] = namespace
+
+        resp = index.delete(**delete_kwargs)
+        logger.debug(
+            f"deleted any content with metadata {RECORD_ID_LABEL}={file_data.identifier} "
+            f"from pinecone index: {resp}"
+        )
+
+    def serverless_delete_by_record_id(self, file_data: FileData) -> None:
+        logger.debug(
+            f"deleting any content with metadata {RECORD_ID_LABEL}={file_data.identifier} "
+            f"from pinecone serverless index"
+        )
+        index = self.connection_config.get_index(pool_threads=MAX_POOL_THREADS)
+        index_stats = index.describe_index_stats()
+        total_vectors = index_stats["total_vector_count"]
+        if total_vectors == 0:
+            return
+        dimension = index_stats["dimension"]
+        query_params = {
+            "filter": {RECORD_ID_LABEL: {"$eq": file_data.identifier}},
+            "vector": [0] * dimension,
+            "top_k": total_vectors,
+        }
+        if namespace := self.upload_config.namespace:
+            query_params["namespace"] = namespace
+        while True:
+            query_results = index.query(**query_params)
+            matches = query_results.get("matches", [])
+            if not matches:
+                break
+            ids = [match["id"] for match in matches]
+            delete_params = {"ids": ids}
+            if namespace := self.upload_config.namespace:
+                delete_params["namespace"] = namespace
+            index.delete(**delete_params)
+        logger.debug(
+            f"deleted any content with metadata {RECORD_ID_LABEL}={file_data.identifier} "
+            f"from pinecone index"
+        )
+
     @requires_dependencies(["pinecone"], extras="pinecone")
     def upsert_batches_async(self, elements_dict: list[dict]):
         from pinecone.exceptions import PineconeApiException
@@ -208,7 +271,15 @@ class PineconeUploader(Uploader):
             f" index named {self.connection_config.index_name}"
             f" with batch size {self.upload_config.batch_size}"
         )
-
+        # Determine if serverless or pod based index
+        pinecone_client = self.connection_config.get_client()
+        index_description = pinecone_client.describe_index(name=self.connection_config.index_name)
+        if "serverless" in index_description.get("spec"):
+            self.serverless_delete_by_record_id(file_data=file_data)
+        elif "pod" in index_description.get("spec"):
+            self.pod_delete_by_record_id(file_data=file_data)
+        else:
+            raise ValueError(f"unexpected spec type in index description: {index_description}")
         self.upsert_batches_async(elements_dict=elements_dict)
 
 

unstructured_ingest/v2/processes/connectors/sharepoint.py

@@ -21,7 +21,6 @@ from unstructured_ingest.v2.interfaces import (
     Indexer,
     IndexerConfig,
     SourceIdentifiers,
-    download_responses,
 )
 from unstructured_ingest.v2.logger import logger
 from unstructured_ingest.v2.processes.connector_registry import (
@@ -426,7 +425,7 @@ class SharepointDownloader(Downloader):
             f.write(etree.tostring(document, encoding="unicode", pretty_print=True))
         return self.generate_download_response(file_data=file_data, download_path=download_path)
 
-    def run(self, file_data: FileData, **kwargs: Any) -> download_responses:
+    def run(self, file_data: FileData, **kwargs: Any) -> DownloadResponse:
         content_type = file_data.additional_metadata.get("sharepoint_content_type")
         if not content_type:
             raise ValueError(
@@ -436,6 +435,8 @@ class SharepointDownloader(Downloader):
             return self.get_document(file_data=file_data)
         elif content_type == SharepointContentType.SITEPAGE.value:
             return self.get_site_page(file_data=file_data)
+        else:
+            raise ValueError(f"content type not recognized: {content_type}")
 
 
 sharepoint_source_entry = SourceRegistryEntry(

unstructured_ingest/v2/processes/connectors/slack.py

@@ -16,9 +16,9 @@ from unstructured_ingest.v2.interfaces import (
     ConnectionConfig,
     Downloader,
     DownloaderConfig,
+    DownloadResponse,
     Indexer,
     IndexerConfig,
-    download_responses,
 )
 from unstructured_ingest.v2.interfaces.file_data import (
     FileData,
@@ -161,7 +161,7 @@ class SlackDownloader(Downloader):
     def run(self, file_data, **kwargs):
         raise NotImplementedError
 
-    async def run_async(self, file_data: FileData, **kwargs) -> download_responses:
+    async def run_async(self, file_data: FileData, **kwargs) -> DownloadResponse:
         # NOTE: Indexer should provide source identifiers required to generate the download path
         download_path = self.get_download_path(file_data)
         if download_path is None:

unstructured_ingest/v2/processes/connectors/sql/postgres.py

@@ -98,20 +98,28 @@ class PostgresDownloader(SQLDownloader):
     download_config: PostgresDownloaderConfig
     connector_type: str = CONNECTOR_TYPE
 
+    @requires_dependencies(["psycopg2"], extras="postgres")
     def query_db(self, file_data: FileData) -> tuple[list[tuple], list[str]]:
+        from psycopg2 import sql
+
         table_name = file_data.additional_metadata["table_name"]
         id_column = file_data.additional_metadata["id_column"]
-        ids = file_data.additional_metadata["ids"]
+        ids = tuple(file_data.additional_metadata["ids"])
+
         with self.connection_config.get_cursor() as cursor:
-            fields = ",".join(self.download_config.fields) if self.download_config.fields else "*"
-            query = "SELECT {fields} FROM {table_name} WHERE {id_column} in ({ids})".format(
+            fields = (
+                sql.SQL(",").join(sql.Identifier(field) for field in self.download_config.fields)
+                if self.download_config.fields
+                else sql.SQL("*")
+            )
+
+            query = sql.SQL("SELECT {fields} FROM {table_name} WHERE {id_column} IN %s").format(
                 fields=fields,
-                table_name=table_name,
-                id_column=id_column,
-                ids=",".join([str(i) for i in ids]),
+                table_name=sql.Identifier(table_name),
+                id_column=sql.Identifier(id_column),
             )
-            logger.debug(f"running query: {query}")
-            cursor.execute(query)
+            logger.debug(f"running query: {cursor.mogrify(query, (ids,))}")
+            cursor.execute(query, (ids,))
             rows = cursor.fetchall()
             columns = [col[0] for col in cursor.description]
             return rows, columns