unstructured-ingest 0.1.1__py3-none-any.whl → 0.2.1__py3-none-any.whl

unstructured_ingest/v2/processes/connectors/slack.py

@@ -0,0 +1,248 @@
+import hashlib
+import time
+import xml.etree.ElementTree as ET
+from dataclasses import dataclass, field
+from datetime import datetime
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Generator, Optional
+
+from pydantic import Field, Secret
+
+from unstructured_ingest.error import SourceConnectionError
+from unstructured_ingest.logger import logger
+from unstructured_ingest.utils.dep_check import requires_dependencies
+from unstructured_ingest.v2.interfaces import (
+    AccessConfig,
+    ConnectionConfig,
+    Downloader,
+    DownloaderConfig,
+    Indexer,
+    IndexerConfig,
+    download_responses,
+)
+from unstructured_ingest.v2.interfaces.file_data import (
+    FileData,
+    FileDataSourceMetadata,
+    SourceIdentifiers,
+)
+from unstructured_ingest.v2.processes.connector_registry import SourceRegistryEntry
+
+if TYPE_CHECKING:
+    from slack_sdk import WebClient
+    from slack_sdk.web.async_client import AsyncWebClient
+
+# NOTE: Pagination limit set to the upper end of the recommended range
+# https://api.slack.com/apis/pagination#facts
+PAGINATION_LIMIT = 200
+
+CONNECTOR_TYPE = "slack"
+
+
+class SlackAccessConfig(AccessConfig):
+    token: str = Field(
+        description="Bot token used to access Slack API, must have channels:history scope for the"
+        " bot user."
+    )
+
+
+class SlackConnectionConfig(ConnectionConfig):
+    access_config: Secret[SlackAccessConfig]
+
+    @requires_dependencies(["slack_sdk"], extras="slack")
+    @SourceConnectionError.wrap
+    def get_client(self) -> "WebClient":
+        from slack_sdk import WebClient
+
+        return WebClient(token=self.access_config.get_secret_value().token)
+
+    @requires_dependencies(["slack_sdk"], extras="slack")
+    @SourceConnectionError.wrap
+    def get_async_client(self) -> "AsyncWebClient":
+        from slack_sdk.web.async_client import AsyncWebClient
+
+        return AsyncWebClient(token=self.access_config.get_secret_value().token)
+
+
+class SlackIndexerConfig(IndexerConfig):
+    channels: list[str] = Field(
+        description="Comma-delimited list of Slack channel IDs to pull messages from, can be"
+        " both public or private channels."
+    )
+    start_date: Optional[datetime] = Field(
+        default=None,
+        description="Start date/time in formats YYYY-MM-DD[T]HH:MM[:SS[.ffffff]][Z or [±]HH[:]MM]"
+        " or YYYY-MM-DD",
+    )
+    end_date: Optional[datetime] = Field(
+        default=None,
+        description="End date/time in formats YYYY-MM-DD[T]HH:MM[:SS[.ffffff]][Z or [±]HH[:]MM]"
+        " or YYYY-MM-DD",
+    )
+
+
+@dataclass
+class SlackIndexer(Indexer):
+    index_config: SlackIndexerConfig
+    connection_config: SlackConnectionConfig
+    connector_type: str = CONNECTOR_TYPE
+
+    def run(self, **kwargs: Any) -> Generator[FileData, None, None]:
+        client = self.connection_config.get_client()
+        for channel in self.index_config.channels:
+            messages = []
+            oldest = (
+                str(self.index_config.start_date.timestamp())
+                if self.index_config.start_date is not None
+                else None
+            )
+            latest = (
+                str(self.index_config.end_date.timestamp())
+                if self.index_config.end_date is not None
+                else None
+            )
+            for conversation_history in client.conversations_history(
+                channel=channel,
+                oldest=oldest,
+                latest=latest,
+                limit=PAGINATION_LIMIT,
+            ):
+                messages = conversation_history.get("messages", [])
+                if messages:
+                    yield self._messages_to_file_data(messages, channel)
+
+    def _messages_to_file_data(
+        self,
+        messages: list[dict],
+        channel: str,
+    ) -> FileData:
+        ts_oldest = min((message["ts"] for message in messages), key=lambda m: float(m))
+        ts_newest = max((message["ts"] for message in messages), key=lambda m: float(m))
+
+        identifier_base = f"{channel}-{ts_oldest}-{ts_newest}"
+        identifier = hashlib.sha256(identifier_base.encode("utf-8")).hexdigest()
+        filename = identifier[:16]
+
+        return FileData(
+            identifier=identifier,
+            connector_type=CONNECTOR_TYPE,
+            source_identifiers=SourceIdentifiers(
+                filename=f"{filename}.xml", fullpath=f"{filename}.xml"
+            ),
+            metadata=FileDataSourceMetadata(
+                date_created=ts_oldest,
+                date_modified=ts_newest,
+                date_processed=str(time.time()),
+                record_locator={
+                    "channel": channel,
+                    "oldest": ts_oldest,
+                    "latest": ts_newest,
+                },
+            ),
+        )
+
+    @SourceConnectionError.wrap
+    def precheck(self) -> None:
+        client = self.connection_config.get_client()
+        for channel in self.index_config.channels:
+            # NOTE: Querying conversations history guarantees that the bot is in the channel
+            client.conversations_history(channel=channel, limit=1)
+
+
+class SlackDownloaderConfig(DownloaderConfig):
+    pass
+
+
+@dataclass
+class SlackDownloader(Downloader):
+    connector_type: str = CONNECTOR_TYPE
+    connection_config: SlackConnectionConfig
+    download_config: SlackDownloaderConfig = field(default_factory=SlackDownloaderConfig)
+
+    def run(self, file_data, **kwargs):
+        raise NotImplementedError
+
+    async def run_async(self, file_data: FileData, **kwargs) -> download_responses:
+        # 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:
+            logger.error(
+                "Generated download path is None, source_identifiers might be missing"
+                "from FileData."
+            )
+            raise ValueError("Generated invalid download path.")
+
+        await self._download_conversation(file_data, download_path)
+        return self.generate_download_response(file_data, download_path)
+
+    def is_async(self):
+        return True
+
+    async def _download_conversation(self, file_data: FileData, download_path: Path) -> None:
+        # NOTE: Indexer should supply the record locator in metadata
+        if (
+            file_data.metadata.record_locator is None
+            or "channel" not in file_data.metadata.record_locator
+            or "oldest" not in file_data.metadata.record_locator
+            or "latest" not in file_data.metadata.record_locator
+        ):
+            logger.error(
+                f"Invalid record locator in metadata: {file_data.metadata.record_locator}."
+                "Keys 'channel', 'oldest' and 'latest' must be present."
+            )
+            raise ValueError("Invalid record locator.")
+
+        client = self.connection_config.get_async_client()
+        messages = []
+        async for conversation_history in await client.conversations_history(
+            channel=file_data.metadata.record_locator["channel"],
+            oldest=file_data.metadata.record_locator["oldest"],
+            latest=file_data.metadata.record_locator["latest"],
+            limit=PAGINATION_LIMIT,
+            # NOTE: In order to get the exact same range of messages as indexer, it provides
+            # timestamps of oldest and newest messages, inclusive=True is necessary to include them
+            inclusive=True,
+        ):
+            messages += conversation_history.get("messages", [])
+
+        conversation = []
+        for message in messages:
+            thread_messages = []
+            async for conversations_replies in await client.conversations_replies(
+                channel=file_data.metadata.record_locator["channel"],
+                ts=message["ts"],
+                limit=PAGINATION_LIMIT,
+            ):
+                thread_messages += conversations_replies.get("messages", [])
+
+            # NOTE: Replies contains the whole thread, including the message references by the `ts`
+            # parameter even if it's the only message (there were no replies).
+            # Reference: https://api.slack.com/methods/conversations.replies#markdown
+            conversation.append(thread_messages)
+
+        conversation_xml = self._conversation_to_xml(conversation)
+        download_path.parent.mkdir(exist_ok=True, parents=True)
+        conversation_xml.write(download_path, encoding="utf-8", xml_declaration=True)
+
+    def _conversation_to_xml(self, conversation: list[list[dict]]) -> ET.ElementTree:
+        root = ET.Element("messages")
+
+        for thread in conversation:
+            message, *replies = thread
+            message_elem = ET.SubElement(root, "message")
+            text_elem = ET.SubElement(message_elem, "text")
+            text_elem.text = message.get("text")
+
+            for reply in replies:
+                reply_msg = reply.get("text", "")
+                text_elem.text = "".join([str(text_elem.text), " <reply> ", reply_msg])
+
+        return ET.ElementTree(root)
+
+
+slack_source_entry = SourceRegistryEntry(
+    indexer=SlackIndexer,
+    indexer_config=SlackIndexerConfig,
+    downloader=SlackDownloader,
+    downloader_config=DownloaderConfig,
+    connection_config=SlackConnectionConfig,
+)

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

