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

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) -> Text:
+        text = Text(overflow="fold", no_wrap=False)
+
+        for obj_type, tracker in self.cls_to_tracker.items():
+            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()
+            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,46 @@
+from typing import List
+
+from rich.text import Text
+
+
+class DiffResult:
+
+    def __init__(self):
+        self.added: List[str] = []
+        self.updated: List[str] = []
+
+    def render(self, deleted_names: List[str], indent="  ") -> Text:
+
+        def added_signage():
+            return Text("Added", style="dim green")
+
+        def updated_signage():
+            return Text("Updated", style="dim yellow")
+
+        def deleted_signage():
+            return Text("Deleted", style="red")
+
+        added = [(added_signage(), name) for name in self.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])
+
+        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,40 @@
+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("{{ 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,141 @@
+import copy
+import glob
+import importlib
+import os
+from typing import Any, List
+
+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
+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

ai/chronon/cli/compile/parse_teams.py

@@ -0,0 +1,238 @@
+import importlib
+import importlib.util
+import os
+import sys
+from copy import deepcopy
+from enum import Enum
+from typing import Any, Dict, Optional, Union
+
+from ai.chronon.api.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
+
+logger = get_logger()
+
+_DEFAULT_CONF_TEAM = "default"
+
+
+def import_module_from_file(file_path):
+    # Get the module name from the file path (without .py extension)
+    module_name = file_path.split("/")[-1].replace(".py", "")
+
+    # Create the module spec
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+
+    # Create the module based on the spec
+    module = importlib.util.module_from_spec(spec)
+
+    # Add the module to sys.modules
+    sys.modules[module_name] = module
+
+    # Execute the module
+    spec.loader.exec_module(module)
+
+    return module
+
+
+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."
+
+    team_module = import_module_from_file(teams_file)
+
+    assert team_module is not None, (
+        f"Team config file {teams_file} is not on the PYTHONPATH. You might need to add the your config "
+        f"directory to the PYTHONPATH."
+    )
+
+    team_dict = {}
+
+    if print:
+        console.print(
+            f"Pulling configuration from [cyan italic]{teams_file}[/cyan italic]"
+        )
+
+    for name, obj in team_module.__dict__.items():
+        if isinstance(obj, Team):
+            obj.name = name
+            team_dict[name] = obj
+
+    return team_dict
+
+
+def update_metadata(obj: Any, team_dict: Dict[str, Team]):
+    assert obj is not None, "Cannot update metadata None object"
+
+    metadata = obj.metaData
+
+    assert obj.metaData is not None, "Cannot update empty metadata"
+
+    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 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 🙏."
+
+    # Only set the outputNamespace if it hasn't been set already
+    if not metadata.outputNamespace:
+        metadata.outputNamespace = team_dict[team].outputNamespace
+
+    if isinstance(obj, Join):
+        join_namespace = obj.metaData.outputNamespace
+
+        # set the metadata for each join part and labelParts
+        def set_group_by_metadata(join_part_gb, output_namespace):
+            if join_part_gb is not None:
+                if join_part_gb.metaData:
+                    # Only set the outputNamespace if it hasn't been set already
+                    if not join_part_gb.metaData.outputNamespace:
+                        join_part_gb.metaData.outputNamespace = output_namespace
+                else:
+                    # If there's no metaData at all, create it and set outputNamespace
+                    join_part_gb.metaData = MetaData()
+                    join_part_gb.metaData.outputNamespace = output_namespace
+
+        if obj.joinParts:
+            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 []):
+                lb.useLongNames = obj.useLongNames
+                set_group_by_metadata(lb.groupBy, join_namespace)
+
+    if metadata.executionInfo is None:
+        metadata.executionInfo = ExecutionInfo()
+
+    merge_team_execution_info(metadata, team_dict, team)
+
+
+def merge_team_execution_info(metadata: MetaData, team_dict: Dict[str, Team], team_name: str):
+    default_team = team_dict.get(_DEFAULT_CONF_TEAM)
+    if not metadata.executionInfo:
+        metadata.executionInfo = ExecutionInfo()
+
+    metadata.executionInfo.env = _merge_mode_maps(
+        default_team.env if default_team else {},
+        team_dict[team_name].env,
+        metadata.executionInfo.env,
+        env_or_config_attribute=EnvOrConfigAttribute.ENV,
+    )
+
+    metadata.executionInfo.conf = _merge_mode_maps(
+        default_team.conf if default_team else {},
+        team_dict[team_name].conf,
+        metadata.executionInfo.conf,
+        env_or_config_attribute=EnvOrConfigAttribute.CONFIG,
+    )
+
+    metadata.executionInfo.clusterConf = _merge_mode_maps(
+        default_team.clusterConf if default_team else {},
+        team_dict[team_name].clusterConf,
+        metadata.executionInfo.clusterConf,
+        env_or_config_attribute=EnvOrConfigAttribute.CLUSTER_CONFIG,
+    )
+
+
+def _merge_maps(*maps: Optional[Dict[str, str]]):
+    """
+    Merges multiple maps into one - with the later maps overriding the earlier ones.
+    """
+
+    result = {}
+
+    for m in maps:
+
+        if m is None:
+            continue
+
+        for key, value in m.items():
+            result[key] = value
+
+    return result
+
+
+class EnvOrConfigAttribute(str, Enum):
+    ENV = "modeEnvironments"
+    CONFIG = "modeConfigs"
+    CLUSTER_CONFIG = "modeClusterConfigs"
+
+
+def _merge_mode_maps(
+        *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):
+        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[_]
+            )
+        return final_mode_map
+
+    filtered_mode_maps = [m for m in mode_maps if m]
+
+    # Initialize the result with the first mode map
+    result = None
+
+    if len(filtered_mode_maps) >= 1:
+        result = push_common_to_modes(filtered_mode_maps[0], env_or_config_attribute)
+
+    # Merge each new mode map into the result
+    for m in filtered_mode_maps[1:]:
+        # We want to prepare the individual modes with `common` in incoming_mode_map
+        incoming_mode_map = push_common_to_modes(m, env_or_config_attribute)
+
+        # create new common
+        incoming_common = incoming_mode_map.common
+        new_common = _merge_maps(result.common, incoming_common)
+        result.common = new_common
+
+        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())
+
+        all_modes_keys = list(set(current_modes_keys + incoming_modes_keys))
+        for mode in all_modes_keys:
+            current_mode = current_modes.get(mode, {})
+
+            # if the incoming_mode is not found, we NEED to default to incoming_common
+            incoming_mode = incoming_modes.get(mode, 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
+            )
+            current_modes[mode] = new_mode
+
+    return result

