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

ai/chronon/cli/compile/display/class_tracker.py

@@ -0,0 +1,102 @@
+import difflib
+from typing import Any, Dict, List
+
+from rich.text import Text
+
+from ai.chronon.cli.compile.display.compiled_obj import CompiledObj
+from ai.chronon.cli.compile.display.diff_result import DiffResult
+
+
+class ClassTracker:
+    """
+    Tracker object per class - Join, StagingQuery, GroupBy etc
+    """
+
+    def __init__(self):
+        self.existing_objs: Dict[str, CompiledObj] = {}  # name to obj
+        self.files_to_obj: Dict[str, List[Any]] = {}
+        self.files_to_errors: Dict[str, List[Exception]] = {}
+        self.new_objs: Dict[str, CompiledObj] = {}  # name to obj
+        self.diff_result = DiffResult()
+        self.deleted_names: List[str] = []
+
+    def add_existing(self, obj: CompiledObj) -> None:
+        self.existing_objs[obj.name] = obj
+
+    def add(self, compiled: CompiledObj) -> None:
+        if compiled.errors:
+            if compiled.file not in self.files_to_errors:
+                self.files_to_errors[compiled.file] = []
+
+            self.files_to_errors[compiled.file].extend(compiled.errors)
+
+        else:
+            if compiled.file not in self.files_to_obj:
+                self.files_to_obj[compiled.file] = []
+
+            self.files_to_obj[compiled.file].append(compiled.obj)
+
+            self.new_objs[compiled.name] = compiled
+            self._update_diff(compiled)
+
+    def _update_diff(self, compiled: CompiledObj) -> None:
+        if compiled.name in self.existing_objs:
+            existing_json = self.existing_objs[compiled.name].tjson
+            new_json = compiled.tjson
+
+            if existing_json != new_json:
+                diff = difflib.unified_diff(
+                    existing_json.splitlines(keepends=True),
+                    new_json.splitlines(keepends=True),
+                    n=2,
+                )
+
+                print(f"Updated object: {compiled.name} in file {compiled.file}")
+                print("".join(diff))
+                print("\n")
+
+                self.diff_result.updated.append(compiled.name)
+
+        else:
+            if not compiled.errors:
+                self.diff_result.added.append(compiled.name)
+
+    def close(self) -> None:
+        self.closed = True
+        self.recent_file = None
+        self.deleted_names = list(self.existing_objs.keys() - self.new_objs.keys())
+
+    def to_status(self) -> Text:
+        text = Text(overflow="fold", no_wrap=False)
+
+        if self.existing_objs:
+            text.append(f"  Parsed {len(self.existing_objs)} previously compiled objects.\n")
+
+        if self.files_to_errors:
+            text.append("  Failed to compile ")
+            text.append(f"{len(self.files_to_errors)} ", style="red")
+            text.append("files.\n")
+
+        return text
+
+    def to_errors(self) -> Text:
+        text = Text(overflow="fold", no_wrap=False)
+
+        if self.files_to_errors:
+            for file, errors in self.files_to_errors.items():
+                text.append("  ERROR ", style="bold red")
+                text.append(f"- {file}:\n")
+
+                for error in errors:
+                    # Format each error properly, handling newlines
+                    error_msg = str(error)
+                    text.append(f"    {error_msg}\n", style="red")
+
+        return text
+
+    # doesn't make sense to show deletes until the very end of compilation
+    def diff(self, ignore_python_errors: bool = False) -> Text:
+        # Don't show diff if there are compile errors - it's confusing
+        if self.files_to_errors and not ignore_python_errors:
+            return Text("\n❗Please fix python errors then retry compilation.\n", style="dim cyan")
+        return self.diff_result.render(deleted_names=self.deleted_names)

ai/chronon/cli/compile/display/compile_status.py

