unstructured-ingest 0.3.15__py3-none-any.whl → 0.4.1__py3-none-any.whl

unstructured_ingest/v2/processes/connectors/confluence.py

@@ -1,3 +1,4 @@
+from contextlib import contextmanager
 from dataclasses import dataclass, field
 from pathlib import Path
 from typing import TYPE_CHECKING, Generator, List, Optional
@@ -17,6 +18,7 @@ from unstructured_ingest.v2.interfaces import (
     Indexer,
     IndexerConfig,
     SourceIdentifiers,
+    download_responses,
 )
 from unstructured_ingest.v2.logger import logger
 from unstructured_ingest.v2.processes.connector_registry import (
@@ -30,38 +32,60 @@ 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")
+    @contextmanager
     def get_client(self) -> "Confluence":
         from atlassian import Confluence
 
         access_configs = self.access_config.get_secret_value()
-        return Confluence(
+        with 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,
+        ) as client:
+            yield client
 
 
 class ConfluenceIndexerConfig(IndexerConfig):
@@ -83,8 +107,8 @@ class ConfluenceIndexer(Indexer):
 
             # Attempt to retrieve a list of spaces with limit=1.
             # This should only succeed if all creds are valid
-            client = self.connection_config.get_client()
-            client.get_all_spaces(limit=1)
+            with self.connection_config.get_client() as client:
+                client.get_all_spaces(limit=1)
             logger.info("Connection to Confluence successful.")
             return True
         except Exception as e:
@@ -96,21 +120,21 @@ class ConfluenceIndexer(Indexer):
         if spaces:
             return spaces
         else:
-            client = self.connection_config.get_client()
-            all_spaces = client.get_all_spaces(limit=self.index_config.max_num_of_spaces)
+            with self.connection_config.get_client() as client:
+                all_spaces = client.get_all_spaces(limit=self.index_config.max_num_of_spaces)
             space_ids = [space["key"] for space in all_spaces["results"]]
             return space_ids
 
     def _get_docs_ids_within_one_space(self, space_id: str) -> List[dict]:
-        client = self.connection_config.get_client()
-        pages = client.get_all_pages_from_space(
-            space=space_id,
-            start=0,
-            limit=self.index_config.max_num_of_docs_from_each_space,
-            expand=None,
-            content_type="page",
-            status=None,
-        )
+        with self.connection_config.get_client() as client:
+            pages = client.get_all_pages_from_space(
+                space=space_id,
+                start=0,
+                limit=self.index_config.max_num_of_docs_from_each_space,
+                expand=None,
+                content_type="page",
+                status=None,
+            )
         doc_ids = [{"space_id": space_id, "doc_id": page["id"]} for page in pages]
         return doc_ids
 
@@ -157,7 +181,18 @@ class ConfluenceIndexer(Indexer):
 
 
 class ConfluenceDownloaderConfig(DownloaderConfig):
-    pass
+    extract_images: bool = Field(
+        default=False,
+        description="if true, will download images and replace "
+        "the html content with base64 encoded images",
+    )
+    extract_files: bool = Field(
+        default=False, description="if true, will download any embedded files"
+    )
+    force_download: bool = Field(
+        default=False,
+        description="if true, will redownload extracted files even if they already exist locally",
+    )
 
 
 @dataclass
@@ -166,14 +201,37 @@ class ConfluenceDownloader(Downloader):
     download_config: ConfluenceDownloaderConfig = field(default_factory=ConfluenceDownloaderConfig)
     connector_type: str = CONNECTOR_TYPE
 
-    def run(self, file_data: FileData, **kwargs) -> DownloadResponse:
+    def download_embedded_files(
+        self, session, html: str, current_file_data: FileData
+    ) -> list[DownloadResponse]:
+        if not self.download_config.extract_files:
+            return []
+        from unstructured_ingest.utils.html import download_embedded_files
+
+        filepath = current_file_data.source_identifiers.relative_path
+        download_path = Path(self.download_dir) / filepath
+        download_dir = download_path.with_suffix("")
+        return download_embedded_files(
+            download_dir=download_dir,
+            original_filedata=current_file_data,
+            original_html=html,
+            session=session,
+            force_download=self.download_config.force_download,
+        )
+
+    def run(self, file_data: FileData, **kwargs) -> download_responses:
+        from bs4 import BeautifulSoup
+
+        from unstructured_ingest.utils.html import convert_image_tags
+
         doc_id = file_data.identifier
         try:
-            client = self.connection_config.get_client()
-            page = client.get_page_by_id(
-                page_id=doc_id,
-                expand="history.lastUpdated,version,body.view",
-            )
+            with self.connection_config.get_client() as client:
+                page = client.get_page_by_id(
+                    page_id=doc_id,
+                    expand="history.lastUpdated,version,body.view",
+                )
+
         except Exception as e:
             logger.error(f"Failed to retrieve page with ID {doc_id}: {e}", exc_info=True)
             raise SourceConnectionError(f"Failed to retrieve page with ID {doc_id}: {e}")
@@ -182,20 +240,52 @@ class ConfluenceDownloader(Downloader):
             raise ValueError(f"Page with ID {doc_id} does not exist.")
 
         content = page["body"]["view"]["value"]
+        # This supports v2 html parsing in unstructured
+        title = page["title"]
+        title_html = f"<title>{title}</title>"
+        content = f"<body class='Document' >{title_html}{content}</body>"
+        if self.download_config.extract_images:
+            with self.connection_config.get_client() as client:
+                content = convert_image_tags(
+                    url=file_data.metadata.url, original_html=content, session=client._session
+                )
 
         filepath = file_data.source_identifiers.relative_path
         download_path = Path(self.download_dir) / filepath
         download_path.parent.mkdir(parents=True, exist_ok=True)
         with open(download_path, "w", encoding="utf8") as f:
-            f.write(content)
+            soup = BeautifulSoup(content, "html.parser")
+            f.write(soup.prettify())
 
         # Update file_data with metadata
         file_data.metadata.date_created = page["history"]["createdDate"]
         file_data.metadata.date_modified = page["version"]["when"]
         file_data.metadata.version = str(page["version"]["number"])
-        file_data.display_name = page["title"]
+        file_data.display_name = title
 
-        return self.generate_download_response(file_data=file_data, download_path=download_path)
+        download_response = self.generate_download_response(
+            file_data=file_data, download_path=download_path
+        )
+        if self.download_config.extract_files:
+            with self.connection_config.get_client() as client:
+                extracted_download_responses = self.download_embedded_files(
+                    html=content,
+                    current_file_data=download_response["file_data"],
+                    session=client._session,
+                )
+                if extracted_download_responses:
+                    for dr in extracted_download_responses:
+                        fd = dr["file_data"]
+                        source_file_path = Path(file_data.source_identifiers.fullpath).with_suffix(
+                            ""
+                        )
+                        new_fullpath = source_file_path / fd.source_identifiers.filename
+                        fd.source_identifiers = SourceIdentifiers(
+                            fullpath=new_fullpath.as_posix(), filename=new_fullpath.name
+                        )
+                    extracted_download_responses.append(download_response)
+                    return extracted_download_responses
+        return download_response
 
 
 confluence_source_entry = SourceRegistryEntry(

unstructured_ingest/v2/processes/connectors/duckdb/base.py

@@ -4,7 +4,7 @@ from typing import Any
 
 import pandas as pd
 
-from unstructured_ingest.utils.data_prep import get_data
+from unstructured_ingest.utils.data_prep import get_data, write_data
 from unstructured_ingest.v2.interfaces import FileData, UploadStager
 from unstructured_ingest.v2.utils import get_enhanced_element_id
 
@@ -96,5 +96,5 @@ class BaseDuckDBUploadStager(UploadStager):
             df[column] = df[column].apply(str)
 
         data = df.to_dict(orient="records")
-        self.write_output(output_path=output_path, data=data)
+        write_data(path=output_path, data=data)
         return output_path

unstructured_ingest/v2/processes/connectors/fsspec/azure.py

@@ -128,22 +128,22 @@ class AzureIndexer(FsspecIndexer):
     def sterilize_info(self, file_data: dict) -> dict:
         return sterilize_dict(data=file_data, default=azure_json_serial)
 
-    def get_metadata(self, file_data: dict) -> FileDataSourceMetadata:
-        path = file_data["name"]
+    def get_metadata(self, file_info: dict) -> FileDataSourceMetadata:
+        path = file_info["name"]
         date_created = (
-            str(file_data.get("creation_time").timestamp())
-            if "creation_time" in file_data
+            str(file_info.get("creation_time").timestamp())
+            if "creation_time" in file_info
             else None
         )
         date_modified = (
-            str(file_data.get("last_modified").timestamp())
-            if "last_modified" in file_data
+            str(file_info.get("last_modified").timestamp())
+            if "last_modified" in file_info
             else None
         )
 
-        file_size = file_data.get("size") if "size" in file_data else None
+        file_size = file_info.get("size") if "size" in file_info else None
 
-        version = file_data.get("etag")
+        version = file_info.get("etag")
         record_locator = {
             "protocol": self.index_config.protocol,
             "remote_file_path": self.index_config.remote_url,

unstructured_ingest/v2/processes/connectors/fsspec/box.py

@@ -104,22 +104,22 @@ class BoxIndexer(FsspecIndexer):
     index_config: BoxIndexerConfig
     connector_type: str = CONNECTOR_TYPE
 
-    def get_metadata(self, file_data: dict) -> FileDataSourceMetadata:
-        path = file_data["name"]
+    def get_metadata(self, file_info: dict) -> FileDataSourceMetadata:
+        path = file_info["name"]
         date_created = None
         date_modified = None
-        if modified_at_str := file_data.get("modified_at"):
+        if modified_at_str := file_info.get("modified_at"):
             date_modified = str(parser.parse(modified_at_str).timestamp())
-        if created_at_str := file_data.get("created_at"):
+        if created_at_str := file_info.get("created_at"):
             date_created = str(parser.parse(created_at_str).timestamp())
 
-        file_size = file_data.get("size") if "size" in file_data else None
+        file_size = file_info.get("size") if "size" in file_info else None
 
-        version = file_data.get("id")
+        version = file_info.get("id")
         record_locator = {
             "protocol": self.index_config.protocol,
             "remote_file_path": self.index_config.remote_url,
-            "file_id": file_data.get("id"),
+            "file_id": file_info.get("id"),
         }
         return FileDataSourceMetadata(
             date_created=date_created,

unstructured_ingest/v2/processes/connectors/fsspec/dropbox.py

@@ -93,15 +93,15 @@ class DropboxIndexer(FsspecIndexer):
     index_config: DropboxIndexerConfig
     connector_type: str = CONNECTOR_TYPE
 
-    def get_path(self, file_data: dict) -> str:
-        return file_data["name"]
+    def get_path(self, file_info: dict) -> str:
+        return file_info["name"]
 
-    def get_metadata(self, file_data: dict) -> FileDataSourceMetadata:
-        path = file_data["name"].lstrip("/")
+    def get_metadata(self, file_info: dict) -> FileDataSourceMetadata:
+        path = file_info["name"].lstrip("/")
         date_created = None
         date_modified = None
-        server_modified = file_data.get("server_modified")
-        client_modified = file_data.get("client_modified")
+        server_modified = file_info.get("server_modified")
+        client_modified = file_info.get("client_modified")
         if server_modified and client_modified and server_modified > client_modified:
             date_created = str(client_modified.timestamp())
             date_modified = str(server_modified.timestamp())
@@ -109,13 +109,13 @@ class DropboxIndexer(FsspecIndexer):
             date_created = str(server_modified.timestamp())
             date_modified = str(client_modified.timestamp())
 
-        file_size = file_data.get("size") if "size" in file_data else None
+        file_size = file_info.get("size") if "size" in file_info else None
 
-        version = file_data.get("content_hash")
+        version = file_info.get("content_hash")
         record_locator = {
             "protocol": self.index_config.protocol,
             "remote_file_path": self.index_config.remote_url,
-            "file_id": file_data.get("id"),
+            "file_id": file_info.get("id"),
         }
         return FileDataSourceMetadata(
             date_created=date_created,

unstructured_ingest/v2/processes/connectors/fsspec/fsspec.py

@@ -119,7 +119,7 @@ class FsspecIndexer(Indexer):
             logger.error(f"failed to validate connection: {e}", exc_info=True)
             raise self.wrap_error(e=e)
 
-    def get_file_data(self) -> list[dict[str, Any]]:
+    def get_file_info(self) -> list[dict[str, Any]]:
         if not self.index_config.recursive:
             # fs.ls does not walk directories
             # directories that are listed in cloud storage can cause problems
@@ -156,24 +156,56 @@ class FsspecIndexer(Indexer):
 
         return random.sample(files, n)
 
-    def get_metadata(self, file_data: dict) -> FileDataSourceMetadata:
+    def get_metadata(self, file_info: dict) -> FileDataSourceMetadata:
         raise NotImplementedError()
 
-    def get_path(self, file_data: dict) -> str:
-        return file_data["name"]
+    def get_path(self, file_info: dict) -> str:
+        return file_info["name"]
 
     def sterilize_info(self, file_data: dict) -> dict:
         return sterilize_dict(data=file_data)
 
+    def create_init_file_data(self, remote_filepath: Optional[str] = None) -> FileData:
+        # Create initial file data that requires no network calls and is constructed purely
+        # with information that exists in the config
+        remote_filepath = remote_filepath or self.index_config.remote_url
+        path_without_protocol = remote_filepath.split("://")[1]
+        rel_path = remote_filepath.replace(path_without_protocol, "").lstrip("/")
+        return FileData(
+            identifier=str(uuid5(NAMESPACE_DNS, remote_filepath)),
+            connector_type=self.connector_type,
+            display_name=remote_filepath,
+            source_identifiers=SourceIdentifiers(
+                filename=Path(remote_filepath).name,
+                rel_path=rel_path or None,
+                fullpath=remote_filepath,
+            ),
+            metadata=FileDataSourceMetadata(url=remote_filepath),
+        )
+
+    def hydrate_file_data(self, init_file_data: FileData):
+        # Get file info
+        with self.connection_config.get_client(protocol=self.index_config.protocol) as client:
+            files = client.ls(self.index_config.path_without_protocol, detail=True)
+        filtered_files = [
+            file for file in files if file.get("size") > 0 and file.get("type") == "file"
+        ]
+        if not filtered_files:
+            raise ValueError(f"{init_file_data} did not reference any valid file")
+        if len(filtered_files) > 1:
+            raise ValueError(f"{init_file_data} referenced more than one file")
+        file_info = filtered_files[0]
+        init_file_data.additional_metadata = self.get_metadata(file_info=file_info)
+
     def run(self, **kwargs: Any) -> Generator[FileData, None, None]:
-        files = self.get_file_data()
-        for file_data in files:
-            file_path = self.get_path(file_data=file_data)
+        files = self.get_file_info()
+        for file_info in files:
+            file_path = self.get_path(file_info=file_info)
             # Note: we remove any remaining leading slashes (Box introduces these)
             # to get a valid relative path
             rel_path = file_path.replace(self.index_config.path_without_protocol, "").lstrip("/")
 
-            additional_metadata = self.sterilize_info(file_data=file_data)
+            additional_metadata = self.sterilize_info(file_data=file_info)
             additional_metadata["original_file_path"] = file_path
             yield FileData(
                 identifier=str(uuid5(NAMESPACE_DNS, file_path)),
@@ -183,7 +215,7 @@ class FsspecIndexer(Indexer):
                     rel_path=rel_path or None,
                     fullpath=file_path,
                 ),
-                metadata=self.get_metadata(file_data=file_data),
+                metadata=self.get_metadata(file_info=file_info),
                 additional_metadata=additional_metadata,
                 display_name=file_path,
             )

unstructured_ingest/v2/processes/connectors/fsspec/gcs.py

@@ -131,22 +131,22 @@ class GcsIndexer(FsspecIndexer):
     index_config: GcsIndexerConfig
     connector_type: str = CONNECTOR_TYPE
 
-    def get_metadata(self, file_data: dict) -> FileDataSourceMetadata:
-        path = file_data["name"]
+    def get_metadata(self, file_info: dict) -> FileDataSourceMetadata:
+        path = file_info["name"]
         date_created = None
         date_modified = None
-        if modified_at_str := file_data.get("updated"):
+        if modified_at_str := file_info.get("updated"):
             date_modified = str(parser.parse(modified_at_str).timestamp())
-        if created_at_str := file_data.get("timeCreated"):
+        if created_at_str := file_info.get("timeCreated"):
             date_created = str(parser.parse(created_at_str).timestamp())
 
-        file_size = file_data.get("size") if "size" in file_data else None
+        file_size = file_info.get("size") if "size" in file_info else None
 
-        version = file_data.get("etag")
+        version = file_info.get("etag")
         record_locator = {
             "protocol": self.index_config.protocol,
             "remote_file_path": self.index_config.remote_url,
-            "file_id": file_data.get("id"),
+            "file_id": file_info.get("id"),
         }
         return FileDataSourceMetadata(
             date_created=date_created,

unstructured_ingest/v2/processes/connectors/fsspec/s3.py

@@ -110,22 +110,22 @@ class S3Indexer(FsspecIndexer):
     def wrap_error(self, e: Exception) -> Exception:
         return self.connection_config.wrap_error(e=e)
 
-    def get_path(self, file_data: dict) -> str:
-        return file_data["Key"]
+    def get_path(self, file_info: dict) -> str:
+        return file_info["Key"]
 
-    def get_metadata(self, file_data: dict) -> FileDataSourceMetadata:
-        path = file_data["Key"]
+    def get_metadata(self, file_info: dict) -> FileDataSourceMetadata:
+        path = file_info["Key"]
         date_created = None
         date_modified = None
-        modified = file_data.get("LastModified")
+        modified = file_info.get("LastModified")
         if modified:
             date_created = str(modified.timestamp())
             date_modified = str(modified.timestamp())
 
-        file_size = file_data.get("size") if "size" in file_data else None
-        file_size = file_size or file_data.get("Size")
+        file_size = file_info.get("size") if "size" in file_info else None
+        file_size = file_size or file_info.get("Size")
 
-        version = file_data.get("ETag").rstrip('"').lstrip('"') if "ETag" in file_data else None
+        version = file_info.get("ETag").rstrip('"').lstrip('"') if "ETag" in file_info else None
         metadata: dict[str, str] = {}
         with contextlib.suppress(AttributeError):
             with self.connection_config.get_client(protocol=self.index_config.protocol) as client:

unstructured_ingest/v2/processes/connectors/fsspec/sftp.py

@@ -107,12 +107,12 @@ class SftpIndexer(FsspecIndexer):
             file.identifier = new_identifier
             yield file
 
-    def get_metadata(self, file_data: dict) -> FileDataSourceMetadata:
-        path = file_data["name"]
-        date_created = str(file_data.get("time").timestamp()) if "time" in file_data else None
-        date_modified = str(file_data.get("mtime").timestamp()) if "mtime" in file_data else None
+    def get_metadata(self, file_info: dict) -> FileDataSourceMetadata:
+        path = file_info["name"]
+        date_created = str(file_info.get("time").timestamp()) if "time" in file_info else None
+        date_modified = str(file_info.get("mtime").timestamp()) if "mtime" in file_info else None
 
-        file_size = file_data.get("size") if "size" in file_data else None
+        file_size = file_info.get("size") if "size" in file_info else None
 
         record_locator = {
             "protocol": self.index_config.protocol,

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

@@ -15,11 +15,14 @@ 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, sqlite_source_entry
+from .vastdb import CONNECTOR_TYPE as VASTDB_CONNECTOR_TYPE
+from .vastdb import vastdb_destination_entry, vastdb_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_source_entry(source_type=SINGLESTORE_CONNECTOR_TYPE, entry=singlestore_source_entry)
+add_source_entry(source_type=VASTDB_CONNECTOR_TYPE, entry=vastdb_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)
@@ -31,3 +34,4 @@ add_destination_entry(
     destination_type=DATABRICKS_DELTA_TABLES_CONNECTOR_TYPE,
     entry=databricks_delta_tables_destination_entry,
 )
+add_destination_entry(destination_type=VASTDB_CONNECTOR_TYPE, entry=vastdb_destination_entry)

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

@@ -3,6 +3,7 @@ from contextlib import contextmanager
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any, Generator, Optional
 
+import pandas as pd
 from pydantic import Field, Secret
 
 from unstructured_ingest.v2.logger import logger
@@ -139,7 +140,7 @@ class SingleStoreUploader(SQLUploader):
                 if isinstance(value, (list, dict)):
                     value = json.dumps(value)
                 if column_name in _DATE_COLUMNS:
-                    if value is None:
+                    if value is None or pd.isna(value):
                         parsed.append(None)
                     else:
                         parsed.append(parse_date_string(value))

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

@@ -170,7 +170,7 @@ class SnowflakeUploader(SQLUploader):
                 f"{self.upload_config.record_id_key}, skipping delete"
             )
         df.replace({np.nan: None}, inplace=True)
-        self._fit_to_schema(df=df, columns=self.get_table_columns())
+        self._fit_to_schema(df=df)
 
         columns = list(df.columns)
         stmt = "INSERT INTO {table_name} ({columns}) VALUES({values})".format(

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

@@ -14,7 +14,7 @@ from dateutil import parser
 from pydantic import BaseModel, Field, Secret
 
 from unstructured_ingest.error import DestinationConnectionError, SourceConnectionError
-from unstructured_ingest.utils.data_prep import get_data, get_data_df, split_dataframe
+from unstructured_ingest.utils.data_prep import get_data, get_data_df, split_dataframe, write_data
 from unstructured_ingest.v2.constants import RECORD_ID_LABEL
 from unstructured_ingest.v2.interfaces import (
     AccessConfig,
@@ -310,9 +310,11 @@ class SQLUploadStager(UploadStager):
         )
         df = self.conform_dataframe(df=df)
 
+        output_filename_suffix = Path(elements_filepath).suffix
+        output_filename = f"{Path(output_filename).stem}{output_filename_suffix}"
         output_path = self.get_output_path(output_filename=output_filename, output_dir=output_dir)
 
-        self.write_output(output_path=output_path, data=df.to_dict(orient="records"))
+        write_data(path=output_path, data=df.to_dict(orient="records"))
         return output_path
 
 
@@ -330,6 +332,7 @@ class SQLUploader(Uploader):
     upload_config: SQLUploaderConfig
     connection_config: SQLConnectionConfig
     values_delimiter: str = "?"
+    _columns: list[str] = field(init=False, default=None)
 
     def precheck(self) -> None:
         try:
@@ -352,7 +355,7 @@ class SQLUploader(Uploader):
             parsed = []
             for column_name, value in zip(columns, row):
                 if column_name in _DATE_COLUMNS:
-                    if value is None:
+                    if value is None or pd.isna(value):  # pandas is nan
                         parsed.append(None)
                     else:
                         parsed.append(parse_date_string(value))
@@ -362,8 +365,9 @@ class SQLUploader(Uploader):
         return output
 
     def _fit_to_schema(self, df: pd.DataFrame) -> pd.DataFrame:
+        table_columns = self.get_table_columns()
         columns = set(df.columns)
-        schema_fields = set(columns)
+        schema_fields = set(table_columns)
         columns_to_drop = columns - schema_fields
         missing_columns = schema_fields - columns
 
@@ -393,8 +397,8 @@ class SQLUploader(Uploader):
                 f"record id column "
                 f"{self.upload_config.record_id_key}, skipping delete"
             )
+        df = self._fit_to_schema(df=df)
         df.replace({np.nan: None}, inplace=True)
-        self._fit_to_schema(df=df)
 
         columns = list(df.columns)
         stmt = "INSERT INTO {table_name} ({columns}) VALUES({values})".format(
@@ -422,9 +426,11 @@ class SQLUploader(Uploader):
                 cursor.executemany(stmt, values)
 
     def get_table_columns(self) -> list[str]:
-        with self.get_cursor() as cursor:
-            cursor.execute(f"SELECT * from {self.upload_config.table_name}")
-            return [desc[0] for desc in cursor.description]
+        if self._columns is None:
+            with self.get_cursor() as cursor:
+                cursor.execute(f"SELECT * from {self.upload_config.table_name} LIMIT 1")
+                self._columns = [desc[0] for desc in cursor.description]
+        return self._columns
 
     def can_delete(self) -> bool:
         return self.upload_config.record_id_key in self.get_table_columns()