awx-zipline-ai 0.2.1__py3-none-any.whl → 0.3.0__py3-none-any.whl

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

@@ -26,14 +26,11 @@ class CompileStatus:
         if self.use_live:
             self.live.console.print(msg)
 
-    def add_object_update_display(
-        self, compiled: CompiledObj, obj_type: str = None
-    ) -> None:
-
+    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}"
+            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()
@@ -43,7 +40,6 @@ class CompileStatus:
         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:
@@ -63,10 +59,14 @@ class CompileStatus:
         if self.use_live:
             self.live.stop()
 
-    def render(self) -> Text:
+    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()
@@ -77,7 +77,7 @@ class CompileStatus:
             if errors:
                 text.append(errors)
 
-            diff = tracker.diff()
+            diff = tracker.diff(ignore_python_errors)
             if diff:
                 text.append(diff)
 

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

@@ -1,36 +1,101 @@
-from typing import List
+from typing import List, Tuple
 
 from rich.text import Text
 
+from ai.chronon.cli.compile.version_utils import parse_name_and_version
 
-class DiffResult:
 
+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")
+            return Text("➕ Added", style="dim green")
 
         def updated_signage():
-            return Text("Updated", style="dim yellow")
+            return Text("❗ Changed in place (no version change)", style="dim yellow")
 
         def deleted_signage():
-            return Text("Deleted", style="red")
+            return Text("🗑️ Deleted", style="red")
 
-        added = [(added_signage(), name) for name in self.added]
+        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]
-
-        result_order = added + updated
-
-        if deleted_names:
-            deleted = [(deleted_signage(), name) for name in deleted_names]
-            result_order += deleted
-
-        result_order = sorted(result_order, key=lambda t: t[1])
+        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:

ai/chronon/cli/compile/fill_templates.py

@@ -1,26 +1,22 @@
+from gen_thrift.api.ttypes import Join, Team
+
 from ai.chronon import utils
-from ai.chronon.api.ttypes import Join, Team
 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("{{ 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)
 
@@ -30,7 +26,6 @@ def set_templated_values(obj, cls, compile_context: CompileContext):
             ]
 
     if cls == Join and obj.labelParts:
