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

ai/chronon/repo/utils.py

@@ -1,3 +1,4 @@
+import functools
 import json
 import os
 import re
@@ -7,6 +8,8 @@ import xml.etree.ElementTree as ET
 from datetime import datetime, timedelta
 from enum import Enum
 
+from click import style
+
 from ai.chronon.cli.compile.parse_teams import EnvOrConfigAttribute
 from ai.chronon.logger import get_logger
 from ai.chronon.repo.constants import (
@@ -17,6 +20,7 @@ from ai.chronon.repo.constants import (
 
 LOG = get_logger()
 
+
 class JobType(Enum):
     SPARK = "spark"
     FLINK = "flink"
@@ -52,9 +56,11 @@ def get_environ_arg(env_name, ignoreError=False) -> str:
         raise ValueError(f"Please set {env_name} environment variable")
     return value
 
+
 def get_customer_warehouse_bucket() -> str:
     return f"zipline-warehouse-{get_customer_id()}"
 
+
 def get_customer_id() -> str:
     return get_environ_arg("CUSTOMER_ID")
 
@@ -94,17 +100,13 @@ def download_only_once(url, path, skip_download=False):
         LOG.info(
             """Files sizes of {url} vs. {path}
     Remote size: {remote_size}
-    Local size : {local_size}""".format(
-                **locals()
-            )
+    Local size : {local_size}""".format(**locals())
         )
         if local_size == remote_size:
             LOG.info("Sizes match. Assuming it's already downloaded.")
             should_download = False
         if should_download:
-            LOG.info(
-                "Different file from remote at local: " + path + ". Re-downloading.."
-            )
+            LOG.info("Different file from remote at local: " + path + ". Re-downloading..")
             check_call("curl {} -o {} --connect-timeout 10".format(url, path))
     else:
         LOG.info("No file at: " + path + ". Downloading..")
@@ -126,9 +128,7 @@ def download_jar(
     )
     scala_version = SCALA_VERSION_FOR_SPARK[spark_version]
     maven_url_prefix = os.environ.get("CHRONON_MAVEN_MIRROR_PREFIX", None)
-    default_url_prefix = (
-        "https://s01.oss.sonatype.org/service/local/repositories/public/content"
-    )
+    default_url_prefix = "https://s01.oss.sonatype.org/service/local/repositories/public/content"
     url_prefix = maven_url_prefix if maven_url_prefix else default_url_prefix
     base_url = "{}/ai/chronon/spark_{}_{}".format(url_prefix, jar_type, scala_version)
     LOG.info("Downloading jar from url: " + base_url)
@@ -137,9 +137,7 @@ def download_jar(
         if version == "latest":
             version = None
         if version is None:
-            metadata_content = check_output(
-                "curl -s {}/maven-metadata.xml".format(base_url)
-            )
+            metadata_content = check_output("curl -s {}/maven-metadata.xml".format(base_url))
             meta_tree = ET.fromstring(metadata_content)
             versions = [
                 node.text
@@ -152,11 +150,13 @@ def download_jar(
                 )
             ]
             version = versions[-1]
-        jar_url = "{base_url}/{version}/spark_{jar_type}_{scala_version}-{version}-assembly.jar".format(
-            base_url=base_url,
-            version=version,
-            scala_version=scala_version,
-            jar_type=jar_type,
+        jar_url = (
+            "{base_url}/{version}/spark_{jar_type}_{scala_version}-{version}-assembly.jar".format(
+                base_url=base_url,
+                version=version,
+                scala_version=scala_version,
+                jar_type=jar_type,
+            )
         )
         jar_path = os.path.join("/tmp", extract_filename_from_path(jar_url))
         download_only_once(jar_url, jar_path, skip_download)
@@ -170,6 +170,7 @@ def get_teams_json_file_path(repo_path):
 def get_teams_py_file_path(repo_path):
     return os.path.join(repo_path, "teams.py")
 
+
 def set_runtime_env_v3(params, conf):
     effective_mode = params.get("mode")
 
@@ -181,9 +182,14 @@ def set_runtime_env_v3(params, conf):
         if os.path.isfile(conf_path):
             with open(conf_path, "r") as infile:
                 conf_json = json.load(infile)
-                metadata = conf_json.get("metaData", {}) or conf_json # user may just pass metadata as the entire json
+                metadata = (
+                    conf_json.get("metaData", {}) or conf_json
+                )  # user may just pass metadata as the entire json
                 env = metadata.get("executionInfo", {}).get("env", {})
-                runtime_env.update(env.get(EnvOrConfigAttribute.ENV,{}).get(effective_mode,{}) or env.get("common", {}))
+                runtime_env.update(
+                    env.get(EnvOrConfigAttribute.ENV, {}).get(effective_mode, {})
+                    or env.get("common", {})
+                )
                 # Also set APP_NAME
                 try:
                     _, conf_type, team, _ = conf.split("/")[-4:]
@@ -206,23 +212,14 @@ def set_runtime_env_v3(params, conf):
                 except Exception:
                     LOG.warn(
                         "Failed to set APP_NAME due to invalid conf path: {}, please ensure to supply the "
-                        "relative path to zipline/ folder".format(
-                           conf
-                        )
+                        "relative path to zipline/ folder".format(conf)
                     )
         else:
             if not params.get("app_name") and not os.environ.get("APP_NAME"):
                 # Provide basic app_name when no conf is defined.
                 # Modes like metadata-upload and metadata-export can rely on conf-type or folder rather than a conf.
                 runtime_env["APP_NAME"] = "_".join(
-                    [
-                        k
-                        for k in [
-                            "chronon",
-                            effective_mode.replace("-", "_")
-                        ]
-                        if k is not None
-                    ]
+                    [k for k in ["chronon", effective_mode.replace("-", "_")] if k is not None]
                 )
     for key, value in runtime_env.items():
         if key not in os.environ and value is not None:
@@ -230,6 +227,7 @@ def set_runtime_env_v3(params, conf):
             print(f"Setting to environment: {key}={value}")
             os.environ[key] = value
 
+
 # TODO: delete this when we cutover
 def set_runtime_env(params):
     """
@@ -263,20 +261,15 @@ def set_runtime_env(params):
     if effective_mode and "streaming" in effective_mode:
         effective_mode = "streaming"
     if params["repo"]:
-
         # Break if teams.json and teams.py exists
         teams_json_file = get_teams_json_file_path(params["repo"])
         teams_py_file = get_teams_py_file_path(params["repo"])
 
         if os.path.exists(teams_json_file) and os.path.exists(teams_py_file):
-            raise ValueError(
-                "Both teams.json and teams.py exist. Please only use teams.py."
-            )
+            raise ValueError("Both teams.json and teams.py exist. Please only use teams.py.")
 
         if os.path.exists(teams_json_file):
-            set_runtime_env_teams_json(
-                environment, params, effective_mode, teams_json_file
-            )
+            set_runtime_env_teams_json(environment, params, effective_mode, teams_json_file)
         if params["app_name"]:
             environment["cli_args"]["APP_NAME"] = params["app_name"]
         else:
@@ -289,11 +282,7 @@ def set_runtime_env(params):
                         for k in [
                             "chronon",
                             conf_type,
-                            (
-                                params["mode"].replace("-", "_")
-                                if params["mode"]
-                                else None
-                            ),
+                            (params["mode"].replace("-", "_") if params["mode"] else None),
                         ]
                         if k is not None
                     ]
@@ -321,6 +310,7 @@ def set_runtime_env(params):
                     LOG.info(f"From <{set_key}> setting {key}={value}")
                     os.environ[key] = value
 
+
 # TODO: delete this when we cutover
 def set_runtime_env_teams_json(environment, params, effective_mode, teams_json_file):
     if os.path.exists(teams_json_file):
@@ -346,9 +336,7 @@ def set_runtime_env_teams_json(environment, params, effective_mode, teams_json_f
                 context = params["env"]
             else:
                 context = "dev"
-            LOG.info(
-                f"Context: {context} -- conf_type: {conf_type} -- team: {team}"
-            )
+            LOG.info(f"Context: {context} -- conf_type: {conf_type} -- team: {team}")
             conf_path = os.path.join(params["repo"], params["conf"])
             if os.path.isfile(conf_path):
                 with open(conf_path, "r") as conf_file:
@@ -362,9 +350,7 @@ def set_runtime_env_teams_json(environment, params, effective_mode, teams_json_f
                     )
 
                     old_env = (
-                        conf_json.get("metaData")
-                        .get("modeToEnvMap", {})
-                        .get(effective_mode, {})
+                        conf_json.get("metaData").get("modeToEnvMap", {}).get(effective_mode, {})
                     )
 
                     environment["conf_env"] = new_env if new_env else old_env
@@ -375,8 +361,8 @@ def set_runtime_env_teams_json(environment, params, effective_mode, teams_json_f
                         "backfill-left",
                         "backfill-final",
                     ]:
-                        environment["conf_env"]["CHRONON_CONFIG_ADDITIONAL_ARGS"] = (
-                            " ".join(custom_json(conf_json).get("additional_args", []))
+                        environment["conf_env"]["CHRONON_CONFIG_ADDITIONAL_ARGS"] = " ".join(
+                            custom_json(conf_json).get("additional_args", [])
                         )
                     environment["cli_args"]["APP_NAME"] = APP_NAME_TEMPLATE.format(
                         mode=effective_mode,
@@ -384,18 +370,14 @@ def set_runtime_env_teams_json(environment, params, effective_mode, teams_json_f
                         context=context,
                         name=conf_json["metaData"]["name"],
                     )
-                environment["team_env"] = (
-                    teams_json[team].get(context, {}).get(effective_mode, {})
-                )
+                environment["team_env"] = teams_json[team].get(context, {}).get(effective_mode, {})
                 # fall-back to prod env even in dev mode when dev env is undefined.
                 environment["production_team_env"] = (
                     teams_json[team].get("production", {}).get(effective_mode, {})
                 )
                 # By default use production env.
                 environment["default_env"] = (
-                    teams_json.get("default", {})
-                    .get("production", {})
-                    .get(effective_mode, {})
+                    teams_json.get("default", {}).get("production", {}).get(effective_mode, {})
                 )
                 environment["cli_args"]["CHRONON_CONF_PATH"] = conf_path
     if params["app_name"]:
@@ -444,9 +426,7 @@ def split_date_range(start_date, end_date, parallelism):
     end_date = datetime.strptime(end_date, "%Y-%m-%d")
     if start_date > end_date:
         raise ValueError("Start date should be earlier than end date")
-    total_days = (
-        end_date - start_date
-    ).days + 1  # +1 to include the end_date in the range
+    total_days = (end_date - start_date).days + 1  # +1 to include the end_date in the range
 
     # Check if parallelism is greater than total_days
     if parallelism > total_days:
@@ -461,12 +441,41 @@ def split_date_range(start_date, end_date, parallelism):
             split_end = end_date
         else:
             split_end = split_start + timedelta(days=split_size - 1)
-        date_ranges.append(
-            (split_start.strftime("%Y-%m-%d"), split_end.strftime("%Y-%m-%d"))
-        )
+        date_ranges.append((split_start.strftime("%Y-%m-%d"), split_end.strftime("%Y-%m-%d")))
     return date_ranges
 
+
 def get_metadata_name_from_conf(repo_path, conf_path):
     with open(os.path.join(repo_path, conf_path), "r") as conf_file:
         data = json.load(conf_file)
-        return data.get("metaData", {}).get("name", None)
+        return data.get("metaData", {}).get("name", None)
+
+
+def handle_conf_not_found(log_error=True, callback=None):
+    def wrapper(func):
+        @functools.wraps(func)
+        def wrapped(*args, **kwargs):
+            try:
+                return func(*args, **kwargs)
+            except FileNotFoundError as e:
+                if log_error:
+                    print(style(f"File not found in {func.__name__}: {e}", fg="red"))
+                if callback:
+                    callback(*args, **kwargs)
+                return
+
+        return wrapped
+
+    return wrapper
+
+
+def print_possible_confs(conf, repo, *args, **kwargs):
+    conf_location = os.path.join(repo, conf)
+    conf_dirname = os.path.dirname(conf_location)
+    if os.path.exists(conf_dirname):
+        print(
+            f"Possible confs from {style(conf_dirname, fg='yellow')}: \n -",
+            "\n - ".join([name for name in os.listdir(conf_dirname)]),
+        )
+    else:
+        print(f"Directory does not exist: {style(conf_dirname, fg='yellow')}")

ai/chronon/repo/zipline.py

@@ -9,24 +9,6 @@ from ai.chronon.repo.hub_runner import hub
 from ai.chronon.repo.init import main as init_main
 from ai.chronon.repo.run import main as run_main
 
-LOGO = """
-     =%%%@:-%%%@=:%%%@+     .%@%@@@@@@%%%%%%: .+#%*.                     -%%%=  -#%#-
-    :@@@@#.@@@@%.%@@@@.     .@@@@@@@@@@@@@@-  -@@@@=                     =@@@+  @@@@@
-   :@@@@*.%@@@#.#@@@%.              .#@@@@:    :==:                      =@@@+   -=-   :
-  =@@@@=-@@@@+:%@@@#.               #@@@%.     :--:  .%%=:+#%@@@%#+-     =@@@+  .-:-.  *%=         #%%*     :=#%@@@@#*-
-.#@@@#-+@@@%-=@@@@-               .%@@@%.      @@@@  .@@@@@@@@%%@@@@%=   =@@@+  +@@@=  *@@@+.      %@@%   :#@@@@%%%@@@@@=
-+**+=-%@@@+-#@@@*----=.          :@@@@#        %@@@  .@@@@%=.   .-#@@@*  =@@@+  +@@@=  *@@@@@*:    %@@%  -@@@%-     .+@@@*
-    +@@@%-+@@@%-=@@@@+          :@@@@*         @@@@  .@@@@.        #@@@: =@@@+  +@@@=  *@@@%@@@*:  %@@%  %@@@#++****+*@@@@-
-  -@@@@+:#@@@*:#@@@#.          -@@@@*          @@@@  .@@@@         *@@@- =@@@+  +@@@=  *@@@.-%@@@#-%@@%  @@@@****#****++++:
- =@@@@--@@@@=:@@@@*           =@@@@+           @@@@  .@@@@#.     .+@@@%  =@@@+  +@@@=  *@@@   -#@@@@@@%  =@@@*.      
-+@@@@--@@@@=:@@@@*           +@@@@@#########+  @@@@  .@@@@@@%*+*#@@@@*   =@@@+  +@@@=  *@@@.    :#@@@@%   =@@@@%     -==+-
-:@@@@* @@@@# @@@@%           *@@@@@@@@@@@@@@@%  @@@@  .@@@@#@@@@@@@%+:    =@@@+  +@@@=  *@@@.      :*@@%    .=#@@@@@@@%*:
-                                                         .@@@%  
-                                                         .@@@%
-                                                         .@@@@
-                                                          ---:
-"""
-
 
 def _set_package_version():
     try:
@@ -37,7 +19,9 @@ def _set_package_version():
     return package_version
 
 
-@click.group(help="The Zipline CLI. A tool for authoring and running Zipline pipelines in the cloud. For more information, see: https://chronon.ai/")
+@click.group(
+    help="The Zipline CLI. A tool for compiling and running Zipline pipelines. For more information, see: https://zipline.ai/docs"
+)
 @click.version_option(version=_set_package_version())
 @click.pass_context
 def zipline(ctx):

ai/chronon/repo/zipline_hub.py

@@ -1,105 +1,277 @@
+import json
 import os
+from datetime import date, datetime, timedelta, timezone
 from typing import Optional
 
 import google.auth
 import requests
 from google.auth.transport.requests import Request
+from google.cloud import iam_credentials_v1
 
 
 class ZiplineHub:
-    def __init__(self, base_url):
+    def __init__(self, base_url, sa_name=None):
         if not base_url:
             raise ValueError("Base URL for ZiplineHub cannot be empty.")
         self.base_url = base_url
-        if self.base_url.startswith("https") and self.base_url.endswith(".app"):
+        if self.base_url.startswith("https"):
             print("\n 🔐 Using Google Cloud authentication for ZiplineHub.")
 
             # First try to get ID token from environment (GitHub Actions)
-            self.id_token = os.getenv('GCP_ID_TOKEN')
+            self.id_token = os.getenv("GCP_ID_TOKEN")
             if self.id_token:
                 print(" 🔑 Using ID token from environment")
-            else:
+            elif sa_name is not None:
                 # Fallback to Google Cloud authentication
+                print(" 🔑 Generating ID token from service account credentials")
+                credentials, project_id = google.auth.default()
+                self.project_id = project_id
+                credentials.refresh(Request())
+
+                self.sa = f"{sa_name}@{project_id}.iam.gserviceaccount.com"
+            else:
                 print(" 🔑 Generating ID token from default credentials")
                 credentials, project_id = google.auth.default()
                 credentials.refresh(Request())
+                self.sa = None
                 self.id_token = credentials.id_token
 
-    def call_diff_api(self, names_to_hashes: dict[str, str]) -> Optional[list[str]]:
-        url = f"{self.base_url}/upload/v1/diff"
+    def _generate_jwt_payload(self, service_account_email: str, resource_url: str) -> str:
+        """Generates JWT payload for service account.
 
-        diff_request = {
-            'namesToHashes': names_to_hashes
+        Creates a properly formatted JWT payload with standard claims (iss, sub, aud,
+        iat, exp) needed for IAP authentication.
+
+        Args:
+            service_account_email (str): Specifies service account JWT is created for.
+            resource_url (str): Specifies scope of the JWT, the URL that the JWT will
+                be allowed to access.
+
+        Returns:
+            str: JSON string containing the JWT payload with properly formatted claims.
+        """
+        # Create current time and expiration time (1 hour later) in UTC
+        iat = datetime.now(tz=timezone.utc)
+        exp = iat + timedelta(seconds=3600)
+
+        # Convert datetime objects to numeric timestamps (seconds since epoch)
+        # as required by JWT standard (RFC 7519)
+        payload = {
+            "iss": service_account_email,
+            "sub": service_account_email,
+            "aud": resource_url,
+            "iat": int(iat.timestamp()),
+            "exp": int(exp.timestamp()),
         }
-        headers = {'Content-Type': 'application/json'}
-        if hasattr(self, 'id_token'):
-            headers['Authorization'] = f'Bearer {self.id_token}'
+
+        return json.dumps(payload)
+
+    def _sign_jwt(self, target_sa: str, resource_url: str) -> str:
+        """Signs JWT payload using ADC and IAM credentials API.
+
+        Uses Google Cloud's IAM Credentials API to sign a JWT. This requires the
+        caller to have iap.webServiceVersions.accessViaIap permission on the target
+        service account.
+
+        Args:
+            target_sa (str): Service Account JWT is being created for.
+                iap.webServiceVersions.accessViaIap permission is required.
+            resource_url (str): Audience of the JWT, and scope of the JWT token.
+                This is the url of the IAP protected application.
+
+        Returns:
+            str: A signed JWT that can be used to access IAP protected apps.
+                Use in Authorization header as: 'Bearer <signed_jwt>'
+        """
+        # Get default credentials from environment or application credentials
+        source_credentials, project_id = google.auth.default()
+
+        # Initialize IAM credentials client with source credentials
+        iam_client = iam_credentials_v1.IAMCredentialsClient(credentials=source_credentials)
+
+        # Generate the service account resource name
+        # Use '-' as placeholder as per API requirements
+        name = iam_client.service_account_path("-", target_sa)
+
+        # Create and sign the JWT payload
+        payload = self._generate_jwt_payload(target_sa, resource_url)
+
+        request = iam_credentials_v1.SignJwtRequest(
+            name=name,
+            payload=payload,
+        )
+        # Sign the JWT using the IAM credentials API
+        response = iam_client.sign_jwt(request=request)
+
+        return response.signed_jwt
+
+    def call_diff_api(self, names_to_hashes: dict[str, str]) -> Optional[list[str]]:
+        url = f"{self.base_url}/upload/v2/diff"
+
+        diff_request = {"namesToHashes": names_to_hashes}
+        headers = {"Content-Type": "application/json"}
+        if self.base_url.startswith("https") and hasattr(self, "sa") and self.sa is not None:
+            headers["Authorization"] = f"Bearer {self._sign_jwt(self.sa, url)}"
+        elif self.base_url.startswith("https"):
+            headers["Authorization"] = f"Bearer {self.id_token}"
         try:
             response = requests.post(url, json=diff_request, headers=headers)
             response.raise_for_status()
             diff_response = response.json()
-            return diff_response['diff']
+            return diff_response["diff"]
         except requests.RequestException as e:
-            print(f" ❌ Error calling diff API: {e}")
+            if e.response is not None and e.response.status_code == 401 and self.sa is None:
+                print(
+                    " ❌  Error calling diff API. Unauthorized and no service account provided. Make sure the environment has default credentials set up or provide a service account name as SA_NAME in teams.py."
+                )
+            elif e.response is not None and e.response.status_code == 401 and self.sa is not None:
+                print(
+                    f" ❌  Error calling diff API. Unauthorized with provided service account: {self.sa}. Make sure the service account has the 'iap.webServiceVersions.accessViaIap' permission."
+                )
+            else:
+                print(f" ❌ Error calling diff API: {e}")
             raise e
 
     def call_upload_api(self, diff_confs, branch: str):
-        url = f"{self.base_url}/upload/v1/confs"
+        url = f"{self.base_url}/upload/v2/confs"
 
         upload_request = {
-            'diffConfs': diff_confs,
-            'branch': branch,
+            "diffConfs": diff_confs,
+            "branch": branch,
         }
-        headers = {'Content-Type': 'application/json'}
-        if hasattr(self, 'id_token'):
-            headers['Authorization'] = f'Bearer {self.id_token}'
+        headers = {"Content-Type": "application/json"}
+        if self.base_url.startswith("https") and hasattr(self, "sa") and self.sa is not None:
+            headers["Authorization"] = f"Bearer {self._sign_jwt(self.sa, url)}"
+        elif self.base_url.startswith("https"):
+            headers["Authorization"] = f"Bearer {self.id_token}"
 
         try:
             response = requests.post(url, json=upload_request, headers=headers)
             response.raise_for_status()
             return response.json()
         except requests.RequestException as e:
-            print(f" ❌ Error calling upload API: {e}")
+            if e.response is not None and e.response.status_code == 401 and self.sa is None:
+                print(
+                    " ❌  Error calling upload API. Unauthorized and no service account provided. Make sure the environment has default credentials set up or provide a service account name as SA_NAME in teams.py."
+                )
+            elif e.response is not None and e.response.status_code == 401 and self.sa is not None:
+                print(
+                    f" ❌  Error calling upload API. Unauthorized with provided service account: {self.sa}. Make sure the service account has the 'iap.webServiceVersions.accessViaIap' permission."
+                )
+            else:
+                print(f" ❌ Error calling upload API: {e}")
+            raise e
+
+    def call_schedule_api(self, modes, branch, conf_name, conf_hash):
+        url = f"{self.base_url}/schedule/v2/schedules"
+
+        schedule_request = {
+            "modeSchedules": modes,
+            "branch": branch,
+            "confName": conf_name,
+            "confHash": conf_hash,
+        }
+
+        headers = {"Content-Type": "application/json"}
+        if self.base_url.startswith("https") and hasattr(self, "sa") and self.sa is not None:
+            headers["Authorization"] = f"Bearer {self._sign_jwt(self.sa, url)}"
+        elif self.base_url.startswith("https"):
+            headers["Authorization"] = f"Bearer {self.id_token}"
+
+        try:
+            response = requests.post(url, json=schedule_request, headers=headers)
+            response.raise_for_status()
+            return response.json()
+        except requests.RequestException as e:
+            if e.response is not None and e.response.status_code == 401 and self.sa is None:
+                print(
+                    " ❌  Error deploying schedule. Unauthorized and no service account provided. Make sure the environment has default credentials set up or provide a service account name as SA_NAME in teams.py."
+                )
+            elif e.response is not None and e.response.status_code == 401 and self.sa is not None:
+                print(
+                    f" ❌  Error deploying schedule. Unauthorized with provided service account: {self.sa}. Make sure the service account has the 'iap.webServiceVersions.accessViaIap' permission."
+                )
+            else:
+                print(f" ❌ Error deploying schedule: {e}")
             raise e
 
     def call_sync_api(self, branch: str, names_to_hashes: dict[str, str]) -> Optional[list[str]]:
-        url = f"{self.base_url}/upload/v1/sync"
+        url = f"{self.base_url}/upload/v2/sync"
 
         sync_request = {
             "namesToHashes": names_to_hashes,
             "branch": branch,
         }
-        headers = {'Content-Type': 'application/json'}
-        if hasattr(self, 'id_token'):
-            headers['Authorization'] = f'Bearer {self.id_token}'
+        headers = {"Content-Type": "application/json"}
+        if self.base_url.startswith("https") and hasattr(self, "sa") and self.sa is not None:
+            headers["Authorization"] = f"Bearer {self._sign_jwt(self.sa, url)}"
+        elif self.base_url.startswith("https"):
+            headers["Authorization"] = f"Bearer {self.id_token}"
+
         try:
             response = requests.post(url, json=sync_request, headers=headers)
             response.raise_for_status()
             return response.json()
         except requests.RequestException as e:
-            print(f" ❌ Error calling diff API: {e}")
+            if e.response is not None and e.response.status_code == 401 and self.sa is None:
+                print(
+                    " ❌  Error calling sync API. Unauthorized and no service account provided. Make sure the environment has default credentials set up or provide a service account name as SA_NAME in teams.py."
+                )
+            elif e.response is not None and e.response.status_code == 401 and self.sa is not None:
+                print(
+                    f" ❌  Error calling sync API. Unauthorized with provided service account: {self.sa}. Make sure the service account has the 'iap.webServiceVersions.accessViaIap' permission."
+                )
+            else:
+                print(f" ❌ Error calling sync API: {e}")
             raise e
 
-    def call_workflow_start_api(self, conf_name, mode, branch, user, start, end, conf_hash):
-        url = f"{self.base_url}/workflow/start"
-
+    def call_workflow_start_api(
+        self,
+        conf_name,
+        mode,
+        branch,
+        user,
+        conf_hash,
+        start=None,
+        end=None,
+        skip_long_running=False,
+    ):
+        url = f"{self.base_url}/workflow/v2/start"
+        end_dt = end.strftime("%Y-%m-%d") if end else date.today().strftime("%Y-%m-%d")
+        start_dt = (
+            start.strftime("%Y-%m-%d")
+            if start
+            else (date.today() - timedelta(days=14)).strftime("%Y-%m-%d")
+        )
         workflow_request = {
-            'confName': conf_name,
-            'confHash': conf_hash,
-            'mode': mode,
-            'branch': branch,
-            'user': user,
-            'start': start,
-            'end': end,
+            "confName": conf_name,
+            "confHash": conf_hash,
+            "mode": mode,
+            "branch": branch,
+            "user": user,
+            "start": start_dt,
+            "end": end_dt,
+            "skipLongRunningNodes": skip_long_running,
         }
-        headers = {'Content-Type': 'application/json'}
-        if hasattr(self, 'id_token'):
-            headers['Authorization'] = f'Bearer {self.id_token}'
+        headers = {"Content-Type": "application/json"}
+        if self.base_url.startswith("https") and hasattr(self, "sa") and self.sa is not None:
+            headers["Authorization"] = f"Bearer {self._sign_jwt(self.sa, url)}"
+        elif self.base_url.startswith("https"):
+            headers["Authorization"] = f"Bearer {self.id_token}"
 
         try:
             response = requests.post(url, json=workflow_request, headers=headers)
             response.raise_for_status()
             return response.json()
         except requests.RequestException as e:
-            print(f" ❌ Error calling workflow start API: {e}")
+            if e.response is not None and e.response.status_code == 401 and self.sa is None:
+                print(
+                    " ❌  Error calling workflow start API. Unauthorized and no service account provided. Make sure the environment has default credentials set up or provide a service account name as SA_NAME in teams.py."
+                )
+            elif e.response is not None and e.response.status_code == 401 and self.sa is not None:
+                print(
+                    f" ❌  Error calling workflow start API. Unauthorized with provided service account: {self.sa}. Make sure the service account has the 'iap.webServiceVersions.accessViaIap' permission."
+                )
+            else:
+                print(f" ❌ Error calling workflow start API: {e}")
             raise e