unstructured-ingest 0.0.0__py3-none-any.whl

unstructured_ingest/connector/delta_table.py

@@ -0,0 +1,203 @@
+import os
+import typing as t
+from dataclasses import dataclass
+from datetime import datetime as dt
+from multiprocessing import Process
+from pathlib import Path
+
+from unstructured_ingest.error import SourceConnectionError, SourceConnectionNetworkError
+from unstructured_ingest.interfaces import (
+    BaseConnectorConfig,
+    BaseDestinationConnector,
+    BaseSingleIngestDoc,
+    BaseSourceConnector,
+    IngestDocCleanupMixin,
+    SourceConnectorCleanupMixin,
+    SourceMetadata,
+    WriteConfig,
+)
+from unstructured_ingest.logger import logger
+from unstructured_ingest.utils.dep_check import requires_dependencies
+
+if t.TYPE_CHECKING:
+    from deltalake import DeltaTable
+
+
+@dataclass
+class SimpleDeltaTableConfig(BaseConnectorConfig):
+    table_uri: t.Union[str, Path]
+    version: t.Optional[int] = None
+    storage_options: t.Optional[t.Dict[str, str]] = None
+    without_files: bool = False
+
+
+@dataclass
+class DeltaTableIngestDoc(IngestDocCleanupMixin, BaseSingleIngestDoc):
+    connector_config: SimpleDeltaTableConfig
+    uri: str
+    modified_date: str
+    created_at: str
+    registry_name: str = "delta-table"
+
+    def uri_filename(self) -> str:
+        basename = os.path.basename(self.uri)
+        return os.path.splitext(basename)[0]
+
+    @property
+    def filename(self):
+        return (Path(self.read_config.download_dir) / f"{self.uri_filename()}.csv").resolve()
+
+    @property
+    def _output_filename(self):
+        """Create filename document id combined with a hash of the query to uniquely identify
+        the output file."""
+        return Path(self.processor_config.output_dir) / f"{self.uri_filename()}.json"
+
+    def _create_full_tmp_dir_path(self):
+        self.filename.parent.mkdir(parents=True, exist_ok=True)
+        self._output_filename.parent.mkdir(parents=True, exist_ok=True)
+
+    @requires_dependencies(["fsspec"], extras="delta-table")
+    def _get_fs_from_uri(self):
+        from fsspec.core import url_to_fs
+
+        try:
+            fs, _ = url_to_fs(self.uri)
+        except ImportError as error:
+            raise ImportError(
+                f"uri {self.uri} may be associated with a filesystem that "
+                f"requires additional dependencies: {error}",
+            )
+        return fs
+
+    def update_source_metadata(self, **kwargs):
+        fs = kwargs.get("fs", self._get_fs_from_uri())
+        version = (
+            fs.checksum(self.uri) if fs.protocol != "gs" else fs.info(self.uri).get("etag", "")
+        )
+        file_exists = fs.exists(self.uri)
+        self.source_metadata = SourceMetadata(
+            date_created=self.created_at,
+            date_modified=self.modified_date,
+            version=version,
+            source_url=self.uri,
+            exists=file_exists,
+        )
+
+    @SourceConnectionError.wrap
+    @BaseSingleIngestDoc.skip_if_file_exists
+    def get_file(self):
+        fs = self._get_fs_from_uri()
+        self.update_source_metadata(fs=fs)
+        logger.info(f"using a {fs} filesystem to collect table data")
+        self._create_full_tmp_dir_path()
+
+        df = self._get_df(filesystem=fs)
+
+        logger.info(f"writing {len(df)} rows to {self.filename}")
+        df.to_csv(self.filename)
+
+    @SourceConnectionNetworkError.wrap
+    def _get_df(self, filesystem):
+        import pyarrow.parquet as pq
+
+        return pq.ParquetDataset(self.uri, filesystem=filesystem).read_pandas().to_pandas()
+
+
+@dataclass
+class DeltaTableSourceConnector(SourceConnectorCleanupMixin, BaseSourceConnector):
+    connector_config: SimpleDeltaTableConfig
+    delta_table: t.Optional["DeltaTable"] = None
+
+    def check_connection(self):
+        pass
+
+    @requires_dependencies(["deltalake"], extras="delta-table")
+    def initialize(self):
+        from deltalake import DeltaTable
+
+        self.delta_table = DeltaTable(
+            table_uri=self.connector_config.table_uri,
+            version=self.connector_config.version,
+            storage_options=self.connector_config.storage_options,
+            without_files=self.connector_config.without_files,
+        )
+        rows = self.delta_table.to_pyarrow_dataset().count_rows()
+        if not rows > 0:
+            raise ValueError(f"no data found at {self.connector_config.table_uri}")
+        logger.info(f"processing {rows} rows of data")
+
+    def get_ingest_docs(self):
+        """Batches the results into distinct docs"""
+        if not self.delta_table:
+            raise ValueError("delta table was never initialized")
+        actions = self.delta_table.get_add_actions().to_pandas()
+        mod_date_dict = {
+            row["path"]: str(row["modification_time"]) for _, row in actions.iterrows()
+        }
+        created_at = dt.fromtimestamp(self.delta_table.metadata().created_time / 1000)
+        return [
+            DeltaTableIngestDoc(
+                connector_config=self.connector_config,
+                processor_config=self.processor_config,
+                read_config=self.read_config,
+                uri=uri,
+                modified_date=mod_date_dict[os.path.basename(uri)],
+                created_at=str(created_at),
+            )
+            for uri in self.delta_table.file_uris()
+        ]
+
+
+@dataclass
+class DeltaTableWriteConfig(WriteConfig):
+    drop_empty_cols: bool = False
+    mode: t.Literal["error", "append", "overwrite", "ignore"] = "error"
+    schema_mode: t.Optional[t.Literal["merge", "overwrite"]] = None
+    engine: t.Literal["pyarrow", "rust"] = "pyarrow"
+
+
+@dataclass
+class DeltaTableDestinationConnector(BaseDestinationConnector):
+    write_config: DeltaTableWriteConfig
+    connector_config: SimpleDeltaTableConfig
+
+    @requires_dependencies(["deltalake"], extras="delta-table")
+    def initialize(self):
+        pass
+
+    def check_connection(self):
+        pass
+
+    @requires_dependencies(["deltalake"], extras="delta-table")
+    def write_dict(self, *args, elements_dict: t.List[t.Dict[str, t.Any]], **kwargs) -> None:
+        from deltalake.writer import write_deltalake
+
+        from unstructured_ingest.utils.table import convert_to_pandas_dataframe
+
+        df = convert_to_pandas_dataframe(
+            elements_dict=elements_dict,
+            drop_empty_cols=self.write_config.drop_empty_cols,
+        )
+        logger.info(
+            f"writing {len(df)} rows to destination table "
+            f"at {self.connector_config.table_uri}\ndtypes: {df.dtypes}",
+        )
+        writer_kwargs = {
+            "table_or_uri": self.connector_config.table_uri,
+            "data": df,
+            "mode": self.write_config.mode,
+            "engine": self.write_config.engine,
+        }
+        if self.write_config.schema_mode is not None:
+            writer_kwargs["schema_mode"] = self.write_config.schema_mode
+        # NOTE: deltalake writer on Linux sometimes can finish but still trigger a SIGABRT and cause
+        # ingest to fail, even though all tasks are completed normally. Putting the writer into a
+        # process mitigates this issue by ensuring python interpreter waits properly for deltalake's
+        # rust backend to finish
+        writer = Process(
+            target=write_deltalake,
+            kwargs=writer_kwargs,
+        )
+        writer.start()
+        writer.join()

