Skip to content
Merged
Show file tree
Hide file tree
Changes from 8 commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion cognite/extractorutils/_inner_util.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@
# limitations under the License.

"""
A module containing utilities meant for use inside the extractor-utils package
A module containing utilities meant for use inside the extractor-utils package.
"""

import json
Expand Down
54 changes: 44 additions & 10 deletions cognite/extractorutils/base.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,6 @@
"""
Module containing the base class for extractors.
"""
# Copyright 2021 Cognite AS
#
# Licensed under the Apache License, Version 2.0 (the "License");
Expand Down Expand Up @@ -35,6 +38,10 @@


class ReloadConfigAction(Enum):
"""
Enum for actions to take when a config file is reloaded.
"""

DO_NOTHING = 1
REPLACE_ATTRIBUTE = 2
SHUTDOWN = 3
Expand Down Expand Up @@ -126,8 +133,10 @@ def __init__(

def _initial_load_config(self, override_path: str | None = None) -> None:
"""
Load a configuration file, either from the specified path, or by a path specified by the user in a command line
arg. Will quit further execution of no path is given.
Load a configuration file.

Either from the specified path, or by a path specified by the user in a command line arg. Will quit further
execution of no path is given.

Args:
override_path: Optional override for file path, ie don't parse command line arguments
Expand All @@ -150,6 +159,11 @@ def config_refresher() -> None:
Thread(target=config_refresher, name="ConfigReloader", daemon=True).start()

def reload_config_callback(self) -> None:
"""
If the reload_config_action was set to CALLBACK, this method will be called when the config file is reloaded.

This method should be overridden in subclasses to provide custom behavior when the config file is reloaded.
"""
self.logger.error("Method for reloading configs has not been overridden in subclass")

def _reload_config(self) -> None:
Expand All @@ -173,9 +187,11 @@ def _reload_config(self) -> None:

def _load_state_store(self) -> None:
"""
Searches through the config object for a StateStoreConfig. If found, it will use that configuration to generate
a state store, if no such config is found it will either create a LocalStateStore or a NoStateStore depending
on whether the ``use_default_state_store`` argument to the constructor was true or false.
Searches through the config object for a StateStoreConfig.

If found, it will use that configuration to generate a state store, if no such config is found it will either
create a LocalStateStore or a NoStateStore depending on whether the ``use_default_state_store`` argument to the
constructor was true or false.

Either way, the state_store attribute is guaranteed to be set after calling this method.
"""
Expand Down Expand Up @@ -213,7 +229,7 @@ def recursive_find_state_store(d: dict[str, Any]) -> StateStoreConfig | None:

def _report_success(self) -> None:
"""
Called on a successful exit of the extractor
Called on a successful exit of the extractor.
"""
if self.extraction_pipeline:
self.logger.info("Reporting new successful run")
Expand All @@ -227,7 +243,7 @@ def _report_success(self) -> None:

def _report_error(self, exception: BaseException) -> None:
"""
Called on an unsuccessful exit of the extractor
Called on an unsuccessful exit of the extractor.

Args:
exception: Exception object that caused the extractor to fail
Expand All @@ -250,7 +266,6 @@ def __enter__(self) -> "Extractor":
Returns:
self
"""

if str(os.getenv("COGNITE_FUNCTION_RUNTIME", False)).lower() != "true":
# Environment Variables
env_file_found = load_dotenv(dotenv_path="./.env", override=True)
Expand Down Expand Up @@ -357,8 +372,9 @@ def __exit__(

def run(self) -> None:
"""
Run the extractor. Ensures that the Extractor is set up correctly (``run`` called within a ``with``) and calls
the ``run_handle``.
Run the extractor.

Ensures that the Extractor is set up correctly (``run`` called within a ``with``) and calls the ``run_handle``.

Can be overrided in subclasses.
"""
Expand All @@ -371,12 +387,30 @@ def run(self) -> None:

@classmethod
def get_current_config(cls) -> CustomConfigClass:
"""
Get the current configuration singleton.

Returns:
The current configuration singleton

Raises:
ValueError: If no configuration singleton has been created, meaning no config file has been loaded.
"""
if Extractor._config_singleton is None: # type: ignore
raise ValueError("No config singleton created. Have a config file been loaded?")
return Extractor._config_singleton # type: ignore

@classmethod
def get_current_statestore(cls) -> AbstractStateStore:
"""
Get the current state store singleton.

Returns:
The current state store singleton

Raises:
ValueError: If no state store singleton has been created, meaning no state store has been loaded.
"""
if Extractor._statestore_singleton is None:
raise ValueError("No state store singleton created. Have a state store been loaded?")
return Extractor._statestore_singleton
3 changes: 1 addition & 2 deletions cognite/extractorutils/configtools/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,8 +13,7 @@
# limitations under the License.

"""
Module containing tools for loading and verifying config files, and a YAML loader to automatically serialize these
dataclasses from a config file.
Module containing tools for loading and verifying config files.

Configs are described as ``dataclass`` es, and use the ``BaseConfig`` class as a superclass to get a few things
built-in: config version, Cognite project and logging. Use type hints to specify types, use the ``Optional`` type to
Expand Down
5 changes: 3 additions & 2 deletions cognite/extractorutils/configtools/_util.py
Original file line number Diff line number Diff line change
Expand Up @@ -27,8 +27,9 @@

def _to_snake_case(dictionary: dict[str, Any], case_style: str) -> dict[str, Any]:
"""
Ensure that all keys in the dictionary follows the snake casing convention (recursively, so any sub-dictionaries are
changed too).
Ensure that all keys in the dictionary follows the snake casing convention.

This function will recursively fix any list or dictionaries.

Args:
dictionary: Dictionary to update.
Expand Down
Loading
Loading