unstructured-ingest 0.3.13__py3-none-any.whl

unstructured_ingest/connector/confluence.py

@@ -0,0 +1,285 @@
+import math
+import typing as t
+from dataclasses import dataclass, field
+from datetime import datetime
+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
+
+if t.TYPE_CHECKING:
+    from atlassian import Confluence
+
+
+@dataclass
+class ConfluenceAccessConfig(AccessConfig):
+    api_token: str = enhanced_field(sensitive=True)
+
+
+@dataclass
+class SimpleConfluenceConfig(BaseConnectorConfig):
+    """Connector config where:
+    user_email is the email to authenticate into Confluence Cloud,
+    api_token is the api token to authenticate into Confluence Cloud,
+    and url is the URL pointing to the Confluence Cloud instance.
+
+    Check https://developer.atlassian.com/cloud/confluence/basic-auth-for-rest-apis/
+    for more info on the api_token.
+    """
+
+    user_email: str
+    access_config: ConfluenceAccessConfig
+    url: str
+    max_num_of_spaces: int = 500
+    max_num_of_docs_from_each_space: int = 100
+    spaces: t.List[str] = field(default_factory=list)
+
+
+@dataclass
+class ConfluenceDocumentMeta:
+    """Metadata specifying:
+    id for the confluence space that the document locates in,
+    and the id of document that is being reached to.
+    """
+
+    space_id: str
+    document_id: str
+
+
+def scroll_wrapper(func):
+    def wrapper(*args, **kwargs):
+        """Wraps a function to obtain scroll functionality."""
+        number_of_items_to_fetch = kwargs["number_of_items_to_fetch"]
+        del kwargs["number_of_items_to_fetch"]
+
+        kwargs["limit"] = min(100, number_of_items_to_fetch)
+        kwargs["start"] = kwargs.get("start", 0)
+
+        all_results = []
+        num_iterations = math.ceil(number_of_items_to_fetch / kwargs["limit"])
+
+        for _ in range(num_iterations):
+            response = func(*args, **kwargs)
+            if isinstance(response, list):
+                all_results += func(*args, **kwargs)
+            elif isinstance(response, dict):
+                all_results += func(*args, **kwargs)["results"]
+
+            kwargs["start"] += kwargs["limit"]
+
+        return all_results[:number_of_items_to_fetch]
+
+    return wrapper
+
+
+@dataclass
+class ConfluenceIngestDoc(IngestDocCleanupMixin, BaseSingleIngestDoc):
+    """Class encapsulating fetching a doc and writing processed results (but not
+    doing the processing).
+
+    Current implementation creates a Confluence connection object
+    to fetch each doc, rather than creating a it for each thread.
+    """
+
+    connector_config: SimpleConfluenceConfig
+    document_meta: ConfluenceDocumentMeta
+    registry_name: str = "confluence"
+
+    # TODO: remove one of filename or _tmp_download_file, using a wrapper
+    @property
+    def filename(self):
+        if not self.read_config.download_dir:
+            return None
+        return (
+            Path(self.read_config.download_dir)
+            / self.document_meta.space_id
+            / f"{self.document_meta.document_id}.html"
+        ).resolve()
+
+    @property
+    def _output_filename(self):
+        """Create output file path based on output directory, space id and document id."""
+        output_file = f"{self.document_meta.document_id}.json"
+        return Path(self.processor_config.output_dir) / self.document_meta.space_id / output_file
+
+    @property
+    def record_locator(self) -> t.Optional[t.Dict[str, t.Any]]:
+        return {
+            "url": self.connector_config.url,
+            "page_id": self.document_meta.document_id,
+        }
+
+    @SourceConnectionNetworkError.wrap
+    @requires_dependencies(["atlassian"], extras="Confluence")
+    def _get_page(self):
+        from atlassian import Confluence
+        from atlassian.errors import ApiError
+
+        try:
+            confluence = Confluence(
+                self.connector_config.url,
+                username=self.connector_config.user_email,
+                password=self.connector_config.access_config.api_token,
+            )
+            result = confluence.get_page_by_id(
+                page_id=self.document_meta.document_id,
+                expand="history.lastUpdated,version,body.view",
+            )
+        except ApiError as e:
+            logger.error(e)
+            return None
+        return result
+
+    def update_source_metadata(self, **kwargs):
+        """Fetches file metadata from the current page."""
+        page = kwargs.get("page", self._get_page())
+        if page is None:
+            self.source_metadata = SourceMetadata(
+                exists=False,
+            )
+            return
+        document_history = page["history"]
+        date_created = datetime.strptime(
+            document_history["createdDate"],
+            "%Y-%m-%dT%H:%M:%S.%fZ",
+        ).isoformat()
+        if last_updated := document_history.get("lastUpdated", {}).get("when", ""):
+            date_modified = datetime.strptime(
+                last_updated,
+                "%Y-%m-%dT%H:%M:%S.%fZ",
+            ).isoformat()
+        else:
+            date_modified = date_created
+        version = page["version"]["number"]
+        self.source_metadata = SourceMetadata(
+            date_created=date_created,
+            date_modified=date_modified,
+            version=version,
+            source_url=page["_links"].get("self", None),
+            exists=True,
+        )
+
+    @SourceConnectionError.wrap
+    @requires_dependencies(["atlassian"], extras="confluence")
+    @BaseSingleIngestDoc.skip_if_file_exists
+    def get_file(self):
+        # TODO: instead of having a separate connection object for each doc,
+        # have a separate connection object for each process
+
+        result = self._get_page()
+        self.update_source_metadata(page=result)
+        if result is None:
+            raise ValueError(f"Failed to retrieve page with ID {self.document_meta.document_id}")
+        self.document = result["body"]["view"]["value"]
+        self.filename.parent.mkdir(parents=True, exist_ok=True)
+        with open(self.filename, "w", encoding="utf8") as f:
+            f.write(self.document)
+
+
+@dataclass
+class ConfluenceSourceConnector(SourceConnectorCleanupMixin, BaseSourceConnector):
+    """Fetches body fields from all documents within all spaces in a Confluence Cloud instance."""
+
+    connector_config: SimpleConfluenceConfig
+    _confluence: t.Optional["Confluence"] = field(init=False, default=None)
+
+    @property
+    def confluence(self) -> "Confluence":
+        from atlassian import Confluence
+
+        if self._confluence is None:
+            self._confluence = Confluence(
+                url=self.connector_config.url,
+                username=self.connector_config.user_email,
+                password=self.connector_config.access_config.api_token,
+            )
+        return self._confluence
+
+    @requires_dependencies(["atlassian", "requests"], extras="Confluence")
+    def check_connection(self):
+        import requests
+
+        url = "rest/api/space"
+        try:
+            self.confluence.request(method="HEAD", path=url)
+        except requests.HTTPError as http_error:
+            logger.error(f"failed to validate connection: {http_error}", exc_info=True)
+            raise SourceConnectionError(f"failed to validate connection: {http_error}")
+
+    @requires_dependencies(["atlassian"], extras="Confluence")
+    def initialize(self):
+        self.list_of_spaces = None
+        if self.connector_config.spaces:
+            self.list_of_spaces = self.connector_config.spaces
+            if self.connector_config.max_num_of_spaces:
+                logger.warning(
+                    """--confluence-list-of-spaces and --confluence-num-of-spaces cannot
+                    be used at the same time. Connector will only fetch the
+                    --confluence-list-of-spaces that you've provided.""",
+                )
+
+    @requires_dependencies(["atlassian"], extras="Confluence")
+    def _get_space_ids(self):
+        """Fetches spaces in a confluence domain."""
+
+        get_spaces_with_scroll = scroll_wrapper(self.confluence.get_all_spaces)
+
+        all_results = get_spaces_with_scroll(
+            number_of_items_to_fetch=self.connector_config.max_num_of_spaces,
+        )
+
+        space_ids = [space["key"] for space in all_results]
+        return space_ids
+
+    @requires_dependencies(["atlassian"], extras="Confluence")
+    def _get_docs_ids_within_one_space(
+        self,
+        space_id: str,
+        content_type: str = "page",
+    ):
+        get_pages_with_scroll = scroll_wrapper(self.confluence.get_all_pages_from_space)
+        results = get_pages_with_scroll(
+            space=space_id,
+            number_of_items_to_fetch=self.connector_config.max_num_of_docs_from_each_space,
+            content_type=content_type,
+        )
+
+        doc_ids = [(space_id, doc["id"]) for doc in results]
+        return doc_ids
+
+    @requires_dependencies(["atlassian"], extras="Confluence")
+    def _get_doc_ids_within_spaces(self):
+        space_ids = self._get_space_ids() if not self.list_of_spaces else self.list_of_spaces
+
+        doc_ids_all = [self._get_docs_ids_within_one_space(space_id=id) for id in space_ids]
+
+        doc_ids_flattened = [
+            (space_id, doc_id)
+            for doc_ids_space in doc_ids_all
+            for space_id, doc_id in doc_ids_space
+        ]
+        return doc_ids_flattened
+
+    def get_ingest_docs(self):
+        """Fetches all documents in a confluence space."""
+        doc_ids = self._get_doc_ids_within_spaces()
+        return [
+            ConfluenceIngestDoc(
+                connector_config=self.connector_config,
+                processor_config=self.processor_config,
+                read_config=self.read_config,
+                document_meta=ConfluenceDocumentMeta(space_id, doc_id),
+            )
+            for space_id, doc_id in doc_ids
+        ]

