fix: get saas token

core/morph/task/utils/connection.py

@@ -719,7 +719,7 @@ def find_builtin_db_connection() -> tuple[str, PostgresqlConnection]:
                 host = match.group("host")
                 database = match.group("database")
 
-                return client.req.workspace_id, PostgresqlConnection(
+                return client.req.project_id, PostgresqlConnection(
                     type=CONNECTION_TYPE.postgres,
                     host=host,
                     user=username,
@@ -735,6 +735,9 @@ def find_cloud_connection(
     ) -> Connection:
         if connection_slug == MORPH_DUCKDB_CONNECTION_SLUG:
             return DuckDBConnection(type=CONNECTION_TYPE.duckdb)
+        elif connection_slug == MORPH_BUILTIN_DB_CONNECTION_SLUG:
+            _, builtin_connection = ConnectionYaml.find_builtin_db_connection()
+            return builtin_connection
         try:
             try:
                 client = MorphApiClient(MorphApiKeyClientImpl)

core/morph_lib/api.py

@@ -1,52 +1,29 @@
 from __future__ import annotations
 
-import os
-from typing import Any, Dict, cast
+from typing import cast
 
-import requests
 import urllib3
 from morph_lib.error import MorphApiError
-from pydantic import BaseModel
 from urllib3.exceptions import InsecureRequestWarning
 
-urllib3.disable_warnings(InsecureRequestWarning)
-
+from morph.api.cloud.client import MorphApiClient, MorphApiKeyClientImpl
+from morph.task.utils.connection import (
+    AirtableConnectionOAuth,
+    AttioConnectionOAuth,
+    BigqueryConnectionOAuth,
+    ConnectionYaml,
+    FreeeConnectionOAuth,
+    GoogleAnalyticsConnectionOAuth,
+    HubspotConnectionOAuth,
+    IntercomConnectionOAuth,
+    LinearConnectionOAuth,
+    MailchimpConnectionOAuth,
+    NotionConnectionOAuth,
+    SalesforceConnectionOAuth,
+    StripeConnectionOAuth,
+)
 
-# ===============================================
-#
-# Implementation
-#
-# ===============================================
-class EnvVars(BaseModel):
-    workspace_id: str
-    base_url: str
-    team_slug: str
-    api_key: str
-
-
-def _canonicalize_base_url(base_url: str) -> str:
-    if base_url.startswith("http"):
-        return base_url
-    else:
-        return f"https://{base_url}"
-
-
-def _read_configuration_from_env() -> EnvVars:
-    """
-    Read configuration from environment variables
-    These variables are loaded from ini file when the morph run command is executed
-    """
-    workspace_id_value = os.getenv("MORPH_PROJECT_ID", "")
-    base_url_value = os.getenv("MORPH_BASE_URL", "")
-    team_slug_value = os.getenv("MORPH_TEAM_SLUG", "")
-    api_key_value = os.getenv("MORPH_API_KEY", "")
-
-    return EnvVars(
-        workspace_id=workspace_id_value,
-        base_url=base_url_value,
-        team_slug=team_slug_value,
-        api_key=api_key_value,
-    )
+urllib3.disable_warnings(InsecureRequestWarning)
 
 
 # ===============================================
@@ -62,37 +39,38 @@ def get_auth_token(connection_slug: str) -> str:
     Make sure to set the environment variables before calling this function.
     @param: connection_slug: The connection slug on morph app
     """
-    config_from_env = _read_configuration_from_env()
-    base_url = config_from_env.base_url
-    team_slug = config_from_env.team_slug
-    api_key = config_from_env.api_key
-
-    headers = {
-        "teamSlug": team_slug,
-        "X-Api-Key": api_key,
-    }
-
-    url = f"{_canonicalize_base_url(base_url)}/external-connection/{connection_slug}"
-
-    try:
-        response = requests.get(url, headers=headers, verify=True)
-    except Exception as e:
-        raise MorphApiError(f"get_auth_token error: {e}")
+    connection_yaml = ConnectionYaml.load_yaml()
+    database_connection = ConnectionYaml.find_connection(
+        connection_yaml, connection_slug
+    )
 
-    if response.status_code != 200:
-        raise MorphApiError(response.text)
-    response_json: Dict[str, Any] = response.json()
-    if (
-        "error" in response_json
-        and "subCode" in response_json
-        and "message" in response_json
+    if database_connection is not None and (
+        isinstance(database_connection, BigqueryConnectionOAuth)
+        or isinstance(database_connection, GoogleAnalyticsConnectionOAuth)
+        or isinstance(database_connection, SalesforceConnectionOAuth)
+        or isinstance(database_connection, NotionConnectionOAuth)
+        or isinstance(database_connection, StripeConnectionOAuth)
+        or isinstance(database_connection, AttioConnectionOAuth)
+        or isinstance(database_connection, AirtableConnectionOAuth)
+        or isinstance(database_connection, FreeeConnectionOAuth)
+        or isinstance(database_connection, HubspotConnectionOAuth)
+        or isinstance(database_connection, IntercomConnectionOAuth)
+        or isinstance(database_connection, LinearConnectionOAuth)
+        or isinstance(database_connection, MailchimpConnectionOAuth)
     ):
-        raise MorphApiError(response_json["message"])
+        return database_connection.access_token or ""
+
+    client = MorphApiClient(MorphApiKeyClientImpl)
+    response = client.req.find_external_connection(connection_slug)
+    if response.is_error():
+        raise MorphApiError(f"Failed to get auth token. {response.text}")
 
+    response_json = response.json()
     if (
         response_json["connectionType"] == "mysql"
         or response_json["connectionType"] == "postgres"
         or response_json["connectionType"] == "redshift"
+        or response_json["connectionType"] == "mssql"
     ):
         raise MorphApiError(f"No auth token in db connection {connection_slug}")
     elif (