databao-context-engine 0.1.1__py3-none-any.whl → 0.1.5__py3-none-any.whl

databao_context_engine/{datasource_config → datasources}/check_config.py

@@ -2,25 +2,26 @@ import logging
 import os
 from dataclasses import dataclass
 from enum import Enum
-from pathlib import Path
 
 from pydantic import ValidationError
 
-from databao_context_engine.pluginlib.build_plugin import BuildDatasourcePlugin, NotSupportedError
-from databao_context_engine.pluginlib.plugin_utils import check_connection_for_datasource
-from databao_context_engine.plugins.plugin_loader import load_plugins
-from databao_context_engine.project.datasource_discovery import (
+from databao_context_engine.datasources.datasource_discovery import (
     discover_datasources,
     get_datasource_descriptors,
     prepare_source,
 )
-from databao_context_engine.project.layout import ensure_project_dir
-from databao_context_engine.project.types import DatasourceId, PreparedConfig
+from databao_context_engine.datasources.types import DatasourceId, PreparedConfig
+from databao_context_engine.pluginlib.build_plugin import BuildDatasourcePlugin, NotSupportedError
+from databao_context_engine.pluginlib.plugin_utils import check_connection_for_datasource
+from databao_context_engine.plugins.plugin_loader import load_plugins
+from databao_context_engine.project.layout import ProjectLayout
 
 logger = logging.getLogger(__name__)
 
 
 class DatasourceConnectionStatus(Enum):
+    """Status of the connection to a datasource."""
+
     VALID = "Valid"
     INVALID = "Invalid"
     UNKNOWN = "Unknown"
@@ -28,6 +29,15 @@ class DatasourceConnectionStatus(Enum):
 
 @dataclass(kw_only=True)
 class CheckDatasourceConnectionResult:
+    """Result of checking the connection status of a datasource.
+
+    Attributes:
+        datasource_id: The id of the datasource.
+        connection_status: The connection status of the datasource.
+        summary: A summary of the connection status' error, or None if the connection is valid.
+        full_message: A detailed message about the connection status' error, or None if the connection is valid.
+    """
+
     datasource_id: DatasourceId
     connection_status: DatasourceConnectionStatus
     summary: str | None
@@ -44,21 +54,19 @@ class CheckDatasourceConnectionResult:
 
 
 def check_datasource_connection(
-    project_dir: Path, *, datasource_ids: list[DatasourceId] | None = None
+    project_layout: ProjectLayout, *, datasource_ids: list[DatasourceId] | None = None
 ) -> dict[DatasourceId, CheckDatasourceConnectionResult]:
-    ensure_project_dir(project_dir)
-
     if datasource_ids:
         logger.info(f"Validating datasource(s): {datasource_ids}")
-        datasources_to_traverse = get_datasource_descriptors(project_dir, datasource_ids)
+        datasources_to_traverse = get_datasource_descriptors(project_layout, datasource_ids)
     else:
-        datasources_to_traverse = discover_datasources(project_dir)
+        datasources_to_traverse = discover_datasources(project_layout)
 
     plugins = load_plugins(exclude_file_plugins=True)
 
     result = {}
     for discovered_datasource in datasources_to_traverse:
-        result_key = DatasourceId.from_datasource_config_file_path(discovered_datasource.path)
+        result_key = DatasourceId.from_datasource_config_file_path(project_layout, discovered_datasource.path)
 
         try:
             prepared_source = prepare_source(discovered_datasource)
@@ -103,12 +111,12 @@ def check_datasource_connection(
                     exc_info=True,
                     stack_info=True,
                 )
-                result[result_key] = get_validation_result_from_error(result_key, e)
+                result[result_key] = _get_validation_result_from_error(result_key, e)
 
     return result
 
 
-def get_validation_result_from_error(datasource_id: DatasourceId, e: Exception):
+def _get_validation_result_from_error(datasource_id: DatasourceId, e: Exception):
     if isinstance(e, ValidationError):
         return CheckDatasourceConnectionResult(
             datasource_id=datasource_id,
@@ -116,16 +124,16 @@ def get_validation_result_from_error(datasource_id: DatasourceId, e: Exception):
             summary="Config file is invalid",
             full_message=str(e),
         )
-    elif isinstance(e, NotImplementedError | NotSupportedError):
+    if isinstance(e, NotImplementedError | NotSupportedError):
         return CheckDatasourceConnectionResult(
             datasource_id=datasource_id,
             connection_status=DatasourceConnectionStatus.UNKNOWN,
             summary="Plugin doesn't support validating its config",
         )
-    else:
-        return CheckDatasourceConnectionResult(
-            datasource_id=datasource_id,
-            connection_status=DatasourceConnectionStatus.INVALID,
-            summary="Connection with the datasource can not be established",
-            full_message=str(e),
-        )
+
+    return CheckDatasourceConnectionResult(
+        datasource_id=datasource_id,
+        connection_status=DatasourceConnectionStatus.INVALID,
+        summary="Connection with the datasource can not be established",
+        full_message=str(e),
+    )

databao_context_engine/datasources/datasource_context.py

@@ -0,0 +1,90 @@
+import logging
+import os
+from dataclasses import dataclass
+from pathlib import Path
+
+import yaml
+
+from databao_context_engine.datasources.types import Datasource, DatasourceId, DatasourceType
+from databao_context_engine.project.layout import ProjectLayout
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(eq=True, frozen=True)
+class DatasourceContext:
+    """A generated Context for a Datasource.
+
+    Attributes:
+        datasource_id: The id of the datasource.
+        context: The context generated for the datasource.
+    """
+
+    datasource_id: DatasourceId
+    # TODO: Read the context as a BuildExecutionResult instead of a Yaml string?
+    context: str
+
+
+def _read_datasource_type_from_context_file(context_path: Path) -> DatasourceType:
+    with context_path.open("r") as context_file:
+        type_key = "datasource_type"
+        for line in context_file:
+            if line.startswith(f"{type_key}: "):
+                datasource_type = yaml.safe_load(line)[type_key]
+                return DatasourceType(full_type=datasource_type)
+
+    raise ValueError(f"Could not find type in context file {context_path}")
+
+
+def get_introspected_datasource_list(project_layout: ProjectLayout) -> list[Datasource]:
+    result = []
+    for dirpath, dirnames, filenames in os.walk(project_layout.output_dir):
+        for context_file_name in filenames:
+            if Path(context_file_name).suffix not in DatasourceId.ALLOWED_YAML_SUFFIXES:
+                continue
+            context_file = Path(dirpath).joinpath(context_file_name)
+            relative_context_file = context_file.relative_to(project_layout.output_dir)
+            try:
+                result.append(
+                    Datasource(
+                        id=DatasourceId.from_datasource_context_file_path(relative_context_file),
+                        type=_read_datasource_type_from_context_file(context_file),
+                    )
+                )
+            except ValueError as e:
+                logger.debug(str(e), exc_info=True, stack_info=True)
+                logger.warning(
+                    f"Ignoring introspected datasource: Failed to read datasource_type from context file at {context_file.resolve()}"
+                )
+
+    return sorted(result, key=lambda ds: str(ds.id).lower())
+
+
+def get_datasource_context(project_layout: ProjectLayout, datasource_id: DatasourceId) -> DatasourceContext:
+    context_path = project_layout.output_dir.joinpath(datasource_id.relative_path_to_context_file())
+    if not context_path.is_file():
+        raise ValueError(f"Context file not found for datasource {str(datasource_id)}")
+
+    context = context_path.read_text()
+    return DatasourceContext(datasource_id=datasource_id, context=context)
+
+
+def get_all_contexts(project_layout: ProjectLayout) -> list[DatasourceContext]:
+    result = []
+    for dirpath, dirnames, filenames in os.walk(project_layout.output_dir):
+        for context_file_name in filenames:
+            if Path(context_file_name).suffix not in DatasourceId.ALLOWED_YAML_SUFFIXES:
+                continue
+            context_file = Path(dirpath).joinpath(context_file_name)
+            relative_context_file = context_file.relative_to(project_layout.output_dir)
+            result.append(
+                DatasourceContext(
+                    datasource_id=DatasourceId.from_datasource_context_file_path(relative_context_file),
+                    context=context_file.read_text(),
+                )
+            )
+    return sorted(result, key=lambda context: str(context.datasource_id).lower())
+
+
+def get_context_header_for_datasource(datasource_id: DatasourceId) -> str:
+    return f"# ===== {str(datasource_id)} ====={os.linesep}"

databao_context_engine/datasources/datasource_discovery.py

@@ -0,0 +1,143 @@
+import logging
+import os
+from pathlib import Path
+from typing import Any
+
+import yaml
+
+from databao_context_engine.datasources.types import (
+    Datasource,
+    DatasourceDescriptor,
+    DatasourceId,
+    DatasourceKind,
+    PreparedConfig,
+    PreparedDatasource,
+    PreparedFile,
+)
+from databao_context_engine.pluginlib.build_plugin import DatasourceType
+from databao_context_engine.project.layout import ProjectLayout
+from databao_context_engine.templating.renderer import render_template
+
+logger = logging.getLogger(__name__)
+
+
+def get_datasource_list(project_layout: ProjectLayout) -> list[Datasource]:
+    result = []
+    for discovered_datasource in discover_datasources(project_layout=project_layout):
+        try:
+            prepared_source = prepare_source(discovered_datasource)
+        except Exception as e:
+            logger.debug(str(e), exc_info=True, stack_info=True)
+            logger.info(f"Invalid source at ({discovered_datasource.path}): {str(e)}")
+            continue
+
+        result.append(
+            Datasource(
+                id=DatasourceId.from_datasource_config_file_path(project_layout, discovered_datasource.path),
+                type=prepared_source.datasource_type,
+            )
+        )
+
+    return result
+
+
+def discover_datasources(project_layout: ProjectLayout) -> list[DatasourceDescriptor]:
+    """Scan the project's src/ directory and return all discovered sources.
+
+    Rules:
+        - Each first-level directory under src/ is treated as a main_type
+        - Unsupported or unreadable entries are skipped.
+        - The returned list is sorted by directory and then filename
+
+    Args:
+        project_layout: ProjectLayout instance representing the project directory and configuration.
+
+    Returns:
+        A list of DatasourceDescriptor instances representing the discovered datasources.
+    """
+    datasources: list[DatasourceDescriptor] = []
+    for dirpath, dirnames, filenames in os.walk(project_layout.src_dir):
+        for config_file_name in filenames:
+            context_file = Path(dirpath).joinpath(config_file_name)
+            datasource = _load_datasource_descriptor(project_layout, context_file)
+            if datasource is not None:
+                datasources.append(datasource)
+
+    return sorted(datasources, key=lambda ds: str(ds.datasource_id.relative_path_to_config_file()).lower())
+
+
+def _is_datasource_file(p: Path) -> bool:
+    # ignore backup files
+    return p.is_file() and not p.suffix.endswith("~")
+
+
+def get_datasource_descriptors(
+    project_layout: ProjectLayout, datasource_ids: list[DatasourceId]
+) -> list[DatasourceDescriptor]:
+    datasources: list[DatasourceDescriptor] = []
+    for datasource_id in sorted(datasource_ids, key=lambda id: str(id)):
+        config_file_path = project_layout.src_dir.joinpath(datasource_id.relative_path_to_config_file())
+        if not config_file_path.is_file():
+            raise ValueError(f"Datasource config file not found: {config_file_path}")
+
+        datasource = _load_datasource_descriptor(project_layout, config_file_path)
+        if datasource is not None:
+            datasources.append(datasource)
+
+    return datasources
+
+
+def _load_datasource_descriptor(project_layout: ProjectLayout, config_file: Path) -> DatasourceDescriptor | None:
+    """Load a single file with src/<parent_name>/ into a DatasourceDescriptor."""
+    if not config_file.is_file():
+        return None
+
+    parent_name = config_file.parent.name
+    extension = config_file.suffix.lower().lstrip(".")
+    relative_config_file = config_file.relative_to(project_layout.src_dir)
+
+    if parent_name == "files" and len(relative_config_file.parts) == 2:
+        datasource_id = DatasourceId.from_datasource_config_file_path(project_layout, config_file)
+        return DatasourceDescriptor(datasource_id=datasource_id, path=config_file.resolve(), kind=DatasourceKind.FILE)
+
+    if extension in {"yaml", "yml"}:
+        datasource_id = DatasourceId.from_datasource_config_file_path(project_layout, config_file)
+        return DatasourceDescriptor(datasource_id=datasource_id, path=config_file.resolve(), kind=DatasourceKind.CONFIG)
+
+    if extension:
+        datasource_id = DatasourceId.from_datasource_config_file_path(project_layout, config_file)
+        return DatasourceDescriptor(datasource_id=datasource_id, path=config_file.resolve(), kind=DatasourceKind.FILE)
+
+    logger.debug("Skipping file without extension: %s", config_file)
+    return None
+
+
+def prepare_source(datasource: DatasourceDescriptor) -> PreparedDatasource:
+    """Convert a discovered datasource into a prepared datasource ready for plugin execution."""
+    if datasource.kind is DatasourceKind.FILE:
+        file_subtype = datasource.path.suffix.lower().lstrip(".")
+        return PreparedFile(
+            datasource_id=datasource.datasource_id,
+            datasource_type=DatasourceType(full_type=file_subtype),
+            path=datasource.path,
+        )
+
+    config = _parse_config_file(datasource.path)
+
+    ds_type = config.get("type")
+    if not ds_type or not isinstance(ds_type, str):
+        raise ValueError("Config missing 'type' at %s - skipping", datasource.path)
+
+    return PreparedConfig(
+        datasource_id=datasource.datasource_id,
+        datasource_type=DatasourceType(full_type=ds_type),
+        path=datasource.path,
+        config=config,
+        datasource_name=datasource.path.stem,
+    )
+
+
+def _parse_config_file(file_path: Path) -> dict[Any, Any]:
+    rendered_file = render_template(file_path.read_text())
+
+    return yaml.safe_load(rendered_file) or {}

databao_context_engine/datasources/types.py

@@ -0,0 +1,194 @@
+from dataclasses import dataclass
+from enum import StrEnum
+from pathlib import Path
+from typing import Any
+
+from databao_context_engine.pluginlib.build_plugin import DatasourceType
+from databao_context_engine.project.layout import ProjectLayout
+
+
+class DatasourceKind(StrEnum):
+    CONFIG = "config"
+    FILE = "file"
+
+
+@dataclass(frozen=True)
+class DatasourceDescriptor:
+    datasource_id: "DatasourceId"
+    path: Path
+    kind: DatasourceKind
+
+
+@dataclass(frozen=True)
+class PreparedConfig:
+    datasource_id: "DatasourceId"
+    datasource_type: DatasourceType
+    path: Path
+    config: dict[Any, Any]
+    datasource_name: str
+
+
+@dataclass(frozen=True)
+class PreparedFile:
+    datasource_id: "DatasourceId"
+    datasource_type: DatasourceType
+    path: Path
+
+
+PreparedDatasource = PreparedConfig | PreparedFile
+
+
+@dataclass(kw_only=True, frozen=True, eq=True)
+class DatasourceId:
+    """The ID of a datasource. The ID is the path to the datasource's config file relative to the src folder in the project.
+
+    e.g: "databases/my_postgres_datasource.yaml"
+
+    Use the provided factory methods `from_string_repr` and `from_datasource_config_file_path` to create a DatasourceId, rather than its constructor.
+
+    Attributes:
+        datasource_path: The path to the datasource relative to project's src folder.
+        config_file_suffix: The suffix of the config (or raw) file.
+    """
+
+    ALLOWED_YAML_SUFFIXES = [".yaml", ".yml"]
+    CONTEXT_FILE_SUFFIX = ".yaml"
+
+    datasource_path: str
+    config_file_suffix: str
+
+    def __post_init__(self):
+        if not self.datasource_path.strip():
+            raise ValueError(f"Invalid DatasourceId ({str(self)}): datasource_path must not be empty")
+        if not self.config_file_suffix.strip():
+            raise ValueError(f"Invalid DatasourceId ({str(self)}): config_file_suffix must not be empty")
+
+        if not self.config_file_suffix.startswith("."):
+            raise ValueError(
+                f'Invalid DatasourceId ({str(self)}): config_file_suffix must start with a dot "." (e.g.: .yaml)'
+            )
+
+        if self.datasource_path.endswith(self.config_file_suffix):
+            raise ValueError(f"Invalid DatasourceId ({str(self)}): datasource_path must not contain the file suffix")
+
+    def __str__(self):
+        return str(self.relative_path_to_config_file())
+
+    def relative_path_to_config_file(self) -> Path:
+        """Return a path to the config file for this datasource.
+
+        The returned path is relative to the src folder in the project.
+
+        Returns:
+            The path to the config file relative to the src folder in the project.
+        """
+        return Path(self.datasource_path + self.config_file_suffix)
+
+    def relative_path_to_context_file(self) -> Path:
+        """Return a path to the config file for this datasource.
+
+        The returned path is relative to an output run folder in the project.
+
+        Returns:
+            The path to the context file relative to the output folder in the project.
+        """
+        # Keep the suffix in the filename if this datasource is a raw file, to handle multiple files with the same name and different extensions
+        suffix = (
+            DatasourceId.CONTEXT_FILE_SUFFIX
+            if self.config_file_suffix in DatasourceId.ALLOWED_YAML_SUFFIXES
+            else (self.config_file_suffix + DatasourceId.CONTEXT_FILE_SUFFIX)
+        )
+
+        return Path(self.datasource_path + suffix)
+
+    @classmethod
+    def from_string_repr(cls, datasource_id_as_string: str) -> "DatasourceId":
+        """Create a DatasourceId from a string representation.
+
+        Args:
+            datasource_id_as_string: The string representation of a DatasourceId.
+                This is the path to the datasource's config file relative to the src folder in the project.
+                (e.g. "databases/my_postgres_datasource.yaml")
+
+        Returns:
+            The DatasourceId instance created from the string representation.
+        """
+        config_file_path = Path(datasource_id_as_string)
+        return DatasourceId._from_relative_datasource_config_file_path(config_file_path)
+
+    @classmethod
+    def from_datasource_config_file_path(
+        cls, project_layout: ProjectLayout, datasource_config_file: Path
+    ) -> "DatasourceId":
+        """Create a DatasourceId from a config file path.
+
+        Args:
+            project_layout: The databao context engine project layout.
+            datasource_config_file: The path to the datasource config file.
+                This path can either be the config file path relative to the src folder or the full path to the config file.
+
+        Returns:
+            The DatasourceId instance created from the config file path.
+        """
+        relative_datasource_config_file = datasource_config_file.relative_to(project_layout.src_dir)
+        return DatasourceId._from_relative_datasource_config_file_path(relative_datasource_config_file)
+
+    @classmethod
+    def _from_relative_datasource_config_file_path(cls, relative_datasource_config_file: Path) -> "DatasourceId":
+        if relative_datasource_config_file.is_absolute():
+            raise ValueError(
+                f"Path to datasource config file should be relative to project's src folder: {relative_datasource_config_file}"
+            )
+        return DatasourceId(
+            datasource_path=_extract_datasource_path(relative_datasource_config_file),
+            config_file_suffix=relative_datasource_config_file.suffix,
+        )
+
+    @classmethod
+    def from_datasource_context_file_path(cls, datasource_context_file: Path) -> "DatasourceId":
+        """Create a DatasourceId from a context file path.
+
+        This factory handles the case where the context was generated from a raw file rather than from a config.
+        In that case, the context file name will look like "<my_datasource_name>.<raw_file_extension>.yaml"
+
+        Args:
+            datasource_context_file: The path to the datasource context file.
+
+        Returns:
+            The DatasourceId instance created from the context file path.
+
+        Raises:
+            ValueError: If the wrong datasource_context_file is provided.
+        """
+        if (
+            len(datasource_context_file.suffixes) > 1
+            and datasource_context_file.suffix == DatasourceId.CONTEXT_FILE_SUFFIX
+        ):
+            # If there is more than 1 suffix, we remove the latest suffix (.yaml) to keep only the actual datasource file suffix
+            context_file_name_without_yaml_extension = datasource_context_file.name[
+                : -len(DatasourceId.CONTEXT_FILE_SUFFIX)
+            ]
+            datasource_context_file = datasource_context_file.with_name(context_file_name_without_yaml_extension)
+        if datasource_context_file.is_absolute():
+            raise ValueError(f"Path to datasource context file should be a relative path: {datasource_context_file}")
+        return DatasourceId(
+            datasource_path=_extract_datasource_path(datasource_context_file),
+            config_file_suffix=datasource_context_file.suffix,
+        )
+
+
+def _extract_datasource_path(datasource_file: Path) -> str:
+    return str(datasource_file.with_suffix(""))
+
+
+@dataclass
+class Datasource:
+    """A datasource contained in the project.
+
+    Attributes:
+        id: The unique identifier of the datasource.
+        type: The type of the datasource.
+    """
+
+    id: DatasourceId
+    type: DatasourceType

databao_context_engine/generate_configs_schemas.py

@@ -38,9 +38,9 @@ def generate_configs_schemas(
     try:
         schema_list = _generate_json_schema_output_for_plugins(include_plugins, exclude_plugins)
 
-        print(f"{os.linesep}{os.linesep}".join(schema_list))
+        click.echo(f"{os.linesep}{os.linesep}".join(schema_list))
     except Exception as e:
-        print(str(e))
+        click.echo(str(e))
 
 
 def _generate_json_schema_output_for_plugins(
@@ -54,10 +54,9 @@ def _generate_json_schema_output_for_plugins(
     if len(filtered_plugins) == 0:
         if include_plugins:
             raise ValueError(f"No plugin found with id in {include_plugins}")
-        elif exclude_plugins:
+        if exclude_plugins:
             raise ValueError(f"No plugin found when excluding {exclude_plugins}")
-        else:
-            raise ValueError("No plugin found")
+        raise ValueError("No plugin found")
 
     results = []
     for plugin in filtered_plugins:

databao_context_engine/init_project.py

@@ -1,17 +1,39 @@
 from pathlib import Path
 
 from databao_context_engine.databao_context_project_manager import DatabaoContextProjectManager
-from databao_context_engine.project.init_project import init_project_dir
+from databao_context_engine.project.init_project import InitProjectError, init_project_dir
 from databao_context_engine.project.layout import is_project_dir_valid
 
 
 def init_dce_project(project_dir: Path) -> DatabaoContextProjectManager:
-    initialised_project_dir = init_project_dir(project_dir=project_dir)
+    """Initialize a Databao Context project in the given directory.
 
-    return DatabaoContextProjectManager(project_dir=initialised_project_dir)
+    Args:
+        project_dir: The directory where the project should be initialized.
+
+    Returns:
+        A DatabaoContextProjectManager for the initialized project.
+
+    Raises:
+        InitProjectError: If the project could not be initialized.
+    """
+    try:
+        initialized_project_dir = init_project_dir(project_dir=project_dir)
+
+        return DatabaoContextProjectManager(project_dir=initialized_project_dir)
+    except InitProjectError as e:
+        raise e
 
 
 def init_or_get_dce_project(project_dir: Path) -> DatabaoContextProjectManager:
+    """Initialize a Databao Context project in the given directory or get a DatabaoContextProjectManager for the project if it already exists.
+
+    Args:
+        project_dir: The directory where the project should be initialized if it doesn't exist yet..
+
+    Returns:
+        A DatabaoContextProjectManager for the project in project_dir.
+    """
     if is_project_dir_valid(project_dir):
         return DatabaoContextProjectManager(project_dir=project_dir)