WIP

Author: ezvz

api/py/ai/chronon/repo/__init__.py

@@ -20,7 +20,7 @@
 MODEL_FOLDER_NAME = "models"
 # TODO - make team part of thrift API?
 TEAMS_FILE_PATH = "teams.json"
-
+OUTPUT_ROOT = "production"
 
 # This is set in the main function -
 # from command line or from env variable during invocation

api/py/ai/chronon/repo/compilev2.py

@@ -71,8 +71,7 @@ def extract_and_convert(
     validator = ChrononRepoValidator(
             chronon_root, os.path.join(chronon_root_path, "production"), log_level=log_level
     )
-
-    successful_writes = {}
+    
     compile_errors = {}
 
     for obj_folder_name in obj_folder_names:
@@ -88,7 +87,6 @@ def extract_and_convert(
             raise ValueError(f"Error in file {target_object_file}: \n {target_file_error}")
 
         compile_errors.update(obj_folder_errors)
-        num_written_objs = 0
 
         full_output_root = os.path.join(chronon_root_path, output_root)
         teams_path = os.path.join(chronon_root_path, TEAMS_FILE_PATH)
@@ -106,12 +104,6 @@ def extract_and_convert(
             if obj_write_errors:
                 compile_errors[origin_file] = obj_write_errors
             else:
-                num_written_objs += 1
-
-                # TODO -- handle target object file and logging in case of failure.
-                if obj == target_object:
-                    target_object_written = True
-
                 # In case of online join, we need to materialize the underlying online group_bys.
                 if obj_class is Join and obj.metaData.online:
                     online_group_bys = {}
@@ -126,10 +118,6 @@ def extract_and_convert(
                         " You can do this by passing the `online=True` argument to the GroupBy constructor."
                         " Fix the following: {}".format(offline_gbs)
                     )
-            
-        if num_written_objs > 0:
-            obj_class_name = obj_class.__name__
-            successful_writes[obj_class_name] = successful_writes.get(obj_class_name, 0) + num_written_objs
 
     if compile_errors:
         create_error_logs(compile_errors, chronon_root)
@@ -155,8 +143,8 @@ def show_lineage_and_schema(target_object):
         raise NotImplementedError("TODO")
     except Exception as e:
 
-        _print_warning("Failed to connect to ZiplineHub (check connection and VPN settings). \n\n " + 
-              "Showing output column names, but cannot show schema/lineage without ZiplineHub.\n\n")
+        _print_warning(f"Failed to connect to ZiplineHub: {str(e)}\n\n" + 
+            "Showing output column names, but cannot show schema/lineage without ZiplineHub.\n\n")
 
         obj_class = target_object.__class__
 
@@ -293,7 +281,3 @@ def _print_error(left, right):
 def _print_warning(string):
     # print in yellow - \u001b[33m
     print(f"\u001b[33m{string}\u001b[0m")
-
-
-if __name__ == "__main__":
-    extract_and_convert()

api/py/ai/chronon/repo/extract_objects.py

@@ -42,8 +42,8 @@ def from_folder(full_path: str, cls: type, log_level=logging.INFO):
 
 def from_folderV2(full_path: str, target_file: str, cls: type):
     """
-    Recursively consumes a folder, and constructs a map
-    Creates a map of object qualifier to
+    Recursively consumes a folder, and constructs a map of 
+    object qualifier to StagingQuery, GroupBy, or Join
     """
     if full_path.endswith("/"):
         full_path = full_path[:-1]
@@ -57,7 +57,6 @@ def from_folderV2(full_path: str, target_file: str, cls: type):
             results_dict = from_file(f, cls, log_level=logging.NOTSET)
             for k, v in results_dict.items():
                 results[k] = (v, f)
-            #results.update(from_file(f, cls, log_level=logging.NOTSET))
         except Exception as e:
             if f == target_file:
                 target_file_error = e

api/py/ai/chronon/repo/runner.py

@@ -2,15 +2,20 @@
 import inspect
 from ai.chronon.repo.compilev2 import extract_and_convert
 from ai.chronon.utils import get_mod_and_var_name_from_gc
-
-# TODO -- how do we want to set this?
-OUTPUT_ROOT = "production"
+from ai.chronon.repo import (
+    FOLDER_NAME_TO_CLASS,
+    JOIN_FOLDER_NAME,
+    GROUP_BY_FOLDER_NAME,
+    STAGING_QUERY_FOLDER_NAME,
+    MODEL_FOLDER_NAME,
+    OUTPUT_ROOT
+)
 
 def get_chronon_root(filepath):
     """
     Infer chronon root from a filepath to a Chronon object
     """
-    target_dirs = {'group_bys', 'joins', 'staging_queries'}
+    target_dirs = FOLDER_NAME_TO_CLASS.keys()
 
     current_path = os.path.dirname(os.path.abspath(filepath))
 
@@ -27,8 +32,13 @@ def get_target_compiled_file(zipline_obj):
     """
     Given a compiled repository, identifies the compiled file
     for this entity
-    """    
-    zipline_obj_type = "group_bys"
+    """        
+    # Validate object is a valid Chronon type
+    if not any(isinstance(zipline_obj, cls) for cls in FOLDER_NAME_TO_CLASS.values()):
+        valid_types = [cls.__name__ for cls in FOLDER_NAME_TO_CLASS.values()]
+        raise ValueError(f"Can only run one of {valid_types}, got {type(zipline_obj).__name__}")
+    
+    zipline_obj_type = next(k for k, v in FOLDER_NAME_TO_CLASS.items() if v == zipline_obj.__class__)
 
     zipline_obj_module, zipline_var_name = get_mod_and_var_name_from_gc(zipline_obj, zipline_obj_type)