unstructured_ingest/connector/databricks_volumes.py

@@ -0,0 +1,137 @@
+import copy
+import json
+import os
+import typing as t
+from dataclasses import dataclass, field
+from io import BytesIO
+from pathlib import PurePath
+
+from unstructured_ingest.enhanced_dataclass import enhanced_field
+from unstructured_ingest.enhanced_dataclass.core import _asdict
+from unstructured_ingest.error import DestinationConnectionError
+from unstructured_ingest.interfaces import (
+    AccessConfig,
+    BaseConnectorConfig,
+    BaseDestinationConnector,
+    BaseSingleIngestDoc,
+    WriteConfig,
+)
+from unstructured_ingest.logger import logger
+from unstructured_ingest.utils.dep_check import requires_dependencies
+
+if t.TYPE_CHECKING:
+    from databricks.sdk import WorkspaceClient
+
+
+@dataclass
+class DatabricksVolumesAccessConfig(AccessConfig):
+    account_id: t.Optional[str] = None
+    username: t.Optional[str] = None
+    password: t.Optional[str] = enhanced_field(default=None, sensitive=True)
+    client_id: t.Optional[str] = None
+    client_secret: t.Optional[str] = enhanced_field(default=None, sensitive=True)
+    token: t.Optional[str] = enhanced_field(default=None, sensitive=True)
+    profile: t.Optional[str] = None
+    azure_workspace_resource_id: t.Optional[str] = None
+    azure_client_secret: t.Optional[str] = enhanced_field(default=None, sensitive=True)
+    azure_client_id: t.Optional[str] = None
+    azure_tenant_id: t.Optional[str] = None
+    azure_environment: t.Optional[str] = None
+    auth_type: t.Optional[str] = None
+    cluster_id: t.Optional[str] = None
+    google_credentials: t.Optional[str] = None
+    google_service_account: t.Optional[str] = None
+
+
+@dataclass
+class SimpleDatabricksVolumesConfig(BaseConnectorConfig):
+    access_config: DatabricksVolumesAccessConfig
+    host: t.Optional[str] = None
+
+
+@dataclass
+class DatabricksVolumesWriteConfig(WriteConfig):
+    volume: str
+    catalog: str
+    volume_path: t.Optional[str] = None
+    overwrite: bool = False
+    encoding: str = "utf-8"
+    schema: str = "default"
+
+    @property
+    def path(self) -> str:
+        path = f"/Volumes/{self.catalog}/{self.schema}/{self.volume}"
+        if self.volume_path:
+            path = f"{path}/{self.volume_path}"
+        return path
+
+
+@dataclass
+class DatabricksVolumesDestinationConnector(BaseDestinationConnector):
+    write_config: DatabricksVolumesWriteConfig
+    connector_config: SimpleDatabricksVolumesConfig
+    _client: t.Optional["WorkspaceClient"] = field(init=False, default=None)
+
+    def to_dict(self, **kwargs):
+        self_cp = copy.copy(self)
+        if hasattr(self_cp, "_client"):
+            setattr(self_cp, "_client", None)
+        return _asdict(self_cp, **kwargs)
+
+    @requires_dependencies(dependencies=["databricks.sdk"], extras="databricks-volumes")
+    def generate_client(self) -> "WorkspaceClient":
+        from databricks.sdk import WorkspaceClient
+
+        return WorkspaceClient(
+            host=self.connector_config.host, **self.connector_config.access_config.to_dict()
+        )
+
+    @property
+    def client(self) -> "WorkspaceClient":
+        if self._client is None:
+            self._client = self.generate_client()
+        return self._client
+
+    def check_connection(self):
+        try:
+            assert self.client.current_user.me().active
+        except Exception as e:
+            logger.error(f"failed to validate connection: {e}", exc_info=True)
+            raise DestinationConnectionError(f"failed to validate connection: {e}")
+
+    def initialize(self):
+        _ = self.client
+
+    def write_dict(
+        self,
+        *args,
+        elements_dict: t.List[t.Dict[str, t.Any]],
+        filename: t.Optional[str] = None,
+        indent: int = 4,
+        encoding: str = "utf-8",
+        **kwargs,
+    ) -> None:
+        output_folder = self.write_config.path
+        output_folder = os.path.join(output_folder)  # Make sure folder ends with file separator
+        filename = (
+            filename.strip(os.sep) if filename else filename
+        )  # Make sure filename doesn't begin with file separator
+        output_path = str(PurePath(output_folder, filename)) if filename else output_folder
+        logger.debug(f"uploading content to {output_path}")
+        self.client.files.upload(
+            file_path=output_path,
+            contents=BytesIO(json.dumps(elements_dict).encode(encoding=self.write_config.encoding)),
+            overwrite=self.write_config.overwrite,
+        )
+
+    def get_elements_dict(self, docs: t.List[BaseSingleIngestDoc]) -> t.List[t.Dict[str, t.Any]]:
+        pass
+
+    def write(self, docs: t.List[BaseSingleIngestDoc]) -> None:
+        for doc in docs:
+            file_path = doc.base_output_filename
+            filename = file_path if file_path else None
+            with open(doc._output_filename) as json_file:
+                logger.debug(f"uploading content from {doc._output_filename}")
+                json_list = json.load(json_file)
+                self.write_dict(elements_dict=json_list, filename=filename)

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()