-
         obj.labelParts.metaData.dependencies = [
             label_dep.replace(
                 "{{ join_backfill_table }}",

ai/chronon/cli/compile/parse_configs.py

@@ -4,8 +4,9 @@ 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.api.ttypes import GroupBy, Join
 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
@@ -13,9 +14,8 @@ from ai.chronon.cli.logger import get_logger
 
 logger = get_logger()
 
-def from_folder(
-    cls: type, input_dir: str, compile_context: CompileContext
-) -> List[CompiledObj]:
+
+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
@@ -26,7 +26,6 @@ def from_folder(
     results = []
 
     for f in python_files:
-
         try:
             results_dict = from_file(f, cls, input_dir)
 
@@ -55,9 +54,7 @@ def from_folder(
                 )
                 results.append(result)
 
-                compile_context.compile_status.add_object_update_display(
-                    result, cls.__name__
-                )
+                compile_context.compile_status.add_object_update_display(result, cls.__name__)
 
         except Exception as e:
             result = CompiledObj(
@@ -71,15 +68,12 @@ def from_folder(
 
             results.append(result)
 
-            compile_context.compile_status.add_object_update_display(
-                result, cls.__name__
-            )
+            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)
@@ -96,16 +90,14 @@ def from_file(file_path: str, cls: type, input_dir: str):
     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]
 
@@ -113,6 +105,7 @@ def from_file(file_path: str, cls: type, input_dir: str):
 
     return result
 
+
 def populate_column_hashes(obj: Any):
     """
     Populate the columnHashes field in the object's metadata with semantic hashes
@@ -135,7 +128,7 @@ def populate_column_hashes(obj: Any):
         obj.metaData.columnHashes = column_hashes
 
         if obj.joinParts:
-            for jp in (obj.joinParts or []):
+            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

ai/chronon/cli/compile/parse_teams.py

@@ -6,13 +6,14 @@ from copy import deepcopy
 from enum import Enum
 from typing import Any, Dict, Optional, Union
 
-from ai.chronon.api.common.ttypes import (
+from gen_thrift.api.ttypes import Join, MetaData, Team
+from gen_thrift.common.ttypes import (
     ClusterConfigProperties,
     ConfigProperties,
     EnvironmentVariables,
     ExecutionInfo,
 )
-from ai.chronon.api.ttypes import Join, MetaData, Team
+
 from ai.chronon.cli.compile.display.console import console
 from ai.chronon.cli.logger import get_logger
 
@@ -43,9 +44,9 @@ def import_module_from_file(file_path):
 def load_teams(conf_root: str, print: bool = True) -> Dict[str, Team]:
     teams_file = os.path.join(conf_root, "teams.py")
 
-    assert os.path.exists(
-        teams_file
-    ), f"Team config file: {teams_file} not found. You might be running this from the wrong directory."
+    assert os.path.exists(teams_file), (
+        f"Team config file: {teams_file} not found. You might be running this from the wrong directory."
+    )
 
     team_module = import_module_from_file(teams_file)
 
@@ -57,9 +58,7 @@ def load_teams(conf_root: str, print: bool = True) -> Dict[str, Team]:
     team_dict = {}
 
     if print:
-        console.print(
-            f"Pulling configuration from [cyan italic]{teams_file}[/cyan italic]"
-        )
+        console.print(f"Pulling configuration from [cyan italic]{teams_file}[/cyan italic]")
 
     for name, obj in team_module.__dict__.items():
         if isinstance(obj, Team):
@@ -79,17 +78,15 @@ def update_metadata(obj: Any, team_dict: Dict[str, Team]):
     name = obj.metaData.name
     team = obj.metaData.team
 
-    assert (
-            team is not None
-    ), f"Team name is required in metadata for {name}. This usually set by compiler. Internal error."
+    assert team is not None, (
+        f"Team name is required in metadata for {name}. This usually set by compiler. Internal error."
+    )
 
-    assert (
-            team in team_dict
-    ), f"Team '{team}' not found in teams.py. Please add an entry 🙏"
+    assert team in team_dict, f"Team '{team}' not found in teams.py. Please add an entry 🙏"
 
-    assert (
-            _DEFAULT_CONF_TEAM in team_dict
-    ), f"'{_DEFAULT_CONF_TEAM}' team not found in teams.py, please add an entry 🙏."
+    assert _DEFAULT_CONF_TEAM in team_dict, (
+        f"'{_DEFAULT_CONF_TEAM}' team not found in teams.py, please add an entry 🙏."
+    )
 
     # Only set the outputNamespace if it hasn't been set already
     if not metadata.outputNamespace:
@@ -111,12 +108,12 @@ def update_metadata(obj: Any, team_dict: Dict[str, Team]):
                     join_part_gb.metaData.outputNamespace = output_namespace
 
         if obj.joinParts:
-            for jp in (obj.joinParts or []):
+            for jp in obj.joinParts or []:
                 jp.useLongNames = obj.useLongNames
                 set_group_by_metadata(jp.groupBy, join_namespace)
 
         if obj.labelParts:
-            for lb in (obj.labelParts.labels or []):
+            for lb in obj.labelParts.labels or []:
                 lb.useLongNames = obj.useLongNames
                 set_group_by_metadata(lb.groupBy, join_namespace)
 
@@ -161,7 +158,6 @@ def _merge_maps(*maps: Optional[Dict[str, str]]):
     result = {}
 
     for m in maps:
-
         if m is None:
             continue
 
@@ -178,22 +174,25 @@ class EnvOrConfigAttribute(str, Enum):
 
 
 def _merge_mode_maps(
-        *mode_maps: Optional[Union[EnvironmentVariables, ConfigProperties, ClusterConfigProperties]],
-        env_or_config_attribute: EnvOrConfigAttribute,
+    *mode_maps: Optional[Union[EnvironmentVariables, ConfigProperties, ClusterConfigProperties]],
+    env_or_config_attribute: EnvOrConfigAttribute,
 ):
     """
     Merges multiple environment variables into one - with the later maps overriding the earlier ones.
     """
 
     # Merge `common` to each individual mode map. Creates a new map
-    def push_common_to_modes(mode_map: Union[EnvironmentVariables, ConfigProperties], mode_key: EnvOrConfigAttribute):
+    def push_common_to_modes(
+        mode_map: Union[EnvironmentVariables, ConfigProperties], mode_key: EnvOrConfigAttribute
+    ):
         final_mode_map = deepcopy(mode_map)
         common = final_mode_map.common
         modes = getattr(final_mode_map, mode_key)
-        for _ in modes:
-            modes[_] = _merge_maps(
-                common, modes[_]
-            )
+
+        if modes:
+            for _ in modes:
+                modes[_] = _merge_maps(common, modes[_])
+
         return final_mode_map
 
     filtered_mode_maps = [m for m in mode_maps if m]
@@ -217,22 +216,27 @@ def _merge_mode_maps(
         current_modes = getattr(result, env_or_config_attribute)
         incoming_modes = getattr(incoming_mode_map, env_or_config_attribute)
 
-        current_modes_keys = list(current_modes.keys())
-        incoming_modes_keys = list(incoming_modes.keys())
+        current_modes_keys = list(current_modes.keys()) if current_modes else []
+        incoming_modes_keys = list(incoming_modes.keys()) if incoming_modes else []
 
         all_modes_keys = list(set(current_modes_keys + incoming_modes_keys))
+
         for mode in all_modes_keys:
-            current_mode = current_modes.get(mode, {})
+            current_mode = current_modes.get(mode, {}) if current_modes else {}
 
             # if the incoming_mode is not found, we NEED to default to incoming_common
-            incoming_mode = incoming_modes.get(mode, incoming_common)
+            incoming_mode = (
+                incoming_modes.get(mode, incoming_common) if incoming_modes else incoming_common
+            )
 
             # first to last with later ones overriding the earlier ones
             # common -> current mode level -> incoming mode level
 
-            new_mode = _merge_maps(
-                new_common, current_mode, incoming_mode
-            )
+            new_mode = _merge_maps(new_common, current_mode, incoming_mode)
+
+            if current_modes is None:
+                current_modes = {}
+
             current_modes[mode] = new_mode
 
     return result

ai/chronon/cli/compile/serializer.py

@@ -45,18 +45,14 @@ class ThriftJSONDecoder(json.JSONDecoder):
                 (_, field_ttype, field_name, field_ttype_info, dummy) = field
                 if field_name not in val:
                     continue
-                converted_val = self._convert(
-                    val[field_name], field_ttype, field_ttype_info
-                )
+                converted_val = self._convert(val[field_name], field_ttype, field_ttype_info)
                 setattr(ret, field_name, converted_val)
         elif ttype == TType.LIST:
             (element_ttype, element_ttype_info, _) = ttype_info
             ret = [self._convert(x, element_ttype, element_ttype_info) for x in val]
         elif ttype == TType.SET:
             (element_ttype, element_ttype_info) = ttype_info
-            ret = set(
-                [self._convert(x, element_ttype, element_ttype_info) for x in val]
-            )
+            ret = set([self._convert(x, element_ttype, element_ttype_info) for x in val])
         elif ttype == TType.MAP:
             (key_ttype, key_ttype_info, val_ttype, val_ttype_info, _) = ttype_info
             ret = dict(
@@ -108,8 +104,6 @@ def file2thrift(path, thrift_class):
 
 
 def thrift_simple_json(obj):
-    simple = TSerialization.serialize(
-        obj, protocol_factory=TSimpleJSONProtocolFactory()
-    )
+    simple = TSerialization.serialize(obj, protocol_factory=TSimpleJSONProtocolFactory())
     parsed = json.loads(simple)
     return json.dumps(parsed, indent=2, sort_keys=True)

ai/chronon/cli/compile/version_utils.py

@@ -0,0 +1,42 @@
+"""
+Utilities for handling config versioning in Chronon.
+"""
+
+from typing import Optional, Tuple
+
+
+def parse_name_and_version(name: str) -> Tuple[str, Optional[int]]:
+    """Parse config name to extract base name and version.
+
+    Args:
+        name: Config name (e.g., 'config_name__1' or 'config_name')
+
+    Returns:
+        Tuple of (base_name, version) where version is None if no version suffix
+    """
+    if "__" in name:
+        parts = name.rsplit("__", 1)
+        if len(parts) == 2 and parts[1].isdigit():
+            return parts[0], int(parts[1])
+    return name, None
+
+
+def is_version_change(deleted_name: str, added_name: str) -> bool:
+    """Check if a deleted/added pair represents a version change.
+
+    Args:
+        deleted_name: Name of deleted config
+        added_name: Name of added config
+
+    Returns:
+        True if this represents a version bump of the same config
+    """
+    deleted_base, deleted_version = parse_name_and_version(deleted_name)
+    added_base, added_version = parse_name_and_version(added_name)
+
+    return (
+        deleted_base == added_base
+        and deleted_version is not None
+        and added_version is not None
+        and deleted_version != added_version
+    )

ai/chronon/cli/git_utils.py

@@ -9,7 +9,6 @@ logger = get_logger()
 
 
 def get_current_branch() -> str:
-
     try:
         subprocess.check_output(["git", "rev-parse", "HEAD"], stderr=subprocess.DEVNULL)
 
@@ -20,7 +19,6 @@ def get_current_branch() -> str:
         )
 
     except subprocess.CalledProcessError as e:
-
         try:
             head_file = Path(".git/HEAD").resolve()
 
@@ -43,7 +41,6 @@ def get_current_branch() -> str:
 
 def get_fork_point(base_branch: str = "main") -> str:
     try:
-
         return (
             subprocess.check_output(["git", "merge-base", base_branch, "HEAD"])
             .decode("utf-8")
@@ -60,9 +57,7 @@ def get_fork_point(base_branch: str = "main") -> str:
 
 def get_file_content_at_commit(file_path: str, commit: str) -> Optional[str]:
     try:
-        return subprocess.check_output(["git", "show", f"{commit}:{file_path}"]).decode(
-            "utf-8"
-        )
+        return subprocess.check_output(["git", "show", f"{commit}:{file_path}"]).decode("utf-8")
     except subprocess.CalledProcessError:
         return None
 
@@ -75,7 +70,6 @@ def get_current_file_content(file_path: str) -> Optional[str]:
 
 
 def get_changes_since_commit(path: str, commit: Optional[str] = None) -> List[str]:
-
     path = Path(path).resolve()
     if not path.exists():
         print(f"⛔ Error: Path does not exist: {path}", file=sys.stderr)
@@ -86,15 +80,12 @@ def get_changes_since_commit(path: str, commit: Optional[str] = None) -> List[st
         commit_range = f"{commit}..HEAD" if commit else "HEAD"
 
         changes = (
-            subprocess.check_output(
-                ["git", "diff", "--name-only", commit_range, "--", str(path)]
-            )
+            subprocess.check_output(["git", "diff", "--name-only", commit_range, "--", str(path)])
             .decode("utf-8")
             .splitlines()
         )
 
     except subprocess.CalledProcessError:
-
         changes = (
             subprocess.check_output(["git", "diff", "--name-only", "--", str(path)])
             .decode("utf-8")
@@ -102,7 +93,6 @@ def get_changes_since_commit(path: str, commit: Optional[str] = None) -> List[st
         )
 
     try:
-
         untracked = (
             subprocess.check_output(
                 ["git", "ls-files", "--others", "--exclude-standard", str(path)]
@@ -114,7 +104,6 @@ def get_changes_since_commit(path: str, commit: Optional[str] = None) -> List[st
         changes.extend(untracked)
 
     except subprocess.CalledProcessError as e:
-
         print(
             f"⛔ Error: {e.stderr.decode('utf-8') if e.stderr else 'Failed to get untracked files'}",
             file=sys.stderr,

ai/chronon/cli/logger.py

@@ -15,9 +15,7 @@ RESET = "\033[0m"
 
 
 class ColorFormatter(logging.Formatter):
-
     def format(self, record):
-
         time_str = datetime.fromtimestamp(record.created).strftime("%H:%M:%S")
         level_color = LEVEL_COLORS.get(record.levelno)
 

ai/chronon/constants.py

@@ -1,3 +1,3 @@
 AIRFLOW_DEPENDENCIES_KEY = "airflowDependencies"
 AIRFLOW_LABEL_DEPENDENCIES_KEY = "airflowLabelDependencies"
-PARTITION_COLUMN_KEY = "spark.chronon.partition.column"
+PARTITION_COLUMN_KEY = "spark.chronon.partition.column"