cosmotech-acceleration-library 2.0.0__py3-none-any.whl → 2.1.0rc1__py3-none-any.whl

cosmotech/coal/__init__.py

@@ -5,4 +5,4 @@
 # etc., to any person is prohibited unless it has been previously and
 # specifically authorized by written means by Cosmo Tech.
 
-__version__ = "2.0.0"
+__version__ = "2.1.0-rc1"

cosmotech/coal/azure/__init__.py

@@ -11,13 +11,13 @@ Azure services integration module.
 This module provides functions for interacting with Azure services like Storage and ADX.
 """
 
+# Re-export blob functions for easier importing
+from cosmotech.coal.azure.blob import (
+    dump_store_to_azure,
+)
+
 # Re-export storage functions for easier importing
 from cosmotech.coal.azure.storage import (
     upload_file,
     upload_folder,
 )
-
-# Re-export blob functions for easier importing
-from cosmotech.coal.azure.blob import (
-    dump_store_to_azure,
-)

cosmotech/coal/azure/adx/__init__.py

@@ -5,22 +5,36 @@
 # etc., to any person is prohibited unless it has been previously and
 # specifically authorized by written means by Cosmo Tech.
 
-from cosmotech.coal.azure.adx.auth import create_kusto_client, create_ingest_client, initialize_clients
-from cosmotech.coal.azure.adx.query import run_query, run_command_query
+from cosmotech.coal.azure.adx.auth import (
+    create_ingest_client,
+    create_kusto_client,
+    initialize_clients,
+)
 from cosmotech.coal.azure.adx.ingestion import (
-    ingest_dataframe,
-    send_to_adx,
+    IngestionStatus,
     check_ingestion_status,
-    monitor_ingestion,
     handle_failures,
-    IngestionStatus,
+    ingest_dataframe,
+    monitor_ingestion,
+    send_to_adx,
 )
-from cosmotech.coal.azure.adx.tables import table_exists, create_table, check_and_create_table, _drop_by_tag
-from cosmotech.coal.azure.adx.utils import type_mapping, create_column_mapping
-from cosmotech.coal.azure.adx.store import send_pyarrow_table_to_adx, send_table_data, process_tables, send_store_to_adx
+from cosmotech.coal.azure.adx.query import run_command_query, run_query
 from cosmotech.coal.azure.adx.runner import (
-    prepare_csv_content,
     construct_create_query,
     insert_csv_files,
+    prepare_csv_content,
     send_runner_data,
 )
+from cosmotech.coal.azure.adx.store import (
+    process_tables,
+    send_pyarrow_table_to_adx,
+    send_store_to_adx,
+    send_table_data,
+)
+from cosmotech.coal.azure.adx.tables import (
+    _drop_by_tag,
+    check_and_create_table,
+    create_table,
+    table_exists,
+)
+from cosmotech.coal.azure.adx.utils import create_column_mapping, type_mapping

cosmotech/coal/azure/adx/ingestion.py

@@ -5,28 +5,24 @@
 # etc., to any person is prohibited unless it has been previously and
 # specifically authorized by written means by Cosmo Tech.
 
+import os
+import time
 from enum import Enum
-from typing import Dict
-from typing import Iterator
-from typing import List
-from typing import Optional
-from typing import Tuple
+from typing import Dict, Iterator, List, Optional, Tuple
 
-import os
 import pandas as pd
-import time
 import tqdm
 from azure.kusto.data import KustoClient
 from azure.kusto.data.data_format import DataFormat
-from azure.kusto.ingest import IngestionProperties
-from azure.kusto.ingest import QueuedIngestClient
-from azure.kusto.ingest import ReportLevel
-from azure.kusto.ingest.status import FailureMessage
-from azure.kusto.ingest.status import KustoIngestStatusQueues
-from azure.kusto.ingest.status import SuccessMessage
+from azure.kusto.ingest import IngestionProperties, QueuedIngestClient, ReportLevel
+from azure.kusto.ingest.status import (
+    FailureMessage,
+    KustoIngestStatusQueues,
+    SuccessMessage,
+)
 from cosmotech.orchestrator.utils.translate import T
 
-from cosmotech.coal.azure.adx.tables import create_table, _drop_by_tag
+from cosmotech.coal.azure.adx.tables import _drop_by_tag, create_table
 from cosmotech.coal.azure.adx.utils import type_mapping
 from cosmotech.coal.utils.logger import LOGGER
 

cosmotech/coal/azure/adx/query.py

@@ -7,9 +7,9 @@
 
 from azure.kusto.data import KustoClient
 from azure.kusto.data.response import KustoResponseDataSet
+from cosmotech.orchestrator.utils.translate import T
 
 from cosmotech.coal.utils.logger import LOGGER
-from cosmotech.orchestrator.utils.translate import T
 
 
 def run_query(client: KustoClient, database: str, query: str) -> KustoResponseDataSet:

cosmotech/coal/azure/adx/utils.py

@@ -5,13 +5,13 @@
 # etc., to any person is prohibited unless it has been previously and
 # specifically authorized by written means by Cosmo Tech.
 
-import dateutil.parser
 from typing import Any, Dict
 
+import dateutil.parser
 import pyarrow
+from cosmotech.orchestrator.utils.translate import T
 
 from cosmotech.coal.utils.logger import LOGGER
-from cosmotech.orchestrator.utils.translate import T
 
 
 def create_column_mapping(data: pyarrow.Table) -> Dict[str, str]:

cosmotech/coal/azure/blob.py

@@ -21,6 +21,7 @@ from azure.storage.blob import BlobServiceClient
 from cosmotech.orchestrator.utils.translate import T
 
 from cosmotech.coal.store.store import Store
+from cosmotech.coal.utils.configuration import Configuration
 from cosmotech.coal.utils.logger import LOGGER
 
 VALID_TYPES = (
@@ -31,42 +32,35 @@ VALID_TYPES = (
 
 
 def dump_store_to_azure(
-    store_folder: str,
-    account_name: str,
-    container_name: str,
-    tenant_id: str,
-    client_id: str,
-    client_secret: str,
-    output_type: str = "sqlite",
-    file_prefix: str = "",
+    configuration: Configuration = Configuration(),
     selected_tables: list[str] = [],
 ) -> None:
     """
     Dump Store data to Azure Blob Storage.
 
     Args:
