150 - Support for yaml extension in project conf files (#175)

* Support for yaml extension in project conf files

* Tests and fixes

* Project tests

* Refactored tests using pytest parametrize decorator

brickflow/__init__.py

@@ -43,9 +43,22 @@ def _insert_before_path_startswith(
     arr.append(new_element)
 
 
+class ConfigFileType(Enum):
+    YAML = "yaml"
+    YML = "yml"
+
+
 class BrickflowProjectConstants(Enum):
-    DEFAULT_MULTI_PROJECT_ROOT_FILE_NAME = ".brickflow-project-root.yml"
-    DEFAULT_MULTI_PROJECT_CONFIG_FILE_NAME = "brickflow-multi-project.yml"
+    DEFAULT_MULTI_PROJECT_ROOT_FILE_NAME = ".brickflow-project-root"
+    DEFAULT_MULTI_PROJECT_CONFIG_FILE_NAME = "brickflow-multi-project"
+    DEFAULT_CONFIG_FILE_TYPE = ConfigFileType.YML.value
+
+
+def get_config_file_type(brickflow_root: str) -> ConfigFileType:
+    for config_file_type in ConfigFileType:
+        if brickflow_root.endswith(config_file_type.value):
+            return config_file_type
+    return ConfigFileType.YAML
 
 
 class BrickflowEnvVars(Enum):
@@ -339,6 +352,8 @@ def get_bundles_project_env() -> str:
     "BrickflowDefaultEnvs",
     "get_default_log_handler",
     "get_brickflow_version",
+    "ConfigFileType",
+    "get_config_file_type",
     "BrickflowProjectConstants",
     "BrickflowProjectDeploymentSettings",
 ]

brickflow/cli/configure.py

@@ -102,7 +102,10 @@ def create_entry_point(working_dir: str, data: str) -> None:
 
 
 def create_brickflow_project_root_marker() -> None:
-    path = Path(BrickflowProjectConstants.DEFAULT_MULTI_PROJECT_ROOT_FILE_NAME.value)
+    path = Path(
+        f"{BrickflowProjectConstants.DEFAULT_MULTI_PROJECT_ROOT_FILE_NAME.value}."
+        f"{BrickflowProjectConstants.DEFAULT_CONFIG_FILE_TYPE.value}"
+    )
     if path.exists():
         click.echo(f"Path: {str(path.absolute())} already exists...")
         # path = Path(working_dir) / "entrypoint.py.new"

brickflow/cli/projects.py

@@ -1,6 +1,5 @@
 import contextlib
 import os
-from enum import Enum
 from pathlib import Path
 from typing import Dict, Optional, List, Generator, Any, Callable
 
@@ -15,6 +14,8 @@
     _ilog,
     BrickflowProjectDeploymentSettings,
     ctx,
+    ConfigFileType,
+    get_config_file_type,
 )
 from brickflow.cli.bundles import (
     bundle_deploy,
@@ -83,11 +84,6 @@ def has_projects(self) -> bool:
         return self.project_roots is not None and len(self.project_roots) > 0
 
 
-class ConfigFileType(Enum):
-    YAML = "yaml"
-    JSON = "json"  # unsupported
-
-
 class MultiProjectManager:
     def __init__(
         self,
@@ -126,8 +122,15 @@ def _load_config(self) -> BrickflowMultiRootProjectConfig:
                 return config
             return BrickflowMultiRootProjectConfig(project_roots={})
 
-    def _root_config_path(self, root: str) -> Path:
-        root_file = BrickflowProjectConstants.DEFAULT_MULTI_PROJECT_ROOT_FILE_NAME.value
+    def _root_config_path(
+        self,
+        root: str,
+        config_file_type: ConfigFileType = BrickflowProjectConstants.DEFAULT_CONFIG_FILE_TYPE,
+    ) -> Path:
+        root_file = (
+            f"{BrickflowProjectConstants.DEFAULT_MULTI_PROJECT_ROOT_FILE_NAME.value}."
+            f"{config_file_type.value}"
+        )
         return self._config_file.parent / root / root_file
 
     def _load_roots(self) -> Dict[str, BrickflowRootProjectConfig]:
@@ -141,7 +144,9 @@ def _load_roots(self) -> Dict[str, BrickflowRootProjectConfig]:
         )
         root_dict = {}
         for root in roots:
-            with self._root_config_path(root).open("r", encoding="utf-8") as f:
+            with self._root_config_path(root, config_file_type=self.file_type).open(
+                "r", encoding="utf-8"
+            ) as f:
                 root_dict[root] = BrickflowRootProjectConfig.parse_obj(
                     yaml.safe_load(f.read())
                 )
@@ -233,22 +238,32 @@ class BrickflowRootNotFound(Exception):
 
 def get_brickflow_root(current_path: Optional[Path] = None) -> Path:
     current_dir = Path(current_path or get_notebook_ws_path(ctx.dbutils) or os.getcwd())
-    potential_config_file_path = (
-        current_dir
-        / BrickflowProjectConstants.DEFAULT_MULTI_PROJECT_CONFIG_FILE_NAME.value
-    )
-    if potential_config_file_path.exists():
-        return potential_config_file_path
-    elif current_dir.parent == current_dir:
+
+    potential_config_files = [
+        f"{BrickflowProjectConstants.DEFAULT_MULTI_PROJECT_CONFIG_FILE_NAME.value}.{cfg_type.value}"
+        for cfg_type in ConfigFileType
+    ]
+    potential_config_file_paths = [current_dir / p for p in potential_config_files]
+
+    for potential_config_file_path in potential_config_file_paths:
+        if potential_config_file_path.exists():
+            return potential_config_file_path
+
+    if current_dir.parent == current_dir:
         # Reached the filesystem root, return just raw file value
         return Path(
-            BrickflowProjectConstants.DEFAULT_MULTI_PROJECT_CONFIG_FILE_NAME.value
+            f"{BrickflowProjectConstants.DEFAULT_MULTI_PROJECT_CONFIG_FILE_NAME.value}."
+            f"{BrickflowProjectConstants.DEFAULT_CONFIG_FILE_TYPE.value}"
         )
     else:
         return get_brickflow_root(current_dir.parent)
 
 
-multi_project_manager = MultiProjectManager(config_file_name=str(get_brickflow_root()))
+brickflow_root_path = get_brickflow_root()
+cfg_file_type = get_config_file_type(str(brickflow_root_path))
+multi_project_manager = MultiProjectManager(
+    config_file_name=str(brickflow_root_path), file_type=cfg_file_type
+)
 
 
 def initialize_project_entrypoint(

brickflow/resolver/__init__.py

@@ -37,8 +37,10 @@ def go_up_till_brickflow_root(cur_path: str) -> str:
     path = pathlib.Path(cur_path).resolve()
 
     valid_roots = [
-        BrickflowProjectConstants.DEFAULT_MULTI_PROJECT_ROOT_FILE_NAME.value,
-        BrickflowProjectConstants.DEFAULT_MULTI_PROJECT_CONFIG_FILE_NAME.value,
+        f"{BrickflowProjectConstants.DEFAULT_MULTI_PROJECT_ROOT_FILE_NAME.value}."
+        f"{BrickflowProjectConstants.DEFAULT_CONFIG_FILE_TYPE.value}",
+        f"{BrickflowProjectConstants.DEFAULT_MULTI_PROJECT_CONFIG_FILE_NAME.value}."
+        f"{BrickflowProjectConstants.DEFAULT_CONFIG_FILE_TYPE.value}",
     ]
 
     # recurse to see if there is a brickflow root and return the path

tests/cli/sample_yaml_project/.brickflow-project-root.yaml

@@ -0,0 +1,9 @@
+version: v1
+projects:
+  test_cli_project:
+    name: test_cli_project
+    brickflow_version: 1.2.1
+    deployment_mode: bundle
+    enable_plugins: false
+    path_from_repo_root_to_project_root: some/test/path
+    path_project_root_to_workflows_dir:  path/to/workflows

tests/cli/sample_yaml_project/brickflow-multi-project.yaml

@@ -0,0 +1,4 @@
+version: v1
+project_roots:
+  test_cli_project:
+    root_yaml_rel_path: .

tests/cli/sample_yml_project/.brickflow-project-root.yml

@@ -0,0 +1,9 @@
+version: v1
+projects:
+  test_cli_project:
+    name: test_cli_project
+    brickflow_version: 1.2.1
+    deployment_mode: bundle
+    enable_plugins: false
+    path_from_repo_root_to_project_root: some/test/path
+    path_project_root_to_workflows_dir:  path/to/workflows

tests/cli/sample_yml_project/brickflow-multi-project.yml

@@ -0,0 +1,4 @@
+version: v1
+project_roots:
+  test_cli_project:
+    root_yaml_rel_path: .

tests/cli/test_projects.py

@@ -0,0 +1,65 @@
+from pathlib import Path
+import shutil
+import os
+import pytest
+from brickflow import ConfigFileType
+from brickflow.cli.projects import MultiProjectManager, get_brickflow_root
+
+
+@pytest.mark.parametrize(
+    "project_folder,extension",
+    [("sample_yml_project", "yml"), ("sample_yaml_project", "yaml")],
+)
+def test_get_brickflow_root(project_folder, extension):
+    cwd = os.getcwd()
+    test_folder = str(Path(__file__).parent)
+
+    # Creating empty test directories
+    os.makedirs(f"{test_folder}/{project_folder}/some/dummy/dir", exist_ok=True)
+    os.chdir(f"{test_folder}/{project_folder}/some/dummy/dir")
+
+    actual = get_brickflow_root()
+    assert actual == Path(
+        f"{test_folder}/{project_folder}/brickflow-multi-project.{extension}"
+    )
+
+    # Cleanup
+    shutil.rmtree(f"{test_folder}/{project_folder}/some")
+    os.chdir(cwd)
+
+
+@pytest.mark.parametrize(
+    "project_folder, config_type",
+    [
+        ("sample_yml_project", ConfigFileType.YML),
+        ("sample_yaml_project", ConfigFileType.YAML),
+    ],
+)
+def test_multi_project_manager_yaml(project_folder, config_type):
+    cwd = os.getcwd()
+    test_folder = str(Path(__file__).parent)
+    os.chdir(test_folder)
+
+    config_file_name = (
+        f"{test_folder}/{project_folder}/brickflow-multi-project.{config_type.value}"
+    )
+    manager = MultiProjectManager(
+        config_file_name=config_file_name, file_type=config_type
+    )
+    assert manager._brickflow_multi_project_config.version == "v1"
+    expected_project_config = {
+        "version": "v1",
+        "projects": {
+            "test_cli_project": {
+                "name": "test_cli_project",
+                "path_from_repo_root_to_project_root": "some/test/path",
+                "path_project_root_to_workflows_dir": "path/to/workflows",
+                "deployment_mode": "bundle",
+                "brickflow_version": "1.2.1",
+                "enable_plugins": False,
+            }
+        },
+    }
+    assert manager._project_config_dict["."].model_dump() == expected_project_config
+
+    os.chdir(cwd)

tests/test_brickflow.py

@@ -1,4 +1,6 @@
 # pylint: disable=unused-import
+import pytest
+from brickflow import get_config_file_type, ConfigFileType
 
 
 def test_imports():
@@ -38,3 +40,16 @@ def test_imports():
         print("All imports Succeeded")
     except ImportError as e:
         print(f"Import failed: {e}")
+
+
+@pytest.mark.parametrize(
+    "config_file_name,expected_extension",
+    [
+        (".brickflow-project-root.yaml", ConfigFileType.YAML),
+        (".brickflow-project-root.yml", ConfigFileType.YML),
+        (".brickflow-project-root.json", ConfigFileType.YAML),
+    ],
+)
+def test_get_config_type(config_file_name, expected_extension):
+    actual = get_config_file_type(f"some/brickflow/root/{config_file_name}")
+    assert actual == expected_extension