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

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

@@ -0,0 +1,162 @@
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+from pathlib import Path
+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.processes.connector_registry import (
+    DestinationRegistryEntry,
+    SourceRegistryEntry,
+)
+from unstructured_ingest.v2.processes.connectors.sql.postgres import (
+    PostgresDownloader,
+    PostgresDownloaderConfig,
+    PostgresIndexer,
+    PostgresIndexerConfig,
+    PostgresUploader,
+    PostgresUploaderConfig,
+    PostgresUploadStager,
+    PostgresUploadStagerConfig,
+)
+from unstructured_ingest.v2.processes.connectors.sql.sql import SQLAccessConfig, SQLConnectionConfig
+
+if TYPE_CHECKING:
+    from snowflake.connector import SnowflakeConnection
+    from snowflake.connector.cursor import SnowflakeCursor
+
+CONNECTOR_TYPE = "snowflake"
+
+
+class SnowflakeAccessConfig(SQLAccessConfig):
+    password: Optional[str] = Field(default=None, description="DB password")
+
+
+class SnowflakeConnectionConfig(SQLConnectionConfig):
+    access_config: Secret[SnowflakeAccessConfig] = Field(
+        default=SnowflakeAccessConfig(), validate_default=True
+    )
+    account: str = Field(
+        default=None,
+        description="Your account identifier. The account identifier "
+        "does not include the snowflakecomputing.com suffix.",
+    )
+    user: Optional[str] = Field(default=None, description="DB username")
+    host: Optional[str] = Field(default=None, description="DB host")
+    port: Optional[int] = Field(default=443, description="DB host connection port")
+    database: str = Field(
+        default=None,
+        description="Database name.",
+    )
+    db_schema: str = Field(default=None, description="Database schema.", alias="schema")
+    role: str = Field(
+        default=None,
+        description="Database role.",
+    )
+    connector_type: str = Field(default=CONNECTOR_TYPE, init=False)
+
+    @contextmanager
+    @requires_dependencies(["snowflake"], extras="snowflake")
+    def get_connection(self) -> Generator["SnowflakeConnection", None, None]:
+        # https://docs.snowflake.com/en/developer-guide/python-connector/python-connector-api#label-snowflake-connector-methods-connect
+        from snowflake.connector import connect
+
+        connect_kwargs = self.model_dump()
+        connect_kwargs["schema"] = connect_kwargs.pop("db_schema")
+        connect_kwargs.pop("access_configs", None)
+        connect_kwargs["password"] = self.access_config.get_secret_value().password
+        # https://peps.python.org/pep-0249/#paramstyle
+        connect_kwargs["paramstyle"] = "qmark"
+        connection = connect(**connect_kwargs)
+        try:
+            yield connection
+        finally:
+            connection.commit()
+            connection.close()
+
+    @contextmanager
+    def get_cursor(self) -> Generator["SnowflakeCursor", None, None]:
+        with self.get_connection() as connection:
+            cursor = connection.cursor()
+            try:
+                yield cursor
+            finally:
+                cursor.close()
+
+
+class SnowflakeIndexerConfig(PostgresIndexerConfig):
+    pass
+
+
+@dataclass
+class SnowflakeIndexer(PostgresIndexer):
+    connection_config: SnowflakeConnectionConfig
+    index_config: SnowflakeIndexerConfig
+    connector_type: str = CONNECTOR_TYPE
+
+
+class SnowflakeDownloaderConfig(PostgresDownloaderConfig):
+    pass
+
+
+@dataclass
+class SnowflakeDownloader(PostgresDownloader):
+    connection_config: SnowflakeConnectionConfig
+    download_config: SnowflakeDownloaderConfig
+    connector_type: str = CONNECTOR_TYPE
+
+
+class SnowflakeUploadStagerConfig(PostgresUploadStagerConfig):
+    pass
+
+
+class SnowflakeUploadStager(PostgresUploadStager):
+    upload_stager_config: SnowflakeUploadStagerConfig
+
+
+class SnowflakeUploaderConfig(PostgresUploaderConfig):
+    pass
+
+
+@dataclass
+class SnowflakeUploader(PostgresUploader):
+    upload_config: SnowflakeUploaderConfig = field(default_factory=SnowflakeUploaderConfig)
+    connection_config: SnowflakeConnectionConfig
+    connector_type: str = CONNECTOR_TYPE
+    values_delimiter: str = "?"
+
+    def upload_contents(self, path: Path) -> None:
+        df = pd.read_json(path, orient="records", lines=True)
+        df.replace({np.nan: None}, inplace=True)
+
+        columns = list(df.columns)
+        stmt = f"INSERT INTO {self.upload_config.table_name} ({','.join(columns)}) VALUES({','.join([self.values_delimiter 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_cursor() as cursor:
+                values = self.prepare_data(columns, tuple(rows.itertuples(index=False, name=None)))
+                # TODO: executemany break on 'Binding data in type (list) is not supported'
+                for val in values:
+                    cursor.execute(stmt, val)
+
+
+snowflake_source_entry = SourceRegistryEntry(
+    connection_config=SnowflakeConnectionConfig,
+    indexer_config=SnowflakeIndexerConfig,
+    indexer=SnowflakeIndexer,
+    downloader_config=SnowflakeDownloaderConfig,
+    downloader=SnowflakeDownloader,
+)
+
+snowflake_destination_entry = DestinationRegistryEntry(
+    connection_config=SnowflakeConnectionConfig,
+    uploader=SnowflakeUploader,
+    uploader_config=SnowflakeUploaderConfig,
+    upload_stager=SnowflakeUploadStager,
+    upload_stager_config=SnowflakeUploadStagerConfig,
+)

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

