unstructured-ingest 0.3.14__py3-none-any.whl → 0.4.0__py3-none-any.whl

test/integration/connectors/sql/test_databricks_delta_tables.py

@@ -0,0 +1,142 @@
+import json
+import os
+import time
+from contextlib import contextmanager
+from pathlib import Path
+from uuid import uuid4
+
+import pytest
+from databricks.sql import connect
+from databricks.sql.client import Connection as DeltaTableConnection
+from databricks.sql.client import Cursor as DeltaTableCursor
+from pydantic import BaseModel, SecretStr
+
+from test.integration.connectors.utils.constants import DESTINATION_TAG, SQL_TAG, env_setup_path
+from test.integration.utils import requires_env
+from unstructured_ingest.v2.interfaces import FileData, SourceIdentifiers
+from unstructured_ingest.v2.logger import logger
+from unstructured_ingest.v2.processes.connectors.sql.databricks_delta_tables import (
+    CONNECTOR_TYPE,
+    DatabrickDeltaTablesAccessConfig,
+    DatabrickDeltaTablesConnectionConfig,
+    DatabrickDeltaTablesUploader,
+    DatabrickDeltaTablesUploaderConfig,
+    DatabrickDeltaTablesUploadStager,
+)
+
+CATALOG = "utic-dev-tech-fixtures"
+
+
+class EnvData(BaseModel):
+    server_hostname: str
+    http_path: str
+    access_token: SecretStr
+
+
+def get_env_data() -> EnvData:
+    return EnvData(
+        server_hostname=os.environ["DATABRICKS_SERVER_HOSTNAME"],
+        http_path=os.environ["DATABRICKS_HTTP_PATH"],
+        access_token=os.environ["DATABRICKS_ACCESS_TOKEN"],
+    )
+
+
+def get_destination_schema(new_table_name: str) -> str:
+    p = Path(env_setup_path / "sql" / "databricks_delta_tables" / "destination" / "schema.sql")
+    with p.open() as f:
+        data_lines = f.readlines()
+    data_lines[0] = data_lines[0].replace("elements", new_table_name)
+    data = "".join([line.strip() for line in data_lines])
+    return data
+
+
+@contextmanager
+def get_connection() -> DeltaTableConnection:
+    env_data = get_env_data()
+    with connect(
+        server_hostname=env_data.server_hostname,
+        http_path=env_data.http_path,
+        access_token=env_data.access_token.get_secret_value(),
+    ) as connection:
+        yield connection
+
+
+@contextmanager
+def get_cursor() -> DeltaTableCursor:
+    with get_connection() as connection:
+        with connection.cursor() as cursor:
+            cursor.execute(f"USE CATALOG '{CATALOG}'")
+            yield cursor
+
+
+@pytest.fixture
+def destination_table() -> str:
+    random_id = str(uuid4())[:8]
+    table_name = f"elements_{random_id}"
+    destination_schema = get_destination_schema(new_table_name=table_name)
+    with get_cursor() as cursor:
+        logger.info(f"creating table: {table_name}")
+        cursor.execute(f"DROP TABLE IF EXISTS {table_name}")
+        cursor.execute(destination_schema)
+
+    yield table_name
+    with get_cursor() as cursor:
+        logger.info(f"dropping table: {table_name}")
+        cursor.execute(f"DROP TABLE IF EXISTS {table_name}")
+
+
+def validate_destination(expected_num_elements: int, table_name: str, retries=30, interval=1):
+    with get_cursor() as cursor:
+        for i in range(retries):
+            cursor.execute(f"SELECT COUNT(*) FROM {table_name}")
+            count = cursor.fetchone()[0]
+            if count == expected_num_elements:
+                break
+            logger.info(f"retry attempt {i}: expected {expected_num_elements} != count {count}")
+            time.sleep(interval)
+        assert (
+            count == expected_num_elements
+        ), f"dest check failed: got {count}, expected {expected_num_elements}"
+
+
+@pytest.mark.asyncio
+@pytest.mark.skip("Resources take too long to spin up to run in CI")
+@pytest.mark.tags(CONNECTOR_TYPE, DESTINATION_TAG, SQL_TAG)
+@requires_env("DATABRICKS_SERVER_HOSTNAME", "DATABRICKS_HTTP_PATH", "DATABRICKS_ACCESS_TOKEN")
+async def test_databricks_delta_tables_destination(
+    upload_file: Path, temp_dir: Path, destination_table: str
+):
+    env_data = get_env_data()
+    mock_file_data = FileData(
+        identifier="mock file data",
+        connector_type=CONNECTOR_TYPE,
+        source_identifiers=SourceIdentifiers(filename=upload_file.name, fullpath=upload_file.name),
+    )
+    stager = DatabrickDeltaTablesUploadStager()
+    staged_path = stager.run(
+        elements_filepath=upload_file,
+        file_data=mock_file_data,
+        output_dir=temp_dir,
+        output_filename=upload_file.name,
+    )
+
+    assert staged_path.suffix == upload_file.suffix
+
+    uploader = DatabrickDeltaTablesUploader(
+        connection_config=DatabrickDeltaTablesConnectionConfig(
+            access_config=DatabrickDeltaTablesAccessConfig(
+                token=env_data.access_token.get_secret_value()
+            ),
+            http_path=env_data.http_path,
+            server_hostname=env_data.server_hostname,
+        ),
+        upload_config=DatabrickDeltaTablesUploaderConfig(
+            catalog=CATALOG, database="default", table_name=destination_table
+        ),
+    )
+    with staged_path.open("r") as f:
+        staged_data = json.load(f)
+    expected_num_elements = len(staged_data)
+    uploader.precheck()
+    uploader.run(path=staged_path, file_data=mock_file_data)
+    validate_destination(expected_num_elements=expected_num_elements, table_name=destination_table)