@@ -2,12 +2,25 @@ from __future__ import annotations
 
 from unstructured_ingest.v2.processes.connector_registry import (
     add_destination_entry,
+    add_source_entry,
 )
 
 from .postgres import CONNECTOR_TYPE as POSTGRES_CONNECTOR_TYPE
-from .postgres import postgres_destination_entry
+from .postgres import postgres_destination_entry, postgres_source_entry
+from .singlestore import CONNECTOR_TYPE as SINGLESTORE_CONNECTOR_TYPE
+from .singlestore import singlestore_destination_entry
+from .snowflake import CONNECTOR_TYPE as SNOWFLAKE_CONNECTOR_TYPE
+from .snowflake import snowflake_destination_entry, snowflake_source_entry
 from .sqlite import CONNECTOR_TYPE as SQLITE_CONNECTOR_TYPE
-from .sqlite import sqlite_destination_entry
+from .sqlite import sqlite_destination_entry, sqlite_source_entry
+
+add_source_entry(source_type=SQLITE_CONNECTOR_TYPE, entry=sqlite_source_entry)
+add_source_entry(source_type=POSTGRES_CONNECTOR_TYPE, entry=postgres_source_entry)
+add_source_entry(source_type=SNOWFLAKE_CONNECTOR_TYPE, entry=snowflake_source_entry)
 
 add_destination_entry(destination_type=SQLITE_CONNECTOR_TYPE, entry=sqlite_destination_entry)
 add_destination_entry(destination_type=POSTGRES_CONNECTOR_TYPE, entry=postgres_destination_entry)