@@ -3,12 +3,14 @@ import json
 import sys
 import uuid
 from abc import ABC, abstractmethod
+from contextlib import contextmanager
 from dataclasses import dataclass, field, replace
 from datetime import date, datetime
 from pathlib import Path
 from time import time
 from typing import Any, Generator, Union
 
+import numpy as np
 import pandas as pd
 from dateutil import parser
 from pydantic import Field, Secret
@@ -94,7 +96,13 @@ class SQLConnectionConfig(ConnectionConfig, ABC):
     access_config: Secret[SQLAccessConfig] = Field(default=SQLAccessConfig(), validate_default=True)
 
     @abstractmethod
-    def get_connection(self) -> Any:
+    @contextmanager
+    def get_connection(self) -> Generator[Any, None, None]:
+        pass
+
+    @abstractmethod
+    @contextmanager
+    def get_cursor(self) -> Generator[Any, None, None]:
         pass
 
 
@@ -108,16 +116,19 @@ class SQLIndexer(Indexer, ABC):
     connection_config: SQLConnectionConfig
     index_config: SQLIndexerConfig
 
-    @abstractmethod
     def _get_doc_ids(self) -> list[str]:
-        pass
+        with self.connection_config.get_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
 
     def precheck(self) -> None:
         try:
-            connection = self.connection_config.get_connection()
-            cursor = connection.cursor()
-            cursor.execute("SELECT 1;")
-            cursor.close()
+            with self.connection_config.get_cursor() as cursor:
+                cursor.execute("SELECT 1;")
         except Exception as e:
             logger.error(f"failed to validate connection: {e}", exc_info=True)
             raise SourceConnectionError(f"failed to validate connection: {e}")
@@ -198,7 +209,7 @@ class SQLDownloader(Downloader, ABC):
             f"Downloading results from table {table_name} and id {record_id} to {download_path}"
         )
         download_path.parent.mkdir(parents=True, exist_ok=True)
-        result.to_csv(download_path)
+        result.to_csv(download_path, index=False)
         copied_file_data = replace(file_data)
         copied_file_data.identifier = filename_id
         copied_file_data.doc_type = "file"
@@ -285,6 +296,7 @@ class SQLUploaderConfig(UploaderConfig):
 class SQLUploader(Uploader):
     upload_config: SQLUploaderConfig
     connection_config: SQLConnectionConfig
+    values_delimiter: str = "?"
 
     def precheck(self) -> None:
         try:
@@ -296,15 +308,42 @@ class SQLUploader(Uploader):
             logger.error(f"failed to validate connection: {e}", exc_info=True)
             raise DestinationConnectionError(f"failed to validate connection: {e}")
 
-    @abstractmethod
     def prepare_data(
         self, columns: list[str], data: tuple[tuple[Any, ...], ...]
     ) -> list[tuple[Any, ...]]:
-        pass
+        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
 
-    @abstractmethod
     def upload_contents(self, path: Path) -> None:
-        pass
+        df = pd.read_json(path, orient="records", lines=True)
+        df.replace({np.nan: None}, inplace=True)
+
+        columns = list(df.columns)
+        stmt = f"INSERT INTO {self.upload_config.table_name} ({','.join(columns)}) VALUES({','.join([self.values_delimiter 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_cursor() as cursor:
+                values = self.prepare_data(columns, tuple(rows.itertuples(index=False, name=None)))
+                # for val in values:
+                #     try:
+                #         cursor.execute(stmt, val)
+                #     except Exception as e:
+                #         print(f"Error: {e}")
+                #         print(f"failed to write {len(columns)}, {len(val)}: {stmt} -> {val}")
+                cursor.executemany(stmt, values)
 
     def run(self, path: Path, file_data: FileData, **kwargs: Any) -> None:
         self.upload_contents(path=path)

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

@@ -1,15 +1,17 @@
 import json
