diff --git a/great_ai/large_file/large_file/large_file_base.py b/great_ai/large_file/large_file/large_file_base.py index 2f9ff2b..5bee7db 100644 --- a/great_ai/large_file/large_file/large_file_base.py +++ b/great_ai/large_file/large_file/large_file_base.py @@ -84,9 +84,11 @@ class LargeFileBase(ABC): @classmethod def configure_credentials_from_file( cls, - secrets_path: Union[Path, str], + secrets: Union[Path, str, ConfigFile], ) -> None: - cls.configure_credentials(**ConfigFile(secrets_path)) + if not isinstance(secrets, ConfigFile): + secrets = ConfigFile(secrets) + cls.configure_credentials(**{k.lower(): v for k, v in secrets.items()}) @classmethod def configure_credentials(cls, **kwargs: str) -> None: diff --git a/great_ai/persistence/mongodb_driver.py b/great_ai/persistence/mongodb_driver.py index 7c2600d..49e7b3e 100644 --- a/great_ai/persistence/mongodb_driver.py +++ b/great_ai/persistence/mongodb_driver.py @@ -90,17 +90,8 @@ class MongodbDriver(TracingDatabaseDriver): query: Dict[str, Any] = { "filter": {}, - "sort": [ - (col.column_id, 1 if col.direction == "asc" else -1) for col in sort_by - ], } - if skip: - query["skip"] = skip - - if take: - query["limit"] = take - and_query: List[Dict[str, Any]] = [{}] and_query.extend({"tags": tag} for tag in conjunctive_tags) and_query.extend( @@ -120,10 +111,21 @@ class MongodbDriver(TracingDatabaseDriver): query["filter"]["$and"] = and_query with MongoClient[Any](self.mongo_connection_string) as client: - values = client[self.mongo_database].traces.find(**query) - documents = [Trace[Any].parse_obj(t) for t in values] + count = client[self.mongo_database].traces.count_documents(**query) - return documents, len(documents) + if skip: + query["skip"] = skip + + if take: + query["limit"] = take + + query["sort"] = [ + (col.column_id, 1 if col.direction == "asc" else -1) for col in sort_by + ] + + with client[self.mongo_database].traces.find(**query) as cursor: + documents = [Trace[Any].parse_obj(t) for t in cursor] + return documents, count def update(self, id: str, new_version: Trace) -> None: serialized = new_version.to_flat_dict() diff --git a/great_ai/persistence/tracing_database_driver.py b/great_ai/persistence/tracing_database_driver.py index 17cccfc..4995361 100644 --- a/great_ai/persistence/tracing_database_driver.py +++ b/great_ai/persistence/tracing_database_driver.py @@ -14,9 +14,11 @@ class TracingDatabaseDriver(ABC): @classmethod def configure_credentials_from_file( cls, - secrets_path: Union[Path, str], + secrets: Union[Path, str, ConfigFile], ) -> None: - cls.configure_credentials(**ConfigFile(secrets_path)) + if not isinstance(secrets, ConfigFile): + secrets = ConfigFile(secrets) + cls.configure_credentials(**{k.lower(): v for k, v in secrets.items()}) @classmethod def configure_credentials( diff --git a/great_ai/utilities/config_file/config_file.py b/great_ai/utilities/config_file/config_file.py index 92a4951..0161842 100644 --- a/great_ai/utilities/config_file/config_file.py +++ b/great_ai/utilities/config_file/config_file.py @@ -6,24 +6,30 @@ from ..logger import get_logger from .parse_error import ParseError from .pattern import pattern -ENVIRONMENT_VARIABLE_KEY_PREFIX = "ENV" - logger = get_logger("ConfigFile") class ConfigFile(Mapping[str, str]): - def __init__(self, path: Union[Path, str]) -> None: + ENVIRONMENT_VARIABLE_KEY_PREFIX = "ENV" + + def __init__(self, path: Union[Path, str], *, ignore_missing: bool = False) -> None: if not isinstance(path, Path): path = Path(path) if not path.exists(): raise FileNotFoundError(path.absolute()) + self._ignore_missing = ignore_missing + self._path = path self._key_values: Dict[str, str] = {} self._parse() + @property + def path(self) -> Path: + return self._path + def _parse(self) -> None: with open(self._path, encoding="utf-8") as f: lines: str = f.read() @@ -39,21 +45,23 @@ class ConfigFile(Mapping[str, str]): already_exists = key in self._key_values if already_exists and not value.startswith( - f"{ENVIRONMENT_VARIABLE_KEY_PREFIX}:" + f"{self.ENVIRONMENT_VARIABLE_KEY_PREFIX}:" ): raise KeyError( f"Key `{key}` has been already defined and its value is `{self._key_values[key]}`" ) - if value.startswith(f"{ENVIRONMENT_VARIABLE_KEY_PREFIX}:"): + if value.startswith(f"{self.ENVIRONMENT_VARIABLE_KEY_PREFIX}:"): _, value = value.split(":") if value not in os.environ: - issue = f'The value of `{key}` contains the "{ENVIRONMENT_VARIABLE_KEY_PREFIX}` prefix but `{value}` is not defined as an environment variable' + issue = f'The value of `{key}` contains the "{self.ENVIRONMENT_VARIABLE_KEY_PREFIX}` prefix but `{value}` is not defined as an environment variable' if already_exists: logger.warning( f"{issue}, using the default value defined above (`{self._key_values[key]}`)" ) continue + elif self._ignore_missing: + logger.warning(issue) else: raise KeyError( f"{issue} and no default value has been provided"