+add_destination_entry(destination_type=SNOWFLAKE_CONNECTOR_TYPE, entry=snowflake_destination_entry)
+add_destination_entry(
+    destination_type=SINGLESTORE_CONNECTOR_TYPE, entry=singlestore_destination_entry
+)

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

@@ -1,17 +1,17 @@
+from contextlib import contextmanager
 from dataclasses import dataclass, field
-from pathlib import Path
-from typing import TYPE_CHECKING, Any, Optional
+from typing import TYPE_CHECKING, Generator, Optional
 
-import numpy as np
-import pandas as pd
 from pydantic import Field, Secret
 
 from unstructured_ingest.utils.dep_check import requires_dependencies
 from unstructured_ingest.v2.interfaces import FileData
 from unstructured_ingest.v2.logger import logger
-from unstructured_ingest.v2.processes.connector_registry import DestinationRegistryEntry
+from unstructured_ingest.v2.processes.connector_registry import (
+    DestinationRegistryEntry,
+    SourceRegistryEntry,
+)
 from unstructured_ingest.v2.processes.connectors.sql.sql import (
-    _DATE_COLUMNS,
     SQLAccessConfig,
     SQLConnectionConfig,
     SQLDownloader,
@@ -22,11 +22,11 @@ from unstructured_ingest.v2.processes.connectors.sql.sql import (
     SQLUploaderConfig,
     SQLUploadStager,
     SQLUploadStagerConfig,
-    parse_date_string,
 )
 
 if TYPE_CHECKING:
     from psycopg2.extensions import connection as PostgresConnection
+    from psycopg2.extensions import cursor as PostgresCursor
 
 CONNECTOR_TYPE = "postgres"
 
@@ -48,18 +48,33 @@ class PostgresConnectionConfig(SQLConnectionConfig):
     port: Optional[int] = Field(default=5432, description="DB host connection port")
     connector_type: str = Field(default=CONNECTOR_TYPE, init=False)
 
+    @contextmanager
     @requires_dependencies(["psycopg2"], extras="postgres")
-    def get_connection(self) -> "PostgresConnection":
+    def get_connection(self) -> Generator["PostgresConnection", None, None]:
         from psycopg2 import connect
 
         access_config = self.access_config.get_secret_value()
-        return connect(
+        connection = connect(
             user=self.username,
             password=access_config.password,
             dbname=self.database,
             host=self.host,
             port=self.port,
         )
+        try:
+            yield connection
+        finally:
+            connection.commit()
+            connection.close()
+
+    @contextmanager
+    def get_cursor(self) -> Generator["PostgresCursor", None, None]:
+        with self.get_connection() as connection:
+            cursor = connection.cursor()
+            try:
+                yield cursor
+            finally:
+                cursor.close()
 
 
 class PostgresIndexerConfig(SQLIndexerConfig):
@@ -72,16 +87,6 @@ class PostgresIndexer(SQLIndexer):
     index_config: PostgresIndexerConfig
     connector_type: str = CONNECTOR_TYPE
 
-    def _get_doc_ids(self) -> list[str]:
-        connection = self.connection_config.get_connection()
-        with connection.cursor() as cursor:
-            cursor.execute(
-                f"SELECT {self.index_config.id_column} FROM {self.index_config.table_name}"
-            )
-            results = cursor.fetchall()
-            ids = [result[0] for result in results]
-            return ids
-
 
 class PostgresDownloaderConfig(SQLDownloaderConfig):
     pass
@@ -97,8 +102,7 @@ class PostgresDownloader(SQLDownloader):
         table_name = file_data.additional_metadata["table_name"]
         id_column = file_data.additional_metadata["id_column"]
         ids = file_data.additional_metadata["ids"]
-        connection = self.connection_config.get_connection()
-        with connection.cursor() as cursor:
+        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=fields,
@@ -130,43 +134,16 @@ class PostgresUploader(SQLUploader):
     upload_config: PostgresUploaderConfig = field(default_factory=PostgresUploaderConfig)
     connection_config: PostgresConnectionConfig
     connector_type: str = CONNECTOR_TYPE
+    values_delimiter: str = "%s"
 
-    def prepare_data(
-        self, columns: list[str], data: tuple[tuple[Any, ...], ...]
-    ) -> list[tuple[Any, ...]]:
-        output = []
-        for row in data:
-            parsed = []
-            for column_name, value in zip(columns, row):
-                if column_name in _DATE_COLUMNS:
-                    if value is None:
-                        parsed.append(None)
-                    else:
-                        parsed.append(parse_date_string(value))
-                else:
-                    parsed.append(value)
-            output.append(tuple(parsed))
-        return output
-
-    def upload_contents(self, path: Path) -> None:
-        df = pd.read_json(path, orient="records", lines=True)
-        logger.debug(f"uploading {len(df)} entries to {self.connection_config.database} ")
-        df.replace({np.nan: None}, inplace=True)
-
-        columns = tuple(df.columns)
-        stmt = f"INSERT INTO {self.upload_config.table_name} ({','.join(columns)}) \
-                VALUES({','.join(['%s' for x in columns])})"  # noqa E501
-
-        for rows in pd.read_json(
-            path, orient="records", lines=True, chunksize=self.upload_config.batch_size
-        ):
-            with self.connection_config.get_connection() as conn:
-                values = self.prepare_data(columns, tuple(rows.itertuples(index=False, name=None)))
-                with conn.cursor() as cur:
-                    cur.executemany(stmt, values)
-
-                conn.commit()
 
+postgres_source_entry = SourceRegistryEntry(
+    connection_config=PostgresConnectionConfig,
+    indexer_config=PostgresIndexerConfig,
+    indexer=PostgresIndexer,
+    downloader_config=PostgresDownloaderConfig,
+    downloader=PostgresDownloader,
+)
 
 postgres_destination_entry = DestinationRegistryEntry(
     connection_config=PostgresConnectionConfig,

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

@@ -0,0 +1,168 @@
+import json
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, Generator, Optional
+
+from pydantic import Field, Secret
+
+from unstructured_ingest.v2.interfaces import FileData
+from unstructured_ingest.v2.logger import logger
+from unstructured_ingest.v2.processes.connector_registry import (
+    DestinationRegistryEntry,
+    SourceRegistryEntry,
+)
+from unstructured_ingest.v2.processes.connectors.sql.sql import (
+    _DATE_COLUMNS,
+    SQLAccessConfig,
+    SQLConnectionConfig,
+    SQLDownloader,
+    SQLDownloaderConfig,
+    SQLIndexer,
+    SQLIndexerConfig,
+    SQLUploader,
+    SQLUploaderConfig,
+    SQLUploadStager,
+    SQLUploadStagerConfig,
+    parse_date_string,
+)
+
+if TYPE_CHECKING:
+    from singlestoredb.connection import Connection as SingleStoreConnection
+    from singlestoredb.connection import Cursor as SingleStoreCursor
+
+CONNECTOR_TYPE = "singlestore"
+
+
+class SingleStoreAccessConfig(SQLAccessConfig):
+    password: Optional[str] = Field(default=None, description="SingleStore password")
+
+
+class SingleStoreConnectionConfig(SQLConnectionConfig):
+    access_config: Secret[SingleStoreAccessConfig]
+    host: Optional[str] = Field(default=None, description="SingleStore host")
+    port: Optional[int] = Field(default=None, description="SingleStore port")
+    user: Optional[str] = Field(default=None, description="SingleStore user")
+    database: Optional[str] = Field(default=None, description="SingleStore database")
+
+    @contextmanager
+    def get_connection(self) -> Generator["SingleStoreConnection", None, None]:
+        import singlestoredb as s2
+
+        connection = s2.connect(
+            host=self.host,
+            port=self.port,
+            database=self.database,
+            user=self.user,
+            password=self.access_config.get_secret_value().password,
+        )
+        try:
+            yield connection
+        finally:
+            connection.commit()
+            connection.close()
+
+    @contextmanager
+    def get_cursor(self) -> Generator["SingleStoreCursor", None, None]:
+        with self.get_connection() as connection:
+            with connection.cursor() as cursor:
+                try:
+                    yield cursor
+                finally:
+                    cursor.close()
+
+
+class SingleStoreIndexerConfig(SQLIndexerConfig):
+    pass
+
+
+@dataclass
+class SingleStoreIndexer(SQLIndexer):
+    connection_config: SingleStoreConnectionConfig
+    index_config: SingleStoreIndexerConfig
+    connector_type: str = CONNECTOR_TYPE
+
+
+class SingleStoreDownloaderConfig(SQLDownloaderConfig):
+    pass
+
+
+@dataclass
+class SingleStoreDownloader(SQLDownloader):
+    connection_config: SingleStoreConnectionConfig
+    download_config: SingleStoreDownloaderConfig
+    connector_type: str = CONNECTOR_TYPE
+
+    def query_db(self, file_data: FileData) -> tuple[list[tuple], list[str]]:
+        table_name = file_data.additional_metadata["table_name"]
+        id_column = file_data.additional_metadata["id_column"]
+        ids = file_data.additional_metadata["ids"]
+        with self.connection_config.get_connection() as sqlite_connection:
+            cursor = sqlite_connection.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=fields,
+                table_name=table_name,
+                id_column=id_column,
+                ids=",".join([str(i) for i in ids]),
+            )
+            logger.debug(f"running query: {query}")
+            cursor.execute(query)
+            rows = cursor.fetchall()
+            columns = [col[0] for col in cursor.description]
+            return rows, columns
+
+
+class SingleStoreUploadStagerConfig(SQLUploadStagerConfig):
+    pass
+
+
+class SingleStoreUploadStager(SQLUploadStager):
+    upload_stager_config: SingleStoreUploadStagerConfig
+
+
+class SingleStoreUploaderConfig(SQLUploaderConfig):
+    pass
+
+
+@dataclass
+class SingleStoreUploader(SQLUploader):
+    upload_config: SingleStoreUploaderConfig = field(default_factory=SingleStoreUploaderConfig)
+    connection_config: SingleStoreConnectionConfig
+    values_delimiter: str = "%s"
+    connector_type: str = CONNECTOR_TYPE
+
+    def prepare_data(
+        self, columns: list[str], data: tuple[tuple[Any, ...], ...]
+    ) -> list[tuple[Any, ...]]:
+        output = []
+        for row in data:
+            parsed = []
+            for column_name, value in zip(columns, row):
+                if isinstance(value, (list, dict)):
+                    value = json.dumps(value)
+                if column_name in _DATE_COLUMNS:
+                    if value is None:
+                        parsed.append(None)
+                    else:
+                        parsed.append(parse_date_string(value))
+                else:
+                    parsed.append(value)
+            output.append(tuple(parsed))
+        return output
+
+
+singlestore_source_entry = SourceRegistryEntry(
+    connection_config=SingleStoreConnectionConfig,
+    indexer_config=SingleStoreIndexerConfig,
+    indexer=SQLIndexer,
+    downloader_config=SingleStoreDownloaderConfig,
+    downloader=SingleStoreDownloader,
+)
+
+singlestore_destination_entry = DestinationRegistryEntry(
+    connection_config=SingleStoreConnectionConfig,
+    uploader=SingleStoreUploader,
+    uploader_config=SingleStoreUploaderConfig,
+    upload_stager=SingleStoreUploadStager,
+    upload_stager_config=SingleStoreUploadStagerConfig,
+)