@@ -0,0 +1,95 @@
+from collections import OrderedDict
+from typing import Dict
+
+from rich.live import Live
+from rich.text import Text
+
+from ai.chronon.cli.compile.display.class_tracker import ClassTracker
+from ai.chronon.cli.compile.display.compiled_obj import CompiledObj
+
+
+class CompileStatus:
+    """
+    Uses rich ui - to consolidate and sink the overview of the compile process to the bottom.
+    """
+
+    def __init__(self, use_live: bool = False):
+        self.cls_to_tracker: Dict[str, ClassTracker] = OrderedDict()
+        self.use_live = use_live
+        # we need vertical_overflow to be visible as the output gets cufoff when our output goes past the termianal window
+        # but then we start seeing duplicates: https://github.com/Textualize/rich/issues/3263
+        if self.use_live:
+            self.live = Live(refresh_per_second=50, vertical_overflow="visible")
+            self.live.start()
+
+    def print_live_console(self, msg: str):
+        if self.use_live:
+            self.live.console.print(msg)
+
+    def add_object_update_display(self, compiled: CompiledObj, obj_type: str = None) -> None:
+        if compiled.obj_type is not None and obj_type is not None:
+            assert compiled.obj_type == obj_type, (
+                f"obj_type mismatch: {compiled.obj_type} != {obj_type}"
+            )
+
+        if obj_type not in self.cls_to_tracker:
+            self.cls_to_tracker[obj_type] = ClassTracker()
+
+        self.cls_to_tracker[obj_type].add(compiled)
+
+        self._update_display()
+
+    def add_existing_object_update_display(self, existing_obj: CompiledObj) -> None:
+        obj_type = existing_obj.obj_type
+
+        if obj_type not in self.cls_to_tracker:
+            self.cls_to_tracker[obj_type] = ClassTracker()
+
+        self.cls_to_tracker[obj_type].add_existing(existing_obj)
+
+        self._update_display()
+
+    def close_cls(self, obj_type: str) -> None:
+        if obj_type in self.cls_to_tracker:
+            self.cls_to_tracker[obj_type].close()
+            self._update_display()
+
+    def close(self) -> None:
+        self._update_display()
+        if self.use_live:
+            self.live.stop()
+
+    def render(self, ignore_python_errors: bool = False) -> Text:
+        text = Text(overflow="fold", no_wrap=False)
+
+        for obj_type, tracker in self.cls_to_tracker.items():
+            # Skip MetaData section
+            if obj_type == "MetaData":
+                continue
+
+            text.append(f"\n{obj_type}-s:\n", style="cyan")
+
+            status = tracker.to_status()
+            if status:
+                text.append(status)
+
+            errors = tracker.to_errors()
+            if errors:
+                text.append(errors)
+
+            diff = tracker.diff(ignore_python_errors)
+            if diff:
+                text.append(diff)
+
+        text.append("\n")
+        return text
+
+    def _update_display(self):
+        # self.live.clear()
+
+        # TODO: add this after live_crop is implemented
+        # text = self.display_text()
+        # if self.use_live:
+        #     self.live.update(text, refresh=True)
+        # return text
+        pass

ai/chronon/cli/compile/display/compiled_obj.py

@@ -0,0 +1,12 @@
+from dataclasses import dataclass
+from typing import Any, List, Optional
+
+
+@dataclass
+class CompiledObj:
+    name: str
+    obj: Any
+    file: str
+    errors: Optional[List[Exception]]
+    obj_type: str
+    tjson: str

ai/chronon/cli/compile/display/console.py

@@ -0,0 +1,3 @@
+from rich.console import Console
+
+console = Console()

ai/chronon/cli/compile/display/diff_result.py

