From fc89a5edb0227c4be28a6352bea32882847540b8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Laurent=20Al=C3=A9p=C3=A9e?= Date: Tue, 29 Sep 2026 15:34:18 +0200 Subject: [PATCH 1/2] add pydocs for commonly used methodes --- cosmotech/coal/cosmotech_api/apis/dataset.py | 25 +++++++++++ cosmotech/coal/cosmotech_api/apis/run.py | 13 ++++++ cosmotech/coal/cosmotech_api/apis/runner.py | 15 +++++++ .../coal/cosmotech_api/apis/workspace.py | 43 +++++++++++++++++++ .../coal/cosmotech_api/objects/connection.py | 11 +++++ .../coal/cosmotech_api/objects/parameters.py | 26 +++++++++++ cosmotech/coal/store/output/aws_channel.py | 13 ++++++ .../coal/store/output/az_storage_channel.py | 11 +++++ .../coal/store/output/channel_interface.py | 7 +++ .../coal/store/output/channel_spliter.py | 15 +++++++ .../coal/store/output/postgres_channel.py | 11 +++++ 11 files changed, 190 insertions(+) diff --git a/cosmotech/coal/cosmotech_api/apis/dataset.py b/cosmotech/coal/cosmotech_api/apis/dataset.py index 528cd9eb..a90e1ef6 100644 --- a/cosmotech/coal/cosmotech_api/apis/dataset.py +++ b/cosmotech/coal/cosmotech_api/apis/dataset.py @@ -35,6 +35,14 @@ def __init__( LOGGER.debug(T("coal.cosmotech_api.initialization.dataset_api_initialized")) def download_dataset(self, dataset_id) -> Dataset: + """Download all parts of a dataset to the configured dataset folder. + + Args: + dataset_id: The ID of the dataset to download + + Returns: + The Dataset object as returned by the API + """ LOGGER.debug(f"Downloading dataset {dataset_id}") dataset = self.get_dataset( organization_id=self.configuration.cosmotech.organization_id, @@ -48,6 +56,14 @@ def download_dataset(self, dataset_id) -> Dataset: return dataset def download_parameter(self, dataset_id) -> Dataset: + """Download all parts of a dataset to the configured parameters folder. + + Args: + dataset_id: The ID of the dataset to download + + Returns: + The Dataset object as returned by the API + """ LOGGER.debug(f"Downloading dataset {dataset_id}") dataset = self.get_dataset( organization_id=self.configuration.cosmotech.organization_id, @@ -80,6 +96,15 @@ def _download_part(self, dataset_id, dataset_part, destination): @staticmethod def path_to_parts(_path, part_type) -> list[tuple[str, Path, DatasetPartTypeEnum]]: + """Turn a file or directory path into a list of dataset part descriptors. + + Args: + _path: Path to a file, or to a directory whose files are each turned into a part + part_type: The DatasetPartTypeEnum to assign to every resulting part + + Returns: + A list of tuples of (relative part name, absolute file path, part type) + """ if (_path := Path(_path)).is_dir(): return list((str(_p.relative_to(_path)), _p, part_type) for _p in _path.rglob("*") if _p.is_file()) return list(((_path.name, _path, part_type),)) diff --git a/cosmotech/coal/cosmotech_api/apis/run.py b/cosmotech/coal/cosmotech_api/apis/run.py index 717e79e4..b2ff08e1 100644 --- a/cosmotech/coal/cosmotech_api/apis/run.py +++ b/cosmotech/coal/cosmotech_api/apis/run.py @@ -34,5 +34,18 @@ def get_run_metadata( include: Optional[list[str]] = None, exclude: Optional[list[str]] = None, ) -> dict[str, Any]: + """Fetch a run and serialize it to a JSON-compatible dictionary. + + Args: + organization_id: The ID of the organization + workspace_id: The ID of the workspace + runner_id: The ID of the runner + run_id: The ID of the run + include: Optional list of field names to include in the result + exclude: Optional list of field names to exclude from the result + + Returns: + A dictionary representation of the run's metadata + """ run = self.get_run(organization_id, workspace_id, runner_id, run_id) return run.model_dump(by_alias=True, exclude_none=True, include=include, exclude=exclude, mode="json") diff --git a/cosmotech/coal/cosmotech_api/apis/runner.py b/cosmotech/coal/cosmotech_api/apis/runner.py index 1ea62c7d..5d56caf1 100644 --- a/cosmotech/coal/cosmotech_api/apis/runner.py +++ b/cosmotech/coal/cosmotech_api/apis/runner.py @@ -33,6 +33,16 @@ def get_runner_metadata( include: Optional[list[str]] = None, exclude: Optional[list[str]] = None, ) -> dict[str, Any]: + """Fetch a runner and serialize it to a JSON-compatible dictionary. + + Args: + runner_id: The ID of the runner, defaults to the configured runner ID + include: Optional list of field names to include in the result + exclude: Optional list of field names to exclude from the result + + Returns: + A dictionary representation of the runner's metadata + """ runner = self.get_runner( self.configuration.cosmotech.organization_id, self.configuration.cosmotech.workspace_id, @@ -45,6 +55,11 @@ def download_runner_data( self, download_datasets: Optional[str] = None, ): + """Download a runner's parameters and, optionally, its associated datasets. + + Args: + download_datasets: If set, also download the runner's base datasets + """ LOGGER.info(T("coal.cosmotech_api.runner.starting_download")) # Get runner data diff --git a/cosmotech/coal/cosmotech_api/apis/workspace.py b/cosmotech/coal/cosmotech_api/apis/workspace.py index 901cf808..30756569 100644 --- a/cosmotech/coal/cosmotech_api/apis/workspace.py +++ b/cosmotech/coal/cosmotech_api/apis/workspace.py @@ -32,6 +32,19 @@ def list_filtered_workspace_files( workspace_id: str, file_prefix: str, ) -> list[str]: + """List workspace files whose name starts with the given prefix. + + Args: + organization_id: The ID of the organization + workspace_id: The ID of the workspace + file_prefix: The prefix to filter workspace file names by + + Returns: + List of matching workspace file names + + Raises: + ValueError: If no workspace file matches the given prefix + """ target_list = [] LOGGER.info(T("coal.cosmotech_api.workspace.target_is_folder")) wsf = self.list_workspace_files(organization_id, workspace_id) @@ -60,6 +73,20 @@ def download_workspace_file( file_name: str, target_dir: Path, ) -> Path: + """Download a single workspace file to a local directory. + + Args: + organization_id: The ID of the organization + workspace_id: The ID of the workspace + file_name: The name of the workspace file to download + target_dir: The local directory to download the file into + + Returns: + The local path of the downloaded file + + Raises: + ValueError: If target_dir is not a directory + """ if target_dir.is_file(): raise ValueError(T("coal.common.file_operations.not_directory").format(target_dir=target_dir)) @@ -85,6 +112,22 @@ def upload_workspace_file( workspace_path: str, overwrite: bool = True, ) -> str: + """Upload a local file to a workspace. + + Args: + organization_id: The ID of the organization + workspace_id: The ID of the workspace + file_path: Local path of the file to upload + workspace_path: Destination path (or directory ending with '/') in the workspace + overwrite: If True, overwrite an existing file at the destination + + Returns: + The name of the uploaded workspace file + + Raises: + ValueError: If file_path does not exist or is not a single file + ApiException: If the API call fails, e.g. because the file already exists + """ target_file = Path(file_path) if not target_file.exists(): LOGGER.error(T("coal.common.file_operations.not_exists").format(file_path=file_path)) diff --git a/cosmotech/coal/cosmotech_api/objects/connection.py b/cosmotech/coal/cosmotech_api/objects/connection.py index 002f5f50..aae6bbab 100644 --- a/cosmotech/coal/cosmotech_api/objects/connection.py +++ b/cosmotech/coal/cosmotech_api/objects/connection.py @@ -53,6 +53,17 @@ def __init__( self.api_client, self.api_type = self.get_api_client() def get_api_client(self) -> (cosmotech_api.ApiClient, str): + """Build an authenticated API client from available environment variables. + + Detects Keycloak, Cosmo Tech API key, or Azure Entra credentials from the + environment and builds the corresponding API client. + + Returns: + A tuple of the configured ApiClient and a label describing the connection type used + + Raises: + EnvironmentError: If no complete set of required environment variables is found + """ existing_keys = set(os.environ.keys()) missing_azure_keys = self.__azure_env_keys - existing_keys missing_api_keys = self.__api_env_keys - existing_keys diff --git a/cosmotech/coal/cosmotech_api/objects/parameters.py b/cosmotech/coal/cosmotech_api/objects/parameters.py index 34f6d356..5170208c 100644 --- a/cosmotech/coal/cosmotech_api/objects/parameters.py +++ b/cosmotech/coal/cosmotech_api/objects/parameters.py @@ -84,6 +84,14 @@ def write_parameters_to_json( self, parameter_folder: str, ) -> str: + """Write the parameters list to a parameters.json file. + + Args: + parameter_folder: Directory in which to create the JSON file + + Returns: + The path of the written JSON file + """ pathlib.Path(parameter_folder).mkdir(exist_ok=True, parents=True) tmp_parameter_file = os.path.join(parameter_folder, "parameters.json") @@ -98,6 +106,14 @@ def write_parameters_to_csv( self, parameter_folder: str, ) -> str: + """Write the parameters list to a parameters.csv file. + + Args: + parameter_folder: Directory in which to create the CSV file + + Returns: + The path of the written CSV file + """ pathlib.Path(parameter_folder).mkdir(exist_ok=True, parents=True) tmp_parameter_file = os.path.join(parameter_folder, "parameters.csv") @@ -116,6 +132,16 @@ def write_parameters( write_csv: bool = True, write_json: bool = False, ) -> Dict[str, str]: + """Write the parameters list to CSV and/or JSON files. + + Args: + parameter_folder: Directory in which to create the files + write_csv: If True, write a parameters.csv file + write_json: If True, write a parameters.json file + + Returns: + Dictionary mapping format name ("csv", "json") to the written file path + """ result = {} if write_csv: diff --git a/cosmotech/coal/store/output/aws_channel.py b/cosmotech/coal/store/output/aws_channel.py index 13590683..eb978eb8 100644 --- a/cosmotech/coal/store/output/aws_channel.py +++ b/cosmotech/coal/store/output/aws_channel.py @@ -15,6 +15,8 @@ class AwsChannel(ChannelInterface): + """Output channel that sends the store's tables or database file to an S3 bucket.""" + required_keys = { "coal": ["store"], "cosmotech": ["runner_id"], @@ -28,7 +30,17 @@ def __init__(self, dct: Dotdict = None): self._s3.file_prefix = self.configuration.cosmotech.runner_id + "/" + self._s3.file_prefix def send(self, filter: Optional[list[str]] = None) -> bool: + """Upload the store's database file or tables to the configured S3 bucket. + + Args: + filter: Optional list of table names to restrict the upload to + + Returns: + True if the data was uploaded successfully + Raises: + ValueError: If the configured output type is not sqlite, csv or parquet + """ _s = Store(configuration=self.configuration) if self._s3.output_type not in ("sqlite", "csv", "parquet"): @@ -67,4 +79,5 @@ def send(self, filter: Optional[list[str]] = None) -> bool: ) def delete(self): + """Delete all objects previously uploaded to the configured S3 bucket.""" self._s3.delete_objects() diff --git a/cosmotech/coal/store/output/az_storage_channel.py b/cosmotech/coal/store/output/az_storage_channel.py index a4ee9a8c..5a7273df 100644 --- a/cosmotech/coal/store/output/az_storage_channel.py +++ b/cosmotech/coal/store/output/az_storage_channel.py @@ -8,6 +8,8 @@ class AzureStorageChannel(ChannelInterface): + """Output channel that sends the store's tables to an Azure Blob Storage container.""" + required_keys = { "coal": ["store"], "cosmotech": ["runner_id"], @@ -31,10 +33,19 @@ def __init__(self, dct: Dotdict = None): self.configuration.azure.file_prefix = runner_id + "/" + prefix def send(self, filter: Optional[list[str]] = None) -> bool: + """Dump the store's tables to the configured Azure Blob Storage container. + + Args: + filter: Optional list of table names to restrict the dump to + + Returns: + True if the data was uploaded successfully + """ dump_store_to_azure( self.configuration, selected_tables=filter, ) def delete(self): + """Delete all blobs previously uploaded to the configured Azure Blob Storage container.""" delete_azure_blobs(self.configuration) diff --git a/cosmotech/coal/store/output/channel_interface.py b/cosmotech/coal/store/output/channel_interface.py index de57fb59..f732bd5f 100644 --- a/cosmotech/coal/store/output/channel_interface.py +++ b/cosmotech/coal/store/output/channel_interface.py @@ -6,6 +6,8 @@ class ChannelInterface: + """Base class for an output channel that can send the store's data to a target and delete it.""" + required_keys = {} requirement_string: str = T("coal.store.output.data_interface.requirements") @@ -21,6 +23,11 @@ def delete(self) -> bool: raise NotImplementedError() def is_available(self) -> bool: + """Check whether all required configuration keys are present, falling back to global config. + + Returns: + True if the channel is fully configured and can be used + """ try: for section in self.required_keys.keys(): for key in self.required_keys[section]: diff --git a/cosmotech/coal/store/output/channel_spliter.py b/cosmotech/coal/store/output/channel_spliter.py index e5ff6739..33b3f9a5 100644 --- a/cosmotech/coal/store/output/channel_spliter.py +++ b/cosmotech/coal/store/output/channel_spliter.py @@ -11,6 +11,8 @@ class ChannelSpliter(ChannelInterface): + """Channel that fans out the store's data to every configured output channel.""" + requirement_string: str = "(Requires any working interface)" targets = list() available_interfaces: dict[str, ChannelInterface] = { @@ -39,6 +41,14 @@ def __init__(self, dct: Dotdict = None): raise AttributeError(T("coal.store.output.split.no_targets")) def send(self, filter: Optional[list[str]] = None) -> bool: + """Send the store's data through every available target channel. + + Args: + filter: Optional list of table names to restrict the send to + + Returns: + True if at least one target channel sent its data successfully + """ any_ok = False for i in self.targets: try: @@ -50,6 +60,11 @@ def send(self, filter: Optional[list[str]] = None) -> bool: return any_ok def delete(self) -> bool: + """Delete the data previously sent through every available target channel. + + Returns: + True if at least one target channel deleted its data successfully + """ any_ok = False for i in self.targets: try: diff --git a/cosmotech/coal/store/output/postgres_channel.py b/cosmotech/coal/store/output/postgres_channel.py index aa8a0233..9a524d6c 100644 --- a/cosmotech/coal/store/output/postgres_channel.py +++ b/cosmotech/coal/store/output/postgres_channel.py @@ -14,6 +14,8 @@ class PostgresChannel(ChannelInterface): + """Output channel that sends the store's tables and runner metadata to a PostgreSQL database.""" + required_keys = { "coal": ["store"], "cosmotech": ["organization_id", "workspace_id", "runner_id"], @@ -34,6 +36,14 @@ def __init__(self, dct: Dotdict = None): self.configuration.setup_db = self.configuration.safe_get("setup_db", True) def send(self, filter: Optional[list[str]] = None) -> bool: + """Send the runner's metadata and the store's tables to the configured PostgreSQL database. + + Args: + filter: Optional list of table names to restrict the send to + + Returns: + True if the data was sent successfully + """ if self.configuration.setup_db: create_metadata(self.configuration) @@ -55,5 +65,6 @@ def send(self, filter: Optional[list[str]] = None) -> bool: add_fk_constraints(self.configuration) def delete(self): + """Delete the runner's metadata, cascading to delete its associated data in PostgreSQL.""" # removing metadata will trigger cascade delete on real data remove_runner_metadata_from_postgresql(self.configuration) From 818c370f949f69eddf400fb05c38239b94565213 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Laurent=20Al=C3=A9p=C3=A9e?= Date: Tue, 29 Sep 2026 15:37:07 +0200 Subject: [PATCH 2/2] Bump version from 2.3.1 to 2.4.0 --- cosmotech/coal/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cosmotech/coal/__init__.py b/cosmotech/coal/__init__.py index 2ba21513..fe6d5078 100644 --- a/cosmotech/coal/__init__.py +++ b/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.3.1" +__version__ = "2.4.0"