ai/chronon/cli/compile/serializer.py

@@ -0,0 +1,115 @@
+#     Copyright (C) 2023 The Chronon Authors.
+#
+#     Licensed under the Apache License, Version 2.0 (the "License");
+#     you may not use this file except in compliance with the License.
+#     You may obtain a copy of the License at
+#
+#         http://www.apache.org/licenses/LICENSE-2.0
+#
+#     Unless required by applicable law or agreed to in writing, software
+#     distributed under the License is distributed on an "AS IS" BASIS,
+#     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#     See the License for the specific language governing permissions and
+#     limitations under the License.
+
+import json
+
+from thrift import TSerialization
+from thrift.protocol.TBinaryProtocol import TBinaryProtocolAccelerated
+from thrift.protocol.TJSONProtocol import TSimpleJSONProtocolFactory
+from thrift.Thrift import TType
+from thrift.transport.TTransport import TMemoryBuffer
+
+
+class ThriftJSONDecoder(json.JSONDecoder):
+    def __init__(self, *args, **kwargs):
+        self._thrift_class = kwargs.pop("thrift_class")
+        super(ThriftJSONDecoder, self).__init__(*args, **kwargs)
+
+    def decode(self, json_str):
+        if isinstance(json_str, dict):
+            dct = json_str
+        else:
+            dct = super(ThriftJSONDecoder, self).decode(json_str)
+        return self._convert(
+            dct, TType.STRUCT, (self._thrift_class, self._thrift_class.thrift_spec)
+        )
+
+    def _convert(self, val, ttype, ttype_info):
+        if ttype == TType.STRUCT:
+            (thrift_class, thrift_spec) = ttype_info
+            ret = thrift_class()
+            for field in thrift_spec:
+                if field is None:
+                    continue
+                (_, 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
+                )
+                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]
+            )
+        elif ttype == TType.MAP:
+            (key_ttype, key_ttype_info, val_ttype, val_ttype_info, _) = ttype_info
+            ret = dict(
+                [
+                    (
+                        self._convert(k, key_ttype, key_ttype_info),
+                        self._convert(v, val_ttype, val_ttype_info),
+                    )
+                    for (k, v) in val.items()
+                ]
+            )
+        elif ttype == TType.STRING:
+            ret = str(val)
+        elif ttype == TType.DOUBLE:
+            ret = float(val)
+        elif ttype == TType.I64:
+            ret = int(val)
+        elif ttype == TType.I32 or ttype == TType.I16 or ttype == TType.BYTE:
+            ret = int(val)
+        elif ttype == TType.BOOL:
+            ret = bool(val)
+        else:
+            raise TypeError("Unrecognized thrift field type: %d" % ttype)
+        return ret
+
+
+def json2thrift(json_str, thrift_class):
+    return json.loads(json_str, cls=ThriftJSONDecoder, thrift_class=thrift_class)
+
+
+def json2binary(json_str, thrift_class):
+    thrift = json2thrift(json_str, thrift_class)
+    transport = TMemoryBuffer()
+    protocol = TBinaryProtocolAccelerated(transport)
+    thrift.write(protocol)
+    # Get the raw bytes representing the object in Thrift binary format
+    return transport.getvalue()
+
+
+def file2thrift(path, thrift_class):
+    try:
+        with open(path, "r") as file:
+            return json2thrift(file.read(), thrift_class)
+    except json.decoder.JSONDecodeError as e:
+        raise Exception(
+            f"Error decoding file into a {thrift_class.__name__}:  {path}. "
+            + f"Please double check that {path} represents a valid {thrift_class.__name__}."
+        ) from e
+
+
+def thrift_simple_json(obj):
+    simple = TSerialization.serialize(
+        obj, protocol_factory=TSimpleJSONProtocolFactory()
+    )
+    parsed = json.loads(simple)
+    return json.dumps(parsed, indent=2, sort_keys=True)