awx-zipline-ai 0.0.32__py3-none-any.whl

ai/chronon/cli/compile/compile_context.py

@@ -0,0 +1,173 @@
+import os
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Type
+
+from gen_thrift.api.ttypes import ConfType, GroupBy, Join, MetaData, Model, StagingQuery, Team
+
+import ai.chronon.cli.compile.parse_teams as teams
+from ai.chronon.cli.compile.conf_validator import ConfValidator
+from ai.chronon.cli.compile.display.compile_status import CompileStatus
+from ai.chronon.cli.compile.display.compiled_obj import CompiledObj
+from ai.chronon.cli.compile.serializer import file2thrift
+from ai.chronon.cli.logger import get_logger, require
+
+logger = get_logger()
+
+
+@dataclass
+class ConfigInfo:
+    folder_name: str
+    cls: Type
+    config_type: Optional[ConfType]
+
+
+@dataclass
+class CompileContext:
+    def __init__(self, ignore_python_errors: bool = False):
+        self.chronon_root: str = os.getenv("CHRONON_ROOT", os.getcwd())
+        self.teams_dict: Dict[str, Team] = teams.load_teams(self.chronon_root)
+        self.compile_dir: str = "compiled"
+        self.ignore_python_errors: bool = ignore_python_errors
+
+        self.config_infos: List[ConfigInfo] = [
+            ConfigInfo(folder_name="joins", cls=Join, config_type=ConfType.JOIN),
+            ConfigInfo(
+                folder_name="group_bys",
+                cls=GroupBy,
+                config_type=ConfType.GROUP_BY,
+            ),
+            ConfigInfo(
+                folder_name="staging_queries",
+                cls=StagingQuery,
+                config_type=ConfType.STAGING_QUERY,
+            ),
+            ConfigInfo(folder_name="models", cls=Model, config_type=ConfType.MODEL),
+            ConfigInfo(
+                folder_name="teams_metadata", cls=MetaData, config_type=None
+            ),  # only for team metadata
+        ]
+
+        self.compile_status = CompileStatus(use_live=False)
+
+        self.existing_confs: Dict[Type, Dict[str, Any]] = {}
+        for config_info in self.config_infos:
+            cls = config_info.cls
+            self.existing_confs[cls] = self._parse_existing_confs(cls)
+
+        self.validator: ConfValidator = ConfValidator(
+            input_root=self.chronon_root,
+            output_root=self.compile_dir,
+            existing_gbs=self.existing_confs[GroupBy],
+            existing_joins=self.existing_confs[Join],
+            existing_staging_queries=self.existing_confs[StagingQuery],
+        )
+
+    def input_dir(self, cls: type) -> str:
+        """
+        - eg., input: group_by class
+        - eg., output: root/group_bys/
+        """
+        config_info = self.config_info_for_class(cls)
+        return os.path.join(self.chronon_root, config_info.folder_name)
+
+    def staging_output_dir(self, cls: type = None) -> str:
+        """
+        - eg., input: group_by class
+        - eg., output: root/compiled_staging/group_bys/
+        """
+        if cls is None:
+            return os.path.join(self.chronon_root, self.compile_dir + "_staging")
+        else:
+            config_info = self.config_info_for_class(cls)
+            return os.path.join(
+                self.chronon_root,
+                self.compile_dir + "_staging",
+                config_info.folder_name,
+            )
+
+    def output_dir(self, cls: type = None) -> str:
+        """
+        - eg., input: group_by class
+        - eg., output: root/compiled/group_bys/
+        """
+        if cls is None:
+            return os.path.join(self.chronon_root, self.compile_dir)
+        else:
+            config_info = self.config_info_for_class(cls)
+            return os.path.join(self.chronon_root, self.compile_dir, config_info.folder_name)
+
+    def staging_output_path(self, compiled_obj: CompiledObj):
+        """
+        - eg., input: group_by with name search.clicks.features.v1
+        - eg., output: root/compiled_staging/group_bys/search/clicks.features.v1
+        """
+
+        output_dir = self.staging_output_dir(compiled_obj.obj.__class__)  # compiled/joins
+
+        team, rest = compiled_obj.name.split(".", 1)  # search, clicks.features.v1
+
+        return os.path.join(
+            output_dir,
+            team,
+            rest,
+        )
+
+    def config_info_for_class(self, cls: type) -> ConfigInfo:
+        for info in self.config_infos:
+            if info.cls == cls:
+                return info
+
+        require(False, f"Class {cls} not found in CONFIG_INFOS")
+
+    def _parse_existing_confs(self, obj_class: type) -> Dict[str, object]:
+        result = {}
+
+        output_dir = self.output_dir(obj_class)
+
+        # Check if output_dir exists before walking
+        if not os.path.exists(output_dir):
+            return result
+
+        for sub_root, _sub_dirs, sub_files in os.walk(output_dir):
+            for f in sub_files:
+                if f.startswith("."):  # ignore hidden files - such as .DS_Store
+                    continue
+
+                full_path = os.path.join(sub_root, f)
+
+                try:
+                    obj = file2thrift(full_path, obj_class)
+
+                    if obj:
+                        if hasattr(obj, "metaData"):
+                            result[obj.metaData.name] = obj
+                            compiled_obj = CompiledObj(
+                                name=obj.metaData.name,
+                                obj=obj,
+                                file=obj.metaData.sourceFile,
+                                errors=None,
+                                obj_type=obj_class.__name__,
+                                tjson=open(full_path).read(),
+                            )
+                            self.compile_status.add_existing_object_update_display(compiled_obj)
+                        elif isinstance(obj, MetaData):
+                            team_metadata_name = ".".join(
+                                full_path.split("/")[-2:]
+                            )  # use the name of the file as team metadata won't have name
+                            result[team_metadata_name] = obj
+                            compiled_obj = CompiledObj(
+                                name=team_metadata_name,
+                                obj=obj,
+                                file=obj.sourceFile,
+                                errors=None,
+                                obj_type=obj_class.__name__,
+                                tjson=open(full_path).read(),
+                            )
+                            self.compile_status.add_existing_object_update_display(compiled_obj)
+                    else:
+                        logger.errors(f"Parsed object from {full_path} has no metaData attribute")
+
+                except Exception as e:
+                    print(f"Failed to parse file {full_path}: {str(e)}", e)
+
+        return result