@@ -0,0 +1,111 @@
+from typing import List, Tuple
+
+from rich.text import Text
+
+from ai.chronon.cli.compile.version_utils import parse_name_and_version
+
+
+class DiffResult:
+    def __init__(self):
+        self.added: List[str] = []
+        self.updated: List[str] = []
+        self.version_bumped: List[
+            Tuple[str, int, int]
+        ] = []  # (base_name, old_version, new_version)
+
+    def detect_version_changes(self, deleted_names: List[str]) -> Tuple[List[str], List[str]]:
+        """Detect version changes between deleted and added names.
+
+        Returns:
+            Tuple of (remaining_deleted, remaining_added) after removing version changes
+        """
+        # Create mappings of base names to versions
+        deleted_base_to_name = {}
+        added_base_to_name = {}
+
+        for name in deleted_names:
+            base_name, version = parse_name_and_version(name)
+            if version is not None:
+                deleted_base_to_name[base_name] = (name, version)
+
+        for name in self.added:
+            base_name, version = parse_name_and_version(name)
+            if version is not None:
+                added_base_to_name[base_name] = (name, version)
+
+        # Find version changes
+        remaining_deleted = []
+        remaining_added = []
+
+        for name in deleted_names:
+            base_name, old_version = parse_name_and_version(name)
+            if (
+                base_name in added_base_to_name
+                and old_version is not None
+                and base_name in deleted_base_to_name
+            ):
+                # This is a version change
+                _, new_version = added_base_to_name[base_name]
+                self.version_bumped.append((base_name, old_version, new_version))
+            else:
+                remaining_deleted.append(name)
+
+        for name in self.added:
+            base_name, version = parse_name_and_version(name)
+            if not (base_name in deleted_base_to_name and version is not None):
+                # This is not part of a version change
+                remaining_added.append(name)
+
+        return remaining_deleted, remaining_added
+
+    def render(self, deleted_names: List[str], indent="  ") -> Text:
+        # Detect version changes first
+        remaining_deleted, remaining_added = self.detect_version_changes(deleted_names)
+
+        def added_signage():
+            return Text("➕ Added", style="dim green")
+
+        def updated_signage():
+            return Text("❗ Changed in place (no version change)", style="dim yellow")
+
+        def deleted_signage():
+            return Text("🗑️ Deleted", style="red")
+
+        def version_bumped_signage():
+            return Text("⬆️ Version changed", style="dim blue")
+
+        added = [(added_signage(), name) for name in remaining_added]
+        updated = [(updated_signage(), name) for name in self.updated]
+        version_bumped = [
+            (version_bumped_signage(), f"{base_name} (v{old_ver} -> v{new_ver})")
+            for base_name, old_ver, new_ver in self.version_bumped
+        ]
+
+        # Put version changes and additions first, changed items at the bottom
+        # Sort each group separately to maintain grouping
+        version_bumped_sorted = sorted(version_bumped, key=lambda t: t[1])
+        added_sorted = sorted(added, key=lambda t: t[1])
+        updated_sorted = sorted(updated, key=lambda t: t[1])
+
+        result_order = version_bumped_sorted + added_sorted
+
+        if remaining_deleted:
+            deleted = [(deleted_signage(), name) for name in remaining_deleted]
+            deleted_sorted = sorted(deleted, key=lambda t: t[1])
+            result_order += deleted_sorted
+
+        # Add updated (changed in place) at the very end
+        result_order += updated_sorted
+
+        text = Text(overflow="fold", no_wrap=False)
+        for signage, name in result_order:
+            text.append(indent)
+            text.append(signage)
+            text.append(" ")
+            text.append(name)
+            text.append("\n")
+
+        if not text:
+            return Text(indent + "No new changes detected\n", style="dim")
+
+        return text

ai/chronon/cli/compile/fill_templates.py

@@ -0,0 +1,35 @@
+from gen_thrift.api.ttypes import Join, Team
+
+from ai.chronon import utils
+from ai.chronon.cli.compile.compile_context import CompileContext
+
+
+def _fill_template(table, obj, namespace):
+    if table:
+        table = table.replace("{{ logged_table }}", utils.log_table_name(obj, full_name=True))
+        table = table.replace("{{ db }}", namespace)
+
+    return table
+
+
+def set_templated_values(obj, cls, compile_context: CompileContext):
+    team_obj: Team = compile_context.teams_dict[obj.team]
+    namespace = team_obj.outputNamespace
+
+    if cls == Join and obj.bootstrapParts:
+        for bootstrap in obj.bootstrapParts:
+            bootstrap.table = _fill_template(bootstrap.table, obj, namespace)
+
+        if obj.metaData.dependencies:
+            obj.metaData.dependencies = [
+                _fill_template(dep, obj, namespace) for dep in obj.metaData.dependencies
+            ]
+
+    if cls == Join and obj.labelParts:
+        obj.labelParts.metaData.dependencies = [
+            label_dep.replace(
+                "{{ join_backfill_table }}",
+                utils.output_table_name(obj, full_name=True),
+            )
+            for label_dep in obj.labelParts.metaData.dependencies
+        ]

ai/chronon/cli/compile/parse_configs.py

