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

ai/chronon/cli/compile/parse_teams.py

@@ -0,0 +1,242 @@
+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 gen_thrift.api.ttypes import Join, MetaData, Team
+from gen_thrift.common.ttypes import (
+    ClusterConfigProperties,
+    ConfigProperties,
+    EnvironmentVariables,
+    ExecutionInfo,
+)
+
+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)
+
+        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]
+
+    # 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()) 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, {}) 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) 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)
+
+            if current_modes is None:
+                current_modes = {}
+
+            current_modes[mode] = new_mode
+
+    return result

ai/chronon/cli/compile/serializer.py

@@ -0,0 +1,109 @@
+#     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)

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

@@ -0,0 +1,145 @@
+import subprocess
+import sys
+from pathlib import Path
+from typing import List, Optional
+
+from ai.chronon.cli.logger import get_logger
+
+logger = get_logger()
+
+
+def get_current_branch() -> str:
+    try:
+        subprocess.check_output(["git", "rev-parse", "HEAD"], stderr=subprocess.DEVNULL)
+
+        return (
+            subprocess.check_output(["git", "rev-parse", "--abbrev-ref", "HEAD"])
+            .decode("utf-8")
+            .strip()
+        )
+
+    except subprocess.CalledProcessError as e:
+        try:
+            head_file = Path(".git/HEAD").resolve()
+
+            if head_file.exists():
+                content = head_file.read_text().strip()
+
+                if content.startswith("ref: refs/heads/"):
+                    return content.split("/")[-1]
+
+        except Exception:
+            pass
+
+        print(
+            f"⛔ Error: {e.stderr.decode('utf-8') if e.stderr else 'Not a git repository or no commits'}",
+            file=sys.stderr,
+        )
+
+        raise
+
+
+def get_fork_point(base_branch: str = "main") -> str:
+    try:
+        return (
+            subprocess.check_output(["git", "merge-base", base_branch, "HEAD"])
+            .decode("utf-8")
+            .strip()
+        )
+
+    except subprocess.CalledProcessError as e:
+        print(
+            f"⛔ Error: {e.stderr.decode('utf-8') if e.stderr else f'Could not determine fork point from {base_branch}'}",
+            file=sys.stderr,
+        )
+        raise
+
+
+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")
+    except subprocess.CalledProcessError:
+        return None
+
+
+def get_current_file_content(file_path: str) -> Optional[str]:
+    try:
+        return Path(file_path).read_text()
+    except Exception:
+        return None
+
+
+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)
+        raise ValueError(f"Path does not exist: {path}")
+
+    try:
+        subprocess.check_output(["git", "rev-parse", "HEAD"], stderr=subprocess.DEVNULL)
+        commit_range = f"{commit}..HEAD" if commit else "HEAD"
+
+        changes = (
+            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")
+            .splitlines()
+        )
+
+    try:
+        untracked = (
+            subprocess.check_output(
+                ["git", "ls-files", "--others", "--exclude-standard", str(path)]
+            )
+            .decode("utf-8")
+            .splitlines()
+        )
+
+        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,
+        )
+
+        raise
+
+    logger.info(f"Changes since commit: {changes}")
+
+    return [change for change in changes if change.strip()]
+
+
+def get_changes_since_fork(path: str, base_branch: str = "main") -> List[str]:
+    try:
+        fork_point = get_fork_point(base_branch)
+        path = Path(path).resolve()
+
+        # Get all potential changes
+        changed_files = set(get_changes_since_commit(str(path), fork_point))
+
+        # Filter out files that are identical to fork point
+        real_changes = []
+        for file in changed_files:
+            fork_content = get_file_content_at_commit(file, fork_point)
+            current_content = get_current_file_content(file)
+
+            if fork_content != current_content:
+                real_changes.append(file)
+
+        logger.info(f"Changes since fork: {real_changes}")
+
+        return real_changes
+
+    except subprocess.CalledProcessError as e:
+        print(
+            f"⛔ Error: {e.stderr.decode('utf-8') if e.stderr else f'Failed to get changes since fork from {base_branch}'}",
+            file=sys.stderr,
+        )
+        raise

ai/chronon/cli/logger.py

@@ -0,0 +1,59 @@
+import logging
+import sys
+from datetime import datetime
+
+TIME_COLOR = "\033[36m"  # Cyan
+LEVEL_COLORS = {
+    logging.DEBUG: "\033[36m",  # Cyan
+    logging.INFO: "\033[32m",  # Green
+    logging.WARNING: "\033[33m",  # Yellow
+    logging.ERROR: "\033[31m",  # Red
+    logging.CRITICAL: "\033[41m",  # White on Red
+}
+FILE_COLOR = "\033[35m"  # Purple
+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)
+
+        return (
+            f"{TIME_COLOR}{time_str}{RESET} "
+            f"{level_color}{record.levelname}{RESET} "
+            f"{FILE_COLOR}{record.filename}:{record.lineno}{RESET} - "
+            f"{record.getMessage()}"
+        )
+
+
+def get_logger(log_level=logging.INFO):
+    logger = logging.getLogger(__name__)
+
+    # no need to reset if a handler already exists
+    if not logger.hasHandlers():
+        handler = logging.StreamHandler(sys.stdout)
+        handler.setFormatter(ColorFormatter())
+
+        logger.addHandler(handler)
+        logger.setLevel(log_level)
+
+    return logger
+
+
+def red(text):
+    return f"\033[1;91m{text}\033[0m"
+
+
+def green(text):
+    return f"\033[1;92m{text}\033[0m"
+
+
+def require(cond, message):
+    if not cond:
+        print(f"X: {message}")
+        sys.exit(1)
+
+
+def done(cond, message):
+    print(f"DONE: {message}")

ai/chronon/constants.py

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