-        store_folder: Folder containing the Store
-        account_name: Azure Storage account name
-        container_name: Azure Storage container name
-        tenant_id: Azure tenant ID
-        client_id: Azure client ID
-        client_secret: Azure client secret
-        output_type: Output file type (sqlite, csv, or parquet)
-        file_prefix: Prefix for uploaded files
+        configuration: Configuration utils class
+        selected_tables: List of tables name
 
     Raises:
         ValueError: If the output type is invalid
     """
-    _s = Store(store_location=store_folder)
+    _s = Store(configuration=configuration)
+    output_type = configuration.safe_get("azure.output_type", default="sqlite")
+    file_prefix = configuration.safe_get("azure.file_prefix", default="")
 
     if output_type not in VALID_TYPES:
         LOGGER.error(T("coal.common.validation.invalid_output_type").format(output_type=output_type))
         raise ValueError(T("coal.common.validation.invalid_output_type").format(output_type=output_type))
 
     container_client = BlobServiceClient(
-        account_url=f"https://{account_name}.blob.core.windows.net/",
-        credential=ClientSecretCredential(tenant_id=tenant_id, client_id=client_id, client_secret=client_secret),
-    ).get_container_client(container_name)
+        account_url=f"https://{configuration.azure.account_name}.blob.core.windows.net/",
+        credential=ClientSecretCredential(
+            tenant_id=configuration.azure.tenant_id,
+            client_id=configuration.azure.client_id,
+            client_secret=configuration.azure.client_secret,
+        ),
+    ).get_container_client(configuration.azure.container_name)
 
     def data_upload(data_stream: BytesIO, file_name: str):
         uploaded_file_name = file_prefix + file_name

cosmotech/coal/cosmotech_api/apis/dataset.py

@@ -35,30 +35,49 @@ class DatasetApi(BaseDatasetApi, Connection):
         LOGGER.debug(T("coal.cosmotech_api.initialization.dataset_api_initialized"))
 
     def download_dataset(self, dataset_id) -> Dataset:
+        LOGGER.debug(f"Downloading dataset {dataset_id}")
         dataset = self.get_dataset(
             organization_id=self.configuration.cosmotech.organization_id,
             workspace_id=self.configuration.cosmotech.workspace_id,
             dataset_id=dataset_id,
         )
+        # send dataset files under dataset id folder
+        destination = Path(self.configuration.cosmotech.dataset_absolute_path) / dataset_id
+        for part in dataset.parts:
+            self._download_part(dataset_id, part, destination)
+        return dataset
 
-        dataset_dir = self.configuration.cosmotech.dataset_absolute_path
-        dataset_dir_path = Path(dataset_dir) / dataset_id
+    def download_parameter(self, dataset_id) -> Dataset:
+        LOGGER.debug(f"Downloading dataset {dataset_id}")
+        dataset = self.get_dataset(
+            organization_id=self.configuration.cosmotech.organization_id,
+            workspace_id=self.configuration.cosmotech.workspace_id,
+            dataset_id=dataset_id,
+        )
+        # send parameters file under parameters_name folder
+        destination = Path(self.configuration.cosmotech.parameters_absolute_path) / dataset_id
         for part in dataset.parts:
-            part_file_path = dataset_dir_path / part.source_name
-            part_file_path.parent.mkdir(parents=True, exist_ok=True)
-            data_part = self.download_dataset_part(
-                organization_id=self.configuration.cosmotech.organization_id,
-                workspace_id=self.configuration.cosmotech.workspace_id,
-                dataset_id=dataset_id,
-                dataset_part_id=part.id,
-            )
-            with open(part_file_path, "wb") as binary_file:
-                binary_file.write(data_part)
-            LOGGER.debug(
-                T("coal.services.dataset.part_downloaded").format(part_name=part.source_name, file_path=part_file_path)
-            )
+            part_dst = destination / part.name
+            self._download_part(dataset_id, part, part_dst)
         return dataset
 
+    def _download_part(self, dataset_id, dataset_part, destination):
+        part_file_path = destination / dataset_part.source_name
+        part_file_path.parent.mkdir(parents=True, exist_ok=True)
+        data_part = self.download_dataset_part(
+            organization_id=self.configuration.cosmotech.organization_id,
+            workspace_id=self.configuration.cosmotech.workspace_id,
+            dataset_id=dataset_id,
+            dataset_part_id=dataset_part.id,
+        )
+        with open(part_file_path, "wb") as binary_file:
+            binary_file.write(data_part)
+        LOGGER.debug(
+            T("coal.services.dataset.part_downloaded").format(
+                part_name=dataset_part.source_name, file_path=part_file_path
+            )
+        )
+
     @staticmethod
     def path_to_parts(_path, part_type) -> list[tuple[str, Path, DatasetPartTypeEnum]]:
         if (_path := Path(_path)).is_dir():
@@ -70,7 +89,21 @@ class DatasetApi(BaseDatasetApi, Connection):
         dataset_name: str,
         as_files: Optional[list[Union[Path, str]]] = (),
         as_db: Optional[list[Union[Path, str]]] = (),
+        tags: Optional[list[str]] = None,
+        additional_data: Optional[dict] = None,
     ) -> Dataset:
+        """Upload a new dataset with optional tags and additional data.
+
+        Args:
+            dataset_name: The name of the dataset to create
+            as_files: List of file paths to upload as FILE type parts
+            as_db: List of file paths to upload as DB type parts
+            tags: Optional list of tags to associate with the dataset
+            additional_data: Optional dictionary of additional metadata
+
+        Returns:
+            The created Dataset object
+        """
         _parts = list()
 
         for _f in as_files:
@@ -81,6 +114,8 @@ class DatasetApi(BaseDatasetApi, Connection):
 
         d_request = DatasetCreateRequest(
             name=dataset_name,
+            tags=tags,
+            additional_data=additional_data,
             parts=list(
                 DatasetPartCreateRequest(
                     name=_p_name,
@@ -92,12 +127,96 @@ class DatasetApi(BaseDatasetApi, Connection):
             ),
         )
 
+        _files = []
+        for _p in _parts:
+            with _p[1].open("rb") as _p_file:
+                _files.append((_p[0], _p_file.read()))
+
         d_ret = self.create_dataset(
             self.configuration.cosmotech.organization_id,
             self.configuration.cosmotech.workspace_id,
             d_request,
-            files=list((_p[0], _p[1].open("rb").read()) for _p in _parts),
+            files=_files,
         )
 
         LOGGER.info(T("coal.services.dataset.dataset_created").format(dataset_id=d_ret.id))
         return d_ret
+
+    def upload_dataset_parts(
+        self,
+        dataset_id: str,
+        as_files: Optional[list[Union[Path, str]]] = (),
+        as_db: Optional[list[Union[Path, str]]] = (),
+        replace_existing: bool = False,
+    ) -> Dataset:
+        """Upload parts to an existing dataset.
+
+        Args:
+            dataset_id: The ID of the existing dataset
+            as_files: List of file paths to upload as FILE type parts
+            as_db: List of file paths to upload as DB type parts
+            replace_existing: If True, replace existing parts with same name
+
+        Returns:
+            The updated Dataset object
+        """
+        # Get current dataset to check existing parts
+        current_dataset = self.get_dataset(
+            organization_id=self.configuration.cosmotech.organization_id,
+            workspace_id=self.configuration.cosmotech.workspace_id,
+            dataset_id=dataset_id,
+        )
+
+        # Build set of existing part names and their IDs for quick lookup
+        existing_parts = {part.source_name: part.id for part in (current_dataset.parts or [])}
+
+        # Collect parts to upload
+        _parts = list()
+        for _f in as_files:
+            _parts.extend(self.path_to_parts(_f, DatasetPartTypeEnum.FILE))
+        for _db in as_db:
+            _parts.extend(self.path_to_parts(_db, DatasetPartTypeEnum.DB))
+
+        # Process each part
+        for _p_name, _p_path, _type in _parts:
+            if _p_name in existing_parts:
+                if replace_existing:
+                    # Delete existing part before creating new one
+                    self.delete_dataset_part(
+                        organization_id=self.configuration.cosmotech.organization_id,
+                        workspace_id=self.configuration.cosmotech.workspace_id,
+                        dataset_id=dataset_id,
+                        dataset_part_id=existing_parts[_p_name],
+                    )
+                    LOGGER.info(T("coal.services.dataset.part_replaced").format(part_name=_p_name))
+                else:
+                    LOGGER.warning(T("coal.services.dataset.part_skipped").format(part_name=_p_name))
+                    continue
+
+            # Create new part
+            part_request = DatasetPartCreateRequest(
+                name=_p_name,
+                description=_p_name,
+                sourceName=_p_name,
+                type=_type,
+            )
+
+            with _p_path.open("rb") as _p_file:
+                self.create_dataset_part(
+                    organization_id=self.configuration.cosmotech.organization_id,
+                    workspace_id=self.configuration.cosmotech.workspace_id,
+                    dataset_id=dataset_id,
+                    dataset_part_create_request=part_request,
+                    file=(_p_name, _p_file.read()),
+                )
+                LOGGER.debug(T("coal.services.dataset.part_uploaded").format(part_name=_p_name))
+
+        # Return updated dataset
+        updated_dataset = self.get_dataset(
+            organization_id=self.configuration.cosmotech.organization_id,
+            workspace_id=self.configuration.cosmotech.workspace_id,
+            dataset_id=dataset_id,
+        )
+
+        LOGGER.info(T("coal.services.dataset.parts_uploaded").format(dataset_id=dataset_id))
+        return updated_dataset

cosmotech/coal/cosmotech_api/apis/runner.py

@@ -29,43 +29,47 @@ class RunnerApi(BaseRunnerApi, Connection):
 
     def get_runner_metadata(
         self,
-        organization_id: str,
-        workspace_id: str,
-        runner_id: str,
+        runner_id: Optional[str] = None,
         include: Optional[list[str]] = None,
         exclude: Optional[list[str]] = None,
     ) -> dict[str, Any]:
-        runner = self.get_runner(organization_id, workspace_id, runner_id)
+        runner = self.get_runner(
+            self.configuration.cosmotech.organization_id,
+            self.configuration.cosmotech.workspace_id,
+            runner_id or self.configuration.cosmotech.runner_id,
+        )
 
         return runner.model_dump(by_alias=True, exclude_none=True, include=include, exclude=exclude, mode="json")
 
     def download_runner_data(
         self,
-        organization_id: str,
-        workspace_id: str,
-        runner_id: str,
-        parameter_folder: str,
-        dataset_folder: Optional[str] = None,
+        download_datasets: Optional[str] = None,
     ):
         LOGGER.info(T("coal.cosmotech_api.runner.starting_download"))
 
         # Get runner data
-        runner_data = self.get_runner(organization_id, workspace_id, runner_id)
+        runner = self.get_runner(
+            self.configuration.cosmotech.organization_id,
+            self.configuration.cosmotech.workspace_id,
+            self.configuration.cosmotech.runner_id,
+        )
 
         # Skip if no parameters found
-        if not runner_data.parameters_values:
+        if not runner.parameters_values:
             LOGGER.warning(T("coal.cosmotech_api.runner.no_parameters"))
         else:
             LOGGER.info(T("coal.cosmotech_api.runner.loaded_data"))
-            parameters = Parameters(runner_data)
-            parameters.write_parameters_to_json(parameter_folder)
+            parameters = Parameters(runner)
+            parameters.write_parameters_to_json(self.configuration.cosmotech.parameters_absolute_path)
 
-        # Download datasets if requested
-        if dataset_folder:
-            datasets_ids = runner_data.datasets.bases
+        if runner.datasets.parameter:
+            ds_api = DatasetApi(self.configuration)
+            ds_api.download_parameter(runner.datasets.parameter)
 
-            if datasets_ids:
-                LOGGER.info(T("coal.cosmotech_api.runner.downloading_datasets").format(count=len(datasets_ids)))
+        # Download datasets if requested
+        if download_datasets:
+            LOGGER.info(T("coal.cosmotech_api.runner.downloading_datasets").format(count=len(runner.datasets.bases)))
+            if runner.datasets.bases:
                 ds_api = DatasetApi(self.configuration)
-                for dataset_id in datasets_ids:
+                for dataset_id in runner.datasets.bases:
                     ds_api.download_dataset(dataset_id)

cosmotech/coal/postgresql/runner.py

@@ -38,8 +38,6 @@ def send_runner_metadata_to_postgresql(
     # Get runner metadata
     _runner_api = RunnerApi(configuration)
     runner = _runner_api.get_runner_metadata(
-        configuration.cosmotech.organization_id,
-        configuration.cosmotech.workspace_id,
         configuration.cosmotech.runner_id,
     )
 
@@ -51,7 +49,7 @@ def send_runner_metadata_to_postgresql(
                 CREATE TABLE IF NOT EXISTS {schema_table}  (
                   id varchar(32) PRIMARY KEY,
                   name varchar(256),
-                  last_run_id varchar(32),
+                  last_csm_run_id varchar(32),
                   run_template_id varchar(32)
                 );
             """