+from contextlib import contextmanager
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Generator
 
-import numpy as np
-import pandas as pd
 from pydantic import Field, Secret, model_validator
 
 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,
@@ -27,6 +29,7 @@ from unstructured_ingest.v2.processes.connectors.sql.sql import (
 
 if TYPE_CHECKING:
     from sqlite3 import Connection as SqliteConnection
+    from sqlite3 import Cursor as SqliteCursor
 
 CONNECTOR_TYPE = "sqlite"
 
@@ -51,10 +54,25 @@ class SQLiteConnectionConfig(SQLConnectionConfig):
             raise ValueError(f"{self.database_path} is not a valid file")
         return self
 
-    def get_connection(self) -> "SqliteConnection":
+    @contextmanager
+    def get_connection(self) -> Generator["SqliteConnection", None, None]:
         from sqlite3 import connect
 
-        return connect(database=self.database_path)
+        connection = connect(database=self.database_path)
+        try:
+            yield connection
+        finally:
+            connection.commit()
+            connection.close()
+
+    @contextmanager
+    def get_cursor(self) -> Generator["SqliteCursor", None, None]:
+        with self.get_connection() as connection:
+            cursor = connection.cursor()
+            try:
+                yield cursor
+            finally:
+                cursor.close()
 
 
 class SQLiteIndexerConfig(SQLIndexerConfig):
@@ -67,16 +85,6 @@ class SQLiteIndexer(SQLIndexer):
     index_config: SQLIndexerConfig
     connector_type: str = CONNECTOR_TYPE
 
-    def _get_doc_ids(self) -> list[str]:
-        with self.connection_config.get_connection() as sqlite_connection:
-            cursor = sqlite_connection.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 SQLiteDownloaderConfig(SQLDownloaderConfig):
     pass
@@ -145,23 +153,14 @@ class SQLiteUploader(SQLUploader):
             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_path} ")
-        df.replace({np.nan: None}, inplace=True)
-
-        columns = tuple(df.columns)
-        stmt = f"INSERT INTO {self.upload_config.table_name} ({','.join(columns)}) \
-                VALUES({','.join(['?' 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)))
-                conn.executemany(stmt, values)
-                conn.commit()
 