ai/chronon/cli/compile/compiler.py

@@ -0,0 +1,183 @@
+import os
+import shutil
+import traceback
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Tuple
+
+from gen_thrift.api.ttypes import ConfType
+
+import ai.chronon.cli.compile.display.compiled_obj
+import ai.chronon.cli.compile.parse_configs as parser
+import ai.chronon.cli.logger as logger
+from ai.chronon.cli.compile import serializer
+from ai.chronon.cli.compile.compile_context import CompileContext, ConfigInfo
+from ai.chronon.cli.compile.display.compiled_obj import CompiledObj
+from ai.chronon.cli.compile.display.console import console
+from ai.chronon.cli.compile.parse_teams import merge_team_execution_info
+from ai.chronon.types import MetaData
+
+logger = logger.get_logger()
+
+
+@dataclass
+class CompileResult:
+    config_info: ConfigInfo
+    obj_dict: Dict[str, Any]
+    error_dict: Dict[str, List[BaseException]]
+
+
+class Compiler:
+    def __init__(self, compile_context: CompileContext):
+        self.compile_context = compile_context
+
+    def compile(self) -> Dict[ConfType, CompileResult]:
+        # Clean staging directory at the start to ensure fresh compilation
+        staging_dir = self.compile_context.staging_output_dir()
+        if os.path.exists(staging_dir):
+            shutil.rmtree(staging_dir)
+
+        config_infos = self.compile_context.config_infos
+
+        compile_results = {}
+        all_compiled_objects = []  # Collect all compiled objects for change validation
+
+        for config_info in config_infos:
+            configs, compiled_objects = self._compile_class_configs(config_info)
+            compile_results[config_info.config_type] = configs
+
+            # Collect compiled objects for change validation
+            all_compiled_objects.extend(compiled_objects)
+
+        # Validate changes once after all classes have been processed
+        self.compile_context.validator.validate_changes(all_compiled_objects)
+
+        # Show the nice display first
+        console.print(
+            self.compile_context.compile_status.render(self.compile_context.ignore_python_errors)
+        )
+
+        # Check for confirmation before finalizing files
+        self.compile_context.validator.check_pending_changes_confirmation(
+            self.compile_context.compile_status
+        )
+
+        # Only proceed with file operations if there are no compilation errors
+        if not self._has_compilation_errors() or self.compile_context.ignore_python_errors:
+            self._compile_team_metadata()
+
+            # check if staging_output_dir exists
+            staging_dir = self.compile_context.staging_output_dir()
+            if os.path.exists(staging_dir):
+                # replace staging_output_dir to output_dir
+                output_dir = self.compile_context.output_dir()
+                if os.path.exists(output_dir):
+                    shutil.rmtree(output_dir)
+                shutil.move(staging_dir, output_dir)
+            else:
+                print(
+                    f"Staging directory {staging_dir} does not exist. "
+                    "Happens when every chronon config fails to compile or when no chronon configs exist."
+                )
+        else:
+            # Clean up staging directory when there are errors (don't move to output)
+            staging_dir = self.compile_context.staging_output_dir()
+            if os.path.exists(staging_dir):
+                shutil.rmtree(staging_dir)
+
+        return compile_results
+
+    def _has_compilation_errors(self):
+        """Check if there are any compilation errors across all class trackers."""
+        for tracker in self.compile_context.compile_status.cls_to_tracker.values():
+            if tracker.files_to_errors:
+                return True
+        return False
+
+    def _compile_team_metadata(self):
+        """
+        Compile the team metadata and return the compiled object.
+        """
+        teams_dict = self.compile_context.teams_dict
+        for team in teams_dict:
+            m = MetaData()
+            merge_team_execution_info(m, teams_dict, team)
+
+            tjson = serializer.thrift_simple_json(m)
+            name = f"{team}.{team}_team_metadata"
+            result = CompiledObj(
+                name=name,
+                obj=m,
+                file=name,
+                errors=None,
+                obj_type=MetaData.__name__,
+                tjson=tjson,
+            )
+            self._write_object(result)
+            self.compile_context.compile_status.add_object_update_display(result, MetaData.__name__)
+
+        # Done writing team metadata, close the class
+        self.compile_context.compile_status.close_cls(MetaData.__name__)
+
+    def _compile_class_configs(
+        self, config_info: ConfigInfo
+    ) -> Tuple[CompileResult, List[CompiledObj]]:
+        compile_result = CompileResult(config_info=config_info, obj_dict={}, error_dict={})
+
+        input_dir = self.compile_context.input_dir(config_info.cls)
+
+        compiled_objects = parser.from_folder(config_info.cls, input_dir, self.compile_context)
+
+        objects, errors = self._write_objects_in_folder(compiled_objects)
+
+        if objects:
+            compile_result.obj_dict.update(objects)
+
+        if errors:
+            compile_result.error_dict.update(errors)
+
+        self.compile_context.compile_status.close_cls(config_info.cls.__name__)
+
+        return compile_result, compiled_objects
+
+    def _write_objects_in_folder(
+        self,
+        compiled_objects: List[ai.chronon.cli.compile.display.compiled_obj.CompiledObj],
+    ) -> Tuple[Dict[str, Any], Dict[str, List[BaseException]]]:
+        error_dict = {}
+        object_dict = {}
+
+        for co in compiled_objects:
+            if co.obj:
+                if co.errors:
+                    error_dict[co.name] = co.errors
+
+                    for error in co.errors:
+                        self.compile_context.compile_status.print_live_console(
+                            f"Error processing conf {co.name}: {error}"
+                        )
+                        traceback.print_exception(type(error), error, error.__traceback__)
+
+                else:
+                    self._write_object(co)
+                    object_dict[co.name] = co.obj
+            else:
+                error_dict[co.file] = co.errors
+
+                self.compile_context.compile_status.print_live_console(
+                    f"Error processing file {co.file}: {co.errors}"
+                )
+                for error in co.errors:
+                    traceback.print_exception(type(error), error, error.__traceback__)
+
+        return object_dict, error_dict
+
+    def _write_object(self, compiled_obj: CompiledObj) -> Optional[List[BaseException]]:
+        output_path = self.compile_context.staging_output_path(compiled_obj)
+
+        folder = os.path.dirname(output_path)
+
+        if not os.path.exists(folder):
+            os.makedirs(folder)
+
+        with open(output_path, "w") as f:
+            f.write(compiled_obj.tjson)