test/integration/connectors/test_confluence.py

@@ -30,10 +30,10 @@ async def test_confluence_source(temp_dir):
     spaces = ["testteamsp", "MFS"]
 
     # Create connection and indexer configurations
-    access_config = ConfluenceAccessConfig(api_token=api_token)
+    access_config = ConfluenceAccessConfig(password=api_token)
     connection_config = ConfluenceConnectionConfig(
         url=confluence_url,
-        user_email=user_email,
+        username=user_email,
         access_config=access_config,
     )
     index_config = ConfluenceIndexerConfig(
@@ -77,10 +77,10 @@ async def test_confluence_source_large(temp_dir):
     spaces = ["testteamsp1"]
 
     # Create connection and indexer configurations
-    access_config = ConfluenceAccessConfig(api_token=api_token)
+    access_config = ConfluenceAccessConfig(password=api_token)
     connection_config = ConfluenceConnectionConfig(
         url=confluence_url,
-        user_email=user_email,
+        username=user_email,
         access_config=access_config,
     )
     index_config = ConfluenceIndexerConfig(

test/integration/connectors/test_pinecone.py

@@ -107,11 +107,15 @@ def pinecone_index() -> Generator[str, None, None]:
 
 
 def validate_pinecone_index(
-    index_name: str, expected_num_of_vectors: int, retries=30, interval=1
+    index_name: str,
+    expected_num_of_vectors: int,
+    retries=30,
+    interval=1,
+    namespace: str = "default",
 ) -> None:
     # Because there's a delay for the index to catch up to the recent writes, add in a retry
     pinecone = Pinecone(api_key=get_api_key())
-    index = pinecone.Index(name=index_name)
+    index = pinecone.Index(name=index_name, namespace=namespace)
     vector_count = -1
     for i in range(retries):
         index_stats = index.describe_index_stats()
@@ -133,11 +137,13 @@ def validate_pinecone_index(
 @pytest.mark.asyncio
 @pytest.mark.tags(CONNECTOR_TYPE, DESTINATION_TAG, VECTOR_DB_TAG)
 async def test_pinecone_destination(pinecone_index: str, upload_file: Path, temp_dir: Path):
+
     file_data = FileData(
         source_identifiers=SourceIdentifiers(fullpath=upload_file.name, filename=upload_file.name),
         connector_type=CONNECTOR_TYPE,
         identifier="pinecone_mock_id",
     )
+
     connection_config = PineconeConnectionConfig(
         index_name=pinecone_index,
         access_config=PineconeAccessConfig(api_key=get_api_key()),
@@ -224,6 +230,66 @@ async def test_pinecone_destination_large_index(
     )
 
 
+@requires_env(API_KEY)
+@pytest.mark.asyncio
+@pytest.mark.tags(CONNECTOR_TYPE, DESTINATION_TAG, VECTOR_DB_TAG)
+async def test_pinecone_destination_namespace(
+    pinecone_index: str, upload_file: Path, temp_dir: Path
+):
+    """
+    tests namespace functionality of destination connector.
+    """
+
+    # creates a file data structure.
+    file_data = FileData(
+        source_identifiers=SourceIdentifiers(fullpath=upload_file.name, filename=upload_file.name),
+        connector_type=CONNECTOR_TYPE,
+        identifier="pinecone_mock_id",
+    )
+
+    connection_config = PineconeConnectionConfig(
+        index_name=pinecone_index,
+        access_config=PineconeAccessConfig(api_key=get_api_key()),
+    )
+
+    stager_config = PineconeUploadStagerConfig()
+
+    stager = PineconeUploadStager(upload_stager_config=stager_config)
+    new_upload_file = stager.run(
+        elements_filepath=upload_file,
+        output_dir=temp_dir,
+        output_filename=upload_file.name,
+        file_data=file_data,
+    )
+
+    # here add namespace defintion
+    upload_config = PineconeUploaderConfig()
+    namespace_test_name = "user-1"
+    upload_config.namespace = namespace_test_name
+    uploader = PineconeUploader(connection_config=connection_config, upload_config=upload_config)
+    uploader.precheck()
+
+    uploader.run(path=new_upload_file, file_data=file_data)
+    with new_upload_file.open() as f:
+        staged_content = json.load(f)
+    expected_num_of_vectors = len(staged_content)
+    logger.info("validating first upload")
+    validate_pinecone_index(
+        index_name=pinecone_index,
+        expected_num_of_vectors=expected_num_of_vectors,
+        namespace=namespace_test_name,
+    )
+
+    # Rerun uploader and make sure no duplicates exist
+    uploader.run(path=new_upload_file, file_data=file_data)
+    logger.info("validating second upload")
+    validate_pinecone_index(
+        index_name=pinecone_index,
+        expected_num_of_vectors=expected_num_of_vectors,
+        namespace=namespace_test_name,
+    )
+
+
 @requires_env(API_KEY)
 @pytest.mark.tags(CONNECTOR_TYPE, DESTINATION_TAG, VECTOR_DB_TAG)
 def test_large_metadata(pinecone_index: str, tmp_path: Path, upload_file: Path):

test/unit/v2/connectors/sql/test_sql.py

@@ -0,0 +1,72 @@
+from pathlib import Path
+
+import pytest
+from pytest_mock import MockerFixture
+
+from unstructured_ingest.v2.interfaces.file_data import FileData, SourceIdentifiers
+from unstructured_ingest.v2.processes.connectors.sql.sql import SQLUploadStager
+
+
+@pytest.fixture
+def mock_instance() -> SQLUploadStager:
+    return SQLUploadStager()
+
+
+@pytest.mark.parametrize(
+    ("input_filepath", "output_filename", "expected"),
+    [
+        (
+            "/path/to/input_file.ndjson",
+            "output_file.ndjson",
+            "output_file.ndjson",
+        ),
+        ("input_file.txt", "output_file.json", "output_file.txt"),
+        ("/path/to/input_file.json", "output_file", "output_file.json"),
+    ],
+)
+def test_run_output_filename_suffix(
+    mocker: MockerFixture,
+    mock_instance: SQLUploadStager,
+    input_filepath: str,
+    output_filename: str,
+    expected: str,
+):
+    output_dir = Path("/tmp/test/output_dir")
+
+    # Mocks
+    mock_get_data = mocker.patch(
+        "unstructured_ingest.v2.processes.connectors.sql.sql.get_data",
+        return_value=[{"key": "value"}, {"key": "value2"}],
+    )
+    mock_conform_dict = mocker.patch.object(
+        SQLUploadStager, "conform_dict", side_effect=lambda element_dict, file_data: element_dict
+    )
+    mock_conform_dataframe = mocker.patch.object(
+        SQLUploadStager, "conform_dataframe", side_effect=lambda df: df
+    )
+    mock_get_output_path = mocker.patch.object(
+        SQLUploadStager, "get_output_path", return_value=output_dir / expected
+    )
+    mock_write_output = mocker.patch.object(SQLUploadStager, "write_output")
+
+    # Act
+    result = mock_instance.run(
+        elements_filepath=Path(input_filepath),
+        file_data=FileData(
+            identifier="test",
+            connector_type="test",
+            source_identifiers=SourceIdentifiers(filename=input_filepath, fullpath=input_filepath),
+        ),
+        output_dir=output_dir,
+        output_filename=output_filename,
+    )
+
+    # Assert
+    mock_get_data.assert_called_once_with(path=Path(input_filepath))
+    assert mock_conform_dict.call_count == 2
+    mock_conform_dataframe.assert_called_once()
+    mock_get_output_path.assert_called_once_with(output_filename=expected, output_dir=output_dir)
+    mock_write_output.assert_called_once_with(
+        output_path=output_dir / expected, data=[{"key": "value"}, {"key": "value2"}]
+    )
+    assert result.name == expected

test/unit/v2/connectors/test_confluence.py

@@ -11,10 +11,10 @@ def test_connection_config_multiple_auth():
     with pytest.raises(ValidationError):
         ConfluenceConnectionConfig(
             access_config=ConfluenceAccessConfig(
-                api_token="api_token",
-                access_token="access_token",
+                password="api_token",
+                token="access_token",
             ),
-            user_email="user_email",
+            username="user_email",
             url="url",
         )
 
@@ -26,14 +26,14 @@ def test_connection_config_no_auth():
 
 def test_connection_config_basic_auth():
     ConfluenceConnectionConfig(
-        access_config=ConfluenceAccessConfig(api_token="api_token"),
+        access_config=ConfluenceAccessConfig(password="api_token"),
         url="url",
-        user_email="user_email",
+        username="user_email",
     )
 
 
 def test_connection_config_pat_auth():
     ConfluenceConnectionConfig(
-        access_config=ConfluenceAccessConfig(access_token="access_token"),
+        access_config=ConfluenceAccessConfig(token="access_token"),
         url="url",
     )

unstructured_ingest/__version__.py

@@ -1 +1 @@
-__version__ = "0.3.14"  # pragma: no cover
+__version__ = "0.4.0"  # pragma: no cover

unstructured_ingest/v2/interfaces/upload_stager.py

@@ -2,7 +2,7 @@ import json
 from abc import ABC
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Any, TypeVar
+from typing import Any, Optional, TypeVar
 
 import ndjson
 from pydantic import BaseModel
@@ -22,10 +22,10 @@ UploadStagerConfigT = TypeVar("UploadStagerConfigT", bound=UploadStagerConfig)
 class UploadStager(BaseProcess, ABC):
     upload_stager_config: UploadStagerConfigT
 
-    def write_output(self, output_path: Path, data: list[dict]) -> None:
+    def write_output(self, output_path: Path, data: list[dict], indent: Optional[int] = 2) -> None:
         if output_path.suffix == ".json":
             with output_path.open("w") as f:
-                json.dump(data, f, indent=2)
+                json.dump(data, f, indent=indent)
         elif output_path.suffix == ".ndjson":
             with output_path.open("w") as f:
                 ndjson.dump(data, f)

unstructured_ingest/v2/processes/connectors/confluence.py

@@ -30,27 +30,45 @@ CONNECTOR_TYPE = "confluence"
 
 
 class ConfluenceAccessConfig(AccessConfig):
-    api_token: Optional[str] = Field(description="Confluence API token", default=None)
-    access_token: Optional[str] = Field(
-        description="Confluence Personal Access Token", default=None
+    password: Optional[str] = Field(
+        description="Confluence password or Cloud API token",
+        default=None,
+    )
+    token: Optional[str] = Field(
+        description="Confluence Personal Access Token",
+        default=None,
     )
 
 
 class ConfluenceConnectionConfig(ConnectionConfig):
     url: str = Field(description="URL of the Confluence instance")
-    user_email: Optional[str] = Field(description="User email for authentication", default=None)
+    username: Optional[str] = Field(
+        description="Username or email for authentication",
+        default=None,
+    )
+    cloud: bool = Field(description="Authenticate to Confluence Cloud", default=False)
     access_config: Secret[ConfluenceAccessConfig] = Field(
         description="Access configuration for Confluence"
     )
 
     def model_post_init(self, __context):
         access_configs = self.access_config.get_secret_value()
-        basic_auth = self.user_email and access_configs.api_token
-        pat_auth = access_configs.access_token
+        basic_auth = self.username and access_configs.password
+        pat_auth = access_configs.token
+        if self.cloud and not basic_auth:
+            raise ValueError(
+                "cloud authentication requires username and API token (--password), "
+                "see: https://atlassian-python-api.readthedocs.io/"
+            )
         if basic_auth and pat_auth:
-            raise ValueError("both forms of auth provided, only one allowed")
+            raise ValueError(
+                "both password and token provided, only one allowed, "
+                "see: https://atlassian-python-api.readthedocs.io/"
+            )
         if not (basic_auth or pat_auth):
-            raise ValueError("neither forms of auth provided")
+            raise ValueError(
+                "no form of auth provided, see: https://atlassian-python-api.readthedocs.io/"
+            )
 
     @requires_dependencies(["atlassian"], extras="confluence")
     def get_client(self) -> "Confluence":
@@ -59,8 +77,10 @@ class ConfluenceConnectionConfig(ConnectionConfig):
         access_configs = self.access_config.get_secret_value()
         return Confluence(
             url=self.url,
-            username=self.user_email,
-            password=access_configs.api_token,
+            username=self.username,
+            password=access_configs.password,
+            token=access_configs.token,
+            cloud=self.cloud,
         )
 
 

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

@@ -25,6 +25,8 @@ from .volumes_native import (
     databricks_native_volumes_destination_entry,
     databricks_native_volumes_source_entry,
 )
+from .volumes_table import CONNECTOR_TYPE as VOLUMES_TABLE_CONNECTOR_TYPE
+from .volumes_table import databricks_volumes_delta_tables_destination_entry
 
 add_source_entry(source_type=VOLUMES_AWS_CONNECTOR_TYPE, entry=databricks_aws_volumes_source_entry)
 add_destination_entry(
@@ -50,3 +52,7 @@ add_source_entry(
 add_destination_entry(
     destination_type=VOLUMES_AZURE_CONNECTOR_TYPE, entry=databricks_azure_volumes_destination_entry
 )
+add_destination_entry(
+    destination_type=VOLUMES_TABLE_CONNECTOR_TYPE,
+    entry=databricks_volumes_delta_tables_destination_entry,
+)

unstructured_ingest/v2/processes/connectors/databricks/volumes.py

@@ -187,6 +187,11 @@ class DatabricksVolumesUploader(Uploader, ABC):
     upload_config: DatabricksVolumesUploaderConfig
     connection_config: DatabricksVolumesConnectionConfig
 
+    def get_output_path(self, file_data: FileData) -> str:
+        return os.path.join(
+            self.upload_config.path, f"{file_data.source_identifiers.filename}.json"
+        )
+
     def precheck(self) -> None:
         try:
             assert self.connection_config.get_client().current_user.me().active
@@ -194,9 +199,7 @@ class DatabricksVolumesUploader(Uploader, ABC):
             raise self.connection_config.wrap_error(e=e)
 
     def run(self, path: Path, file_data: FileData, **kwargs: Any) -> None:
-        output_path = os.path.join(
-            self.upload_config.path, f"{file_data.source_identifiers.filename}.json"
-        )
+        output_path = self.get_output_path(file_data=file_data)
         with open(path, "rb") as elements_file:
             try:
                 self.connection_config.get_client().files.upload(

unstructured_ingest/v2/processes/connectors/databricks/volumes_table.py

@@ -0,0 +1,106 @@
+import json
+import os
+from contextlib import contextmanager
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Generator, Optional
+
+from pydantic import Field
+
+from unstructured_ingest.v2.interfaces import FileData, Uploader, UploaderConfig
+from unstructured_ingest.v2.logger import logger
+from unstructured_ingest.v2.processes.connector_registry import (
+    DestinationRegistryEntry,
+)
+from unstructured_ingest.v2.processes.connectors.databricks.volumes import DatabricksPathMixin
+from unstructured_ingest.v2.processes.connectors.sql.databricks_delta_tables import (
+    DatabrickDeltaTablesConnectionConfig,
+    DatabrickDeltaTablesUploadStager,
+    DatabrickDeltaTablesUploadStagerConfig,
+)
+
+CONNECTOR_TYPE = "databricks_volume_delta_tables"
+
+
+class DatabricksVolumeDeltaTableUploaderConfig(UploaderConfig, DatabricksPathMixin):
+    database: str = Field(description="Database name", default="default")
+    table_name: str = Field(description="Table name")
+
+
+@dataclass
+class DatabricksVolumeDeltaTableStager(DatabrickDeltaTablesUploadStager):
+    def write_output(self, output_path: Path, data: list[dict], indent: Optional[int] = 2) -> None:
+        # To avoid new line issues when migrating from volumes into delta tables, omit indenting
+        # and always write it as a json file
+        with output_path.with_suffix(".json").open("w") as f:
+            json.dump(data, f)
+
+
+@dataclass
+class DatabricksVolumeDeltaTableUploader(Uploader):
+    connection_config: DatabrickDeltaTablesConnectionConfig
+    upload_config: DatabricksVolumeDeltaTableUploaderConfig
+    connector_type: str = CONNECTOR_TYPE
+
+    def precheck(self) -> None:
+        with self.connection_config.get_cursor() as cursor:
+            cursor.execute("SHOW CATALOGS")
+            catalogs = [r[0] for r in cursor.fetchall()]
+            if self.upload_config.catalog not in catalogs:
+                raise ValueError(
+                    "Catalog {} not found in {}".format(
+                        self.upload_config.catalog, ", ".join(catalogs)
+                    )
+                )
+            cursor.execute(f"USE CATALOG '{self.upload_config.catalog}'")
+            cursor.execute("SHOW DATABASES")
+            databases = [r[0] for r in cursor.fetchall()]
+            if self.upload_config.database not in databases:
+                raise ValueError(
+                    "Database {} not found in {}".format(
+                        self.upload_config.database, ", ".join(databases)
+                    )
+                )
+            cursor.execute("SHOW TABLES")
+            table_names = [r[1] for r in cursor.fetchall()]
+            if self.upload_config.table_name not in table_names:
+                raise ValueError(
+                    "Table {} not found in {}".format(
+                        self.upload_config.table_name, ", ".join(table_names)
+                    )
+                )
+
+    def get_output_path(self, file_data: FileData, suffix: str = ".json") -> str:
+        filename = Path(file_data.source_identifiers.filename)
+        adjusted_filename = filename if filename.suffix == suffix else f"{filename}{suffix}"
+        return os.path.join(self.upload_config.path, f"{adjusted_filename}")
+
+    @contextmanager
+    def get_cursor(self, **connect_kwargs) -> Generator[Any, None, None]:
+        with self.connection_config.get_cursor(**connect_kwargs) as cursor:
+            cursor.execute(f"USE CATALOG '{self.upload_config.catalog}'")
+            yield cursor
+
+    def run(self, path: Path, file_data: FileData, **kwargs: Any) -> None:
+        with self.get_cursor(staging_allowed_local_path=str(path.parent)) as cursor:
+            catalog_path = self.get_output_path(file_data=file_data)
+            logger.debug(f"uploading {path.as_posix()} to {catalog_path}")
+            cursor.execute(f"PUT '{path.as_posix()}' INTO '{catalog_path}' OVERWRITE")
+            logger.debug(
+                f"migrating content from {catalog_path} to table {self.upload_config.table_name}"
+            )
+            with path.open() as f:
+                data = json.load(f)
+                columns = data[0].keys()
+            column_str = ", ".join(columns)
+            sql_statment = f"INSERT INTO `{self.upload_config.table_name}` ({column_str}) SELECT {column_str} FROM json.`{catalog_path}`"  # noqa: E501
+            cursor.execute(sql_statment)
+
+
+databricks_volumes_delta_tables_destination_entry = DestinationRegistryEntry(
+    connection_config=DatabrickDeltaTablesConnectionConfig,
+    uploader=DatabricksVolumeDeltaTableUploader,
+    uploader_config=DatabricksVolumeDeltaTableUploaderConfig,
+    upload_stager=DatabricksVolumeDeltaTableStager,
+    upload_stager_config=DatabrickDeltaTablesUploadStagerConfig,
+)