+sqlite_source_entry = SourceRegistryEntry(
+    connection_config=SQLiteConnectionConfig,
+    indexer_config=SQLiteIndexerConfig,
+    indexer=SQLIndexer,
+    downloader_config=SQLiteDownloaderConfig,
+    downloader=SQLiteDownloader,
+)
 
 sqlite_destination_entry = DestinationRegistryEntry(
     connection_config=SQLiteConnectionConfig,

unstructured_ingest/v2/unstructured_api.py

@@ -26,7 +26,7 @@ def create_partition_request(filename: Path, parameters_dict: dict) -> "Partitio
     # NOTE(austin): PartitionParameters is a Pydantic model in v0.26.0
     # Prior to this it was a dataclass which doesn't have .__fields
     try:
-        possible_fields = PartitionParameters.__fields__
+        possible_fields = PartitionParameters.model_fields
     except AttributeError:
         possible_fields = [f.name for f in fields(PartitionParameters)]
 

{unstructured_ingest-0.1.1.dist-info → unstructured_ingest-0.2.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: unstructured-ingest
-Version: 0.1.1
+Version: 0.2.1
 Summary: A library that prepares raw documents for downstream ML tasks.
 Home-page: https://github.com/Unstructured-IO/unstructured-ingest
 Author: Unstructured Technologies
@@ -22,30 +22,30 @@ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Requires-Python: >=3.9.0,<3.13
 Description-Content-Type: text/markdown
 License-File: LICENSE.md
+Requires-Dist: tqdm
 Requires-Dist: python-dateutil
 Requires-Dist: pandas
+Requires-Dist: click
 Requires-Dist: pydantic>=2.7
 Requires-Dist: dataclasses-json
 Requires-Dist: opentelemetry-sdk
-Requires-Dist: click
-Requires-Dist: tqdm
 Provides-Extra: airtable
 Requires-Dist: pyairtable; extra == "airtable"
 Provides-Extra: astradb
 Requires-Dist: astrapy; extra == "astradb"
 Provides-Extra: azure
-Requires-Dist: adlfs; extra == "azure"
 Requires-Dist: fsspec; extra == "azure"
+Requires-Dist: adlfs; extra == "azure"
 Provides-Extra: azure-cognitive-search
 Requires-Dist: azure-search-documents; extra == "azure-cognitive-search"
 Provides-Extra: bedrock
 Requires-Dist: boto3; extra == "bedrock"
 Provides-Extra: biomed
-Requires-Dist: bs4; extra == "biomed"
 Requires-Dist: requests; extra == "biomed"
+Requires-Dist: bs4; extra == "biomed"
 Provides-Extra: box
-Requires-Dist: boxfs; extra == "box"
 Requires-Dist: fsspec; extra == "box"
+Requires-Dist: boxfs; extra == "box"
 Provides-Extra: chroma
 Requires-Dist: chromadb; extra == "chroma"
 Provides-Extra: clarifai
@@ -60,8 +60,8 @@ Requires-Dist: unstructured[tsv]; extra == "csv"
 Provides-Extra: databricks-volumes
 Requires-Dist: databricks-sdk; extra == "databricks-volumes"
 Provides-Extra: delta-table
-Requires-Dist: deltalake; extra == "delta-table"
 Requires-Dist: fsspec; extra == "delta-table"
+Requires-Dist: deltalake; extra == "delta-table"
 Provides-Extra: discord
 Requires-Dist: discord-py; extra == "discord"
 Provides-Extra: doc
@@ -88,8 +88,8 @@ Provides-Extra: epub
 Requires-Dist: unstructured[epub]; extra == "epub"
 Provides-Extra: gcs
 Requires-Dist: gcsfs; extra == "gcs"
-Requires-Dist: bs4; extra == "gcs"
 Requires-Dist: fsspec; extra == "gcs"
+Requires-Dist: bs4; extra == "gcs"
 Provides-Extra: github
 Requires-Dist: pygithub>1.58.0; extra == "github"
 Requires-Dist: requests; extra == "github"
@@ -98,8 +98,8 @@ Requires-Dist: python-gitlab; extra == "gitlab"
 Provides-Extra: google-drive
 Requires-Dist: google-api-python-client; extra == "google-drive"
 Provides-Extra: hubspot
-Requires-Dist: hubspot-api-client; extra == "hubspot"
 Requires-Dist: urllib3; extra == "hubspot"
+Requires-Dist: hubspot-api-client; extra == "hubspot"
 Provides-Extra: jira
 Requires-Dist: atlassian-python-api; extra == "jira"
 Provides-Extra: kafka
@@ -115,16 +115,16 @@ Requires-Dist: pymongo; extra == "mongodb"
 Provides-Extra: msg
 Requires-Dist: unstructured[msg]; extra == "msg"
 Provides-Extra: notion
-Requires-Dist: notion-client; extra == "notion"
-Requires-Dist: htmlBuilder; extra == "notion"
-Requires-Dist: backoff; extra == "notion"
 Requires-Dist: httpx; extra == "notion"
+Requires-Dist: backoff; extra == "notion"
+Requires-Dist: htmlBuilder; extra == "notion"
+Requires-Dist: notion-client; extra == "notion"
 Provides-Extra: odt
 Requires-Dist: unstructured[odt]; extra == "odt"
 Provides-Extra: onedrive
-Requires-Dist: Office365-REST-Python-Client; extra == "onedrive"
-Requires-Dist: bs4; extra == "onedrive"
 Requires-Dist: msal; extra == "onedrive"
+Requires-Dist: bs4; extra == "onedrive"
+Requires-Dist: Office365-REST-Python-Client; extra == "onedrive"
 Provides-Extra: openai
 Requires-Dist: openai; extra == "openai"
 Requires-Dist: tiktoken; extra == "openai"
@@ -133,8 +133,8 @@ Requires-Dist: opensearch-py; extra == "opensearch"
 Provides-Extra: org
 Requires-Dist: unstructured[org]; extra == "org"
 Provides-Extra: outlook
-Requires-Dist: Office365-REST-Python-Client; extra == "outlook"
 Requires-Dist: msal; extra == "outlook"
+Requires-Dist: Office365-REST-Python-Client; extra == "outlook"
 Provides-Extra: pdf
 Requires-Dist: unstructured[pdf]; extra == "pdf"
 Provides-Extra: pinecone
@@ -164,12 +164,14 @@ Provides-Extra: sftp
 Requires-Dist: fsspec; extra == "sftp"
 Requires-Dist: paramiko; extra == "sftp"
 Provides-Extra: sharepoint
-Requires-Dist: Office365-REST-Python-Client; extra == "sharepoint"
 Requires-Dist: msal; extra == "sharepoint"
+Requires-Dist: Office365-REST-Python-Client; extra == "sharepoint"
 Provides-Extra: singlestore
 Requires-Dist: singlestoredb; extra == "singlestore"
 Provides-Extra: slack
-Requires-Dist: slack-sdk; extra == "slack"
+Requires-Dist: slack-sdk[optional]; extra == "slack"
+Provides-Extra: snowflake
+Requires-Dist: snowflake; extra == "snowflake"
 Provides-Extra: togetherai
 Requires-Dist: together; extra == "togetherai"
 Provides-Extra: tsv