unstructured_ingest/connector/discord.py

@@ -0,0 +1,180 @@
+import datetime as dt
+import typing as t
+from dataclasses import dataclass
+from pathlib import Path
+
+from unstructured_ingest.enhanced_dataclass import enhanced_field
+from unstructured_ingest.error import SourceConnectionError, SourceConnectionNetworkError
+from unstructured_ingest.interfaces import (
+    AccessConfig,
+    BaseConnectorConfig,
+    BaseSingleIngestDoc,
+    BaseSourceConnector,
+    IngestDocCleanupMixin,
+    SourceConnectorCleanupMixin,
+    SourceMetadata,
+)
+from unstructured_ingest.logger import logger
+from unstructured_ingest.utils.dep_check import (
+    requires_dependencies,
+)
+
+
+@dataclass
+class DiscordAccessConfig(AccessConfig):
+    token: str = enhanced_field(sensitive=True)
+
+
+@dataclass
+class SimpleDiscordConfig(BaseConnectorConfig):
+    """Connector config where channels is a comma separated list of
+    Discord channels to pull messages from.
+    """
+
+    # Discord Specific Options
+    access_config: DiscordAccessConfig
+    channels: t.List[str]
+    period: t.Optional[int] = None
+
+
+@dataclass
+class DiscordIngestDoc(IngestDocCleanupMixin, BaseSingleIngestDoc):
+    """Class encapsulating fetching a doc and writing processed results (but not
+    doing the processing!).
+    Also includes a cleanup method. When things go wrong and the cleanup
+    method is not called, the file is left behind on the filesystem to assist debugging.
+    """
+
+    connector_config: SimpleDiscordConfig
+    channel: str
+    days: t.Optional[int] = None
+    registry_name: str = "discord"
+
+    # NOTE(crag): probably doesn't matter,  but intentionally not defining tmp_download_file
+    # __post_init__ for multiprocessing simplicity (no Path objects in initially
+    # instantiated object)
+    def _tmp_download_file(self):
+        channel_file = self.channel + ".txt"
+        return Path(self.read_config.download_dir) / channel_file
+
+    @property
+    def _output_filename(self):
+        output_file = self.channel + ".json"
+        return Path(self.processor_config.output_dir) / output_file
+
+    def _create_full_tmp_dir_path(self):
+        self._tmp_download_file().parent.mkdir(parents=True, exist_ok=True)
+
+    @SourceConnectionNetworkError.wrap
+    @requires_dependencies(dependencies=["discord"], extras="discord")
+    def _get_messages(self):
+        """Actually fetches the data from discord."""
+        import discord
+        from discord.ext import commands
+
+        messages: t.List[discord.Message] = []
+        jumpurl: t.List[str] = []
+        intents = discord.Intents.default()
+        intents.message_content = True
+        bot = commands.Bot(command_prefix=">", intents=intents)
+
+        @bot.event
+        async def on_ready():
+            try:
+                after_date = None
+                if self.days:
+                    after_date = dt.datetime.utcnow() - dt.timedelta(days=self.days)
+                channel = bot.get_channel(int(self.channel))
+                jumpurl.append(channel.jump_url)  # type: ignore
+                async for msg in channel.history(after=after_date):  # type: ignore
+                    messages.append(msg)
+                await bot.close()
+            except Exception:
+                logger.error("Error fetching messages")
+                await bot.close()
+                raise
+
+        bot.run(self.connector_config.access_config.token)
+        jump_url = None if len(jumpurl) < 1 else jumpurl[0]
+        return messages, jump_url
+
+    def update_source_metadata(self, **kwargs):
+        messages, jump_url = kwargs.get("messages_tuple", self._get_messages())
+        if messages == []:
+            self.source_metadata = SourceMetadata(
+                exists=False,
+            )
+            return
+        dates = [m.created_at for m in messages if m.created_at]
+        dates.sort()
+        self.source_metadata = SourceMetadata(
+            date_created=dates[0].isoformat(),
+            date_modified=dates[-1].isoformat(),
+            source_url=jump_url,
+            exists=True,
+        )
+
+    @SourceConnectionError.wrap
+    @BaseSingleIngestDoc.skip_if_file_exists
+    def get_file(self):
+        self._create_full_tmp_dir_path()
+
+        messages, jump_url = self._get_messages()
+        self.update_source_metadata(messages_tuple=(messages, jump_url))
+        if messages == []:
+            raise ValueError(f"Failed to retrieve messages from Discord channel {self.channel}")
+        self._tmp_download_file().parent.mkdir(parents=True, exist_ok=True)
+        with open(self._tmp_download_file(), "w") as f:
+            for m in messages:
+                f.write(m.content + "\n")
+
+    @property
+    def filename(self):
+        """The filename of the file created from a discord channel"""
+        return self._tmp_download_file()
+
+    @property
+    def version(self) -> t.Optional[str]:
+        return None
+
+    @property
+    def record_locator(self) -> t.Optional[t.Dict[str, t.Any]]:
+        return {
+            "channel": self.channel,
+        }
+
+
+class DiscordSourceConnector(SourceConnectorCleanupMixin, BaseSourceConnector):
+    """Objects of this class support fetching document(s) from"""
+
+    connector_config: SimpleDiscordConfig
+
+    def initialize(self):
+        pass
+
+    @requires_dependencies(dependencies=["discord"], extras="discord")
+    def check_connection(self):
+        import asyncio
+
+        import discord
+        from discord.client import Client
+
+        intents = discord.Intents.default()
+        try:
+            client = Client(intents=intents)
+            asyncio.run(client.start(token=self.connector_config.access_config.token))
+        except Exception as e:
+            logger.error(f"failed to validate connection: {e}", exc_info=True)
+            raise SourceConnectionError(f"failed to validate connection: {e}")
+
+    def get_ingest_docs(self):
+        return [
+            DiscordIngestDoc(
+                connector_config=self.connector_config,
+                processor_config=self.processor_config,
+                read_config=self.read_config,
+                channel=channel,
+                days=self.connector_config.period,
+            )
+            for channel in self.connector_config.channels
+        ]