@@ -0,0 +1,134 @@
+import copy
+import glob
+import importlib
+import os
+from typing import Any, List
+
+from gen_thrift.api.ttypes import GroupBy, Join
+
+from ai.chronon import airflow_helpers
+from ai.chronon.cli.compile import parse_teams, serializer
+from ai.chronon.cli.compile.compile_context import CompileContext
+from ai.chronon.cli.compile.display.compiled_obj import CompiledObj
+from ai.chronon.cli.logger import get_logger
+
+logger = get_logger()
+
+
+def from_folder(cls: type, input_dir: str, compile_context: CompileContext) -> List[CompiledObj]:
+    """
+    Recursively consumes a folder, and constructs a map of
+    object qualifier to StagingQuery, GroupBy, or Join
+    """
+
+    python_files = glob.glob(os.path.join(input_dir, "**/*.py"), recursive=True)
+
+    results = []
+
+    for f in python_files:
+        try:
+            results_dict = from_file(f, cls, input_dir)
+
+            for name, obj in results_dict.items():
+                parse_teams.update_metadata(obj, compile_context.teams_dict)
+                # Populate columnHashes field with semantic hashes
+                populate_column_hashes(obj)
+
+                # Airflow deps must be set AFTER updating metadata
+                airflow_helpers.set_airflow_deps(obj)
+
+                obj.metaData.sourceFile = os.path.relpath(f, compile_context.chronon_root)
+
+                tjson = serializer.thrift_simple_json(obj)
+
+                # Perform validation
+                errors = compile_context.validator.validate_obj(obj)
+
+                result = CompiledObj(
+                    name=name,
+                    obj=obj,
+                    file=f,
+                    errors=errors if len(errors) > 0 else None,
+                    obj_type=cls.__name__,
+                    tjson=tjson,
+                )
+                results.append(result)
+
+                compile_context.compile_status.add_object_update_display(result, cls.__name__)
+
+        except Exception as e:
+            result = CompiledObj(
+                name=None,
+                obj=None,
+                file=f,
+                errors=[e],
+                obj_type=cls.__name__,
+                tjson=None,
+            )
+
+            results.append(result)
+
+            compile_context.compile_status.add_object_update_display(result, cls.__name__)
+
+    return results
+
+
+def from_file(file_path: str, cls: type, input_dir: str):
+    # this is where the python path should have been set to
+    chronon_root = os.path.dirname(input_dir)
+    rel_path = os.path.relpath(file_path, chronon_root)
+
+    rel_path_without_extension = os.path.splitext(rel_path)[0]
+
+    module_name = rel_path_without_extension.replace("/", ".")
+
+    conf_type, team_name_with_path = module_name.split(".", 1)
+    mod_path = team_name_with_path.replace("/", ".")
+
+    module = importlib.import_module(module_name)
+
+    result = {}
+
+    for var_name, obj in list(module.__dict__.items()):
+        if isinstance(obj, cls):
+            copied_obj = copy.deepcopy(obj)
+
+            name = f"{mod_path}.{var_name}"
+
+            # Add version suffix if version is set
+            name = name + "__" + str(copied_obj.metaData.version)
+
+            copied_obj.metaData.name = name
+            copied_obj.metaData.team = mod_path.split(".")[0]
+
+            result[name] = copied_obj
+
+    return result
+
+
+def populate_column_hashes(obj: Any):
+    """
+    Populate the columnHashes field in the object's metadata with semantic hashes
+    for each output column.
+    """
+    # Import here to avoid circular imports
+    from ai.chronon.cli.compile.column_hashing import (
+        compute_group_by_columns_hashes,
+        compute_join_column_hashes,
+    )
+
+    if isinstance(obj, GroupBy):
+        # For GroupBy objects, get column hashes
+        column_hashes = compute_group_by_columns_hashes(obj, exclude_keys=False)
+        obj.metaData.columnHashes = column_hashes
+
+    elif isinstance(obj, Join):
+        # For Join objects, get column hashes
+        column_hashes = compute_join_column_hashes(obj)
+        obj.metaData.columnHashes = column_hashes
+
+        if obj.joinParts:
+            for jp in obj.joinParts or []:
+                group_by = jp.groupBy
+                group_by_hashes = compute_group_by_columns_hashes(group_by)
+                group_by.metaData.columnHashes = group_by_hashes