@@ -60,24 +58,25 @@ def send_runner_metadata_to_postgresql(
             conn.commit()
             LOGGER.info(T("coal.services.postgresql.metadata"))
             sql_upsert = f"""
-                INSERT INTO {schema_table} (id, name, last_run_id, run_template_id)
-                  VALUES(%s, %s, %s, %s)
+                INSERT INTO {schema_table} (id, name, last_csm_run_id, run_template_id)
+                  VALUES ($1, $2, $3, $4)
                   ON CONFLICT (id)
                   DO
-                    UPDATE SET name = EXCLUDED.name, last_run_id = EXCLUDED.last_run_id;
+                    UPDATE SET name = EXCLUDED.name, last_csm_run_id = EXCLUDED.last_csm_run_id;
             """
+            LOGGER.debug(runner)
             curs.execute(
                 sql_upsert,
                 (
                     runner.get("id"),
                     runner.get("name"),
-                    runner.get("lastRunId"),
+                    runner.get("lastRunInfo").get("lastRunId"),
                     runner.get("runTemplateId"),
                 ),
             )
             conn.commit()
             LOGGER.info(T("coal.services.postgresql.metadata_updated"))
-    return runner.get("lastRunId")
+    return runner.get("lastRunInfo").get("lastRunId")
 
 
 def remove_runner_metadata_from_postgresql(
@@ -97,8 +96,6 @@ def remove_runner_metadata_from_postgresql(
     # Get runner metadata
     _runner_api = RunnerApi(configuration)
     runner = _runner_api.get_runner_metadata(
-        configuration.cosmotech.organization_id,
-        configuration.cosmotech.workspace_id,
         configuration.cosmotech.runner_id,
     )
 
@@ -108,7 +105,7 @@ def remove_runner_metadata_from_postgresql(
             schema_table = f"{_psql.db_schema}.{_psql.table_prefix}RunnerMetadata"
             sql_delete_from_metatable = f"""
                 DELETE FROM {schema_table}
-                WHERE last_run_id={runner.get("lastRunId")};
+                WHERE last_csm_run_id={runner.get("lastRunId")};
             """
             curs.execute(sql_delete_from_metatable)
             conn.commit()

cosmotech/coal/postgresql/store.py

@@ -53,24 +53,27 @@ def dump_store_to_postgresql(
         selected_tables: list of tables to send
         fk_id: foreign key id to add to all table on all rows
     """
-    _c = Configuration()
-    _c.postgres.host = postgres_host
-    _c.postgres.port = postgres_port
-    _c.postgres.db_name = postgres_db
-    _c.postgres.db_schema = postgres_schema
-    _c.postgres.user_name = postgres_user
-    _c.postgres.user_password = postgres_password
-    _c.postgres.password_encoding = force_encode
-    _c.postgres.table_prefix = table_prefix
-
-    dump_store_to_postgresql_from_conf(
-        configuration=_c, store_folder=store_folder, replace=replace, selected_tables=selected_tables, fk_id=fk_id
+    _c = Configuration(
+        {
+            "coal": {"store": store_folder},
+            "postgres": {
+                "host": postgres_host,
+                "port": postgres_port,
+                "db_name": postgres_db,
+                "db_schema": postgres_schema,
+                "user_name": postgres_user,
+                "user_password": postgres_password,
+                "password_encoding": force_encode,
+                "table_prefix": table_prefix,
+            },
+        }
     )
 
+    dump_store_to_postgresql_from_conf(configuration=_c, replace=replace, selected_tables=selected_tables, fk_id=fk_id)
+
 
 def dump_store_to_postgresql_from_conf(
     configuration: Configuration,
-    store_folder: str,
     replace: bool = True,
     selected_tables: list[str] = [],
     fk_id: str = None,
@@ -80,14 +83,12 @@ def dump_store_to_postgresql_from_conf(
 
     Args:
         configuration: coal Configuration
-        store_folder: Folder containing the Store
         replace: Whether to replace existing tables
         selected_tables: list of tables to send
         fk_id: foreign key id to add to all table on all rows
     """
     _psql = PostgresUtils(configuration)
-    print(_psql.send_pyarrow_table_to_postgresql)
-    _s = Store(store_location=store_folder)
+    _s = Store(configuration=configuration)
 
     tables = list(_s.list_tables())
     if selected_tables:
@@ -100,18 +101,12 @@ def dump_store_to_postgresql_from_conf(
             _s_time = perf_counter()
             target_table_name = f"{_psql.table_prefix}{table_name}"
             LOGGER.info(T("coal.services.database.table_entry").format(table=target_table_name))
-            if fk_id:
-                _s.execute_query(
-                    f"""
-                    ALTER TABLE {_psql.table_prefix}{table_name}
-                    ADD run_id TEXT NOT NULL
-                    DEFAULT ({fk_id})
-                    """
-                )
             data = _s.get_table(table_name)
             if not len(data):
                 LOGGER.info(T("coal.services.database.no_rows"))
                 continue
+            if fk_id:
+                data = data.append_column("csm_run_id", [[fk_id] * data.num_rows])
             _dl_time = perf_counter()
             rows = _psql.send_pyarrow_table_to_postgresql(
                 data,
@@ -120,7 +115,7 @@ def dump_store_to_postgresql_from_conf(
             )
             if fk_id and _psql.is_metadata_exists():
                 metadata_table = f"{_psql.table_prefix}RunnerMetadata"
-                _psql.add_fk_constraint(table_name, "run_id", metadata_table, "last_run_id")
+                _psql.add_fk_constraint(table_name, "csm_run_id", metadata_table, "last_csm_run_id")
 
             total_rows += rows
             _up_time = perf_counter()

cosmotech/coal/postgresql/utils.py

@@ -75,7 +75,8 @@ class PostgresUtils:
             f"/{self.db_name}"
         )
 
-    def metadata_table_name(self):
+    @property
+    def metadata_table_name(self) -> str:
         return f"{self.table_prefix}RunnerMetadata"
 
     def get_postgresql_table_schema(self, target_table_name: str) -> Optional[pa.Schema]:

cosmotech/coal/singlestore/store.py

@@ -12,15 +12,16 @@ This module provides functions for interacting with SingleStore databases
 for store operations.
 """
 
+import csv
 import pathlib
 import time
-import csv
+
 import singlestoredb as s2
+from cosmotech.orchestrator.utils.translate import T
 
 from cosmotech.coal.store.csv import store_csv_file
 from cosmotech.coal.store.store import Store
 from cosmotech.coal.utils.logger import LOGGER
-from cosmotech.orchestrator.utils.translate import T
 
 
 def _get_data(table_name: str, output_directory: str, cursor) -> None: