mage-ai 0.9.44__py3-none-any.whl → 0.9.46__py3-none-any.whl

mage_ai/api/resources/OauthResource.py

@@ -1,22 +1,13 @@
-import json
-import os
 import urllib.parse
-import uuid
 from datetime import datetime, timedelta
 from urllib.parse import parse_qs, urlparse
 
-import aiohttp
-
 from mage_ai.api.errors import ApiError
 from mage_ai.api.resources.GenericResource import GenericResource
 from mage_ai.authentication.oauth2 import generate_access_token
 from mage_ai.authentication.oauth.constants import (
-    ACTIVE_DIRECTORY_CLIENT_ID,
-    GHE_CLIENT_ID_ENV_VAR,
-    GHE_CLIENT_SECRET_ENV_VAR,
     GITHUB_CLIENT_ID,
     GITHUB_STATE,
-    OAUTH_PROVIDER_ACTIVE_DIRECTORY,
     OAUTH_PROVIDER_GHE,
     OAUTH_PROVIDER_GITHUB,
     VALID_OAUTH_PROVIDERS,
@@ -26,13 +17,35 @@ from mage_ai.authentication.oauth.utils import (
     access_tokens_for_client,
     add_access_token_to_query,
 )
+from mage_ai.authentication.providers.constants import NAME_TO_PROVIDER
 from mage_ai.data_preparation.git.api import get_oauth_client_id
 from mage_ai.orchestration.db import safe_db_query
 from mage_ai.orchestration.db.models.oauth import Oauth2AccessToken, Oauth2Application
-from mage_ai.settings import ACTIVE_DIRECTORY_DIRECTORY_ID, ROUTES_BASE_PATH
 
 
 class OauthResource(GenericResource):
+    @classmethod
+    @safe_db_query
+    async def collection(self, query, meta, user, **kwargs):
+        redirect_uri = query.get('redirect_uri', [None])
+        if redirect_uri:
+            redirect_uri = redirect_uri[0]
+
+        oauths = []
+        for provider, provider_class in NAME_TO_PROVIDER.items():
+            try:
+                provider_instance = provider_class()
+                auth_url_response = provider_instance.get_auth_url_response(
+                    redirect_uri=redirect_uri
+                )
+                if auth_url_response:
+                    auth_url_response['provider'] = provider
+                    oauths.append(auth_url_response)
+            except Exception:
+                continue
+
+        return self.build_result_set(oauths, user, **kwargs)
+
     @classmethod
     @safe_db_query
     def create(self, payload, user, **kwargs):
@@ -107,6 +120,12 @@ class OauthResource(GenericResource):
             provider = OAUTH_PROVIDER_GHE
         else:
             provider = pk
+
+        provider_class = NAME_TO_PROVIDER.get(provider)
+        provider_instance = None
+        if provider_class is not None:
+            provider_instance = provider_class()
+
         access_tokens = access_tokens_for_client(
             get_oauth_client_id(provider),
             user=user,
@@ -121,38 +140,22 @@ class OauthResource(GenericResource):
         # If an oauth code is provided, we need to exchange it for an access token for
         # the provider and return the redirect uri.
         elif code:
-            if OAUTH_PROVIDER_GHE == pk:
-                # Fetch access token for GitHub Enterprise and add it to the redirect URI.
-
-                # 1. Get GHE client credentials from environment variables.
-                # 2. Obtain an access token from the GitHub Enterprise OAuth service.
-                # 3. Update uri query parameters with provider and access token from GHE API.
-                # 4. Recreate and return the uri to include the access token and provider.
+            if provider_instance is not None:
+                # 1. Obtain an access token from the Oauth provider.
+                # 2. Update uri query parameters with provider and access token from response.
+                # 3. Recreate and return the uri to include the access token and provider.
                 parsed_url = urlparse(urllib.parse.unquote(redirect_uri))
                 parsed_url_query = parse_qs(parsed_url.query)
 
-                query = {'provider': pk}
+                query = {'provider': provider}
                 for k, v in parsed_url_query.items():
-                    if type(v) is list:
+                    if isinstance(v, list):
                         v = ','.join(v)
                     query[k] = v
 
-                data = dict()
-                async with aiohttp.ClientSession() as session:
-                    async with session.post(
-                        f'{ghe_hostname}/login/oauth/access_token',
-                        headers={
-                            'Accept': 'application/json',
-                        },
-                        data=dict(
-                            client_id=os.getenv(GHE_CLIENT_ID_ENV_VAR),
-                            client_secret=os.getenv(GHE_CLIENT_SECRET_ENV_VAR),
-                            code=code,
-                        ),
-                        timeout=20,
-                    ) as response:
-                        data = await response.json()
-
+                data = await provider_instance.get_access_token_response(
+                    code, redirect_uri=redirect_uri
+                )
                 query = add_access_token_to_query(data, query)
 
                 parts = redirect_uri.split('?')
@@ -168,74 +171,33 @@ class OauthResource(GenericResource):
                 model['url'] = redirect_uri_final
         # Otherwise, return the authorization url to start the oauth flow.
         else:
-            if OAUTH_PROVIDER_GITHUB == pk:
-                if ghe_hostname:
-                    parsed_url = urlparse(urllib.parse.unquote(redirect_uri))
-                    base_url = parsed_url.scheme + '://' + parsed_url.netloc
-                    if ROUTES_BASE_PATH:
-                        base_url += f'/{ROUTES_BASE_PATH}'
-                    redirect_uri_query = urllib.parse.urlencode(
-                        dict(
-                            provider=OAUTH_PROVIDER_GHE,
-                            redirect_uri=redirect_uri,
+            if OAUTH_PROVIDER_GITHUB == provider:
+                query = dict(
+                    client_id=GITHUB_CLIENT_ID,
+                    redirect_uri=urllib.parse.quote_plus(
+                        '?'.join(
+                            [
+                                f'https://api.mage.ai/v1/oauth/{pk}',
+                                f'redirect_uri={urllib.parse.unquote(redirect_uri)}',
+                            ]
                         )
-                    )
-                    query = dict(
-                        client_id=os.getenv(GHE_CLIENT_ID_ENV_VAR),
-                        redirect_uri=urllib.parse.quote_plus(
-                            f'{base_url}/oauth?{redirect_uri_query}',
-                        ),
-                        scope='repo',
-                        state=uuid.uuid4().hex,
-                    )
-                    host = ghe_hostname
-                else:
-                    query = dict(
-                        client_id=GITHUB_CLIENT_ID,
-                        redirect_uri=urllib.parse.quote_plus(
-                            '?'.join(
-                                [
-                                    f'https://api.mage.ai/v1/oauth/{pk}',
-                                    f'redirect_uri={urllib.parse.unquote(redirect_uri)}',
-                                ]
-                            )
-                        ),
-                        scope='repo',
-                        state=GITHUB_STATE,
-                    )
-                    host = 'https://github.com'
+                    ),
+                    scope='repo',
+                    state=GITHUB_STATE,
+                )
                 query_strings = []
                 for k, v in query.items():
                     query_strings.append(f'{k}={v}')
 
-                model['url'] = f"{host}/login/oauth/authorize?{'&'.join(query_strings)}"
-            elif OAUTH_PROVIDER_ACTIVE_DIRECTORY == pk:
-                ad_directory_id = ACTIVE_DIRECTORY_DIRECTORY_ID
-                if ad_directory_id:
-                    from requests.models import PreparedRequest
-
-                    req = PreparedRequest()
-                    req.prepare_url(redirect_uri, dict(provider=pk))
-                    query = dict(
-                        client_id=ACTIVE_DIRECTORY_CLIENT_ID,
-                        redirect_uri=f'https://api.mage.ai/v1/oauth/{pk}',
-                        response_type='code',
-                        scope='User.Read',
-                        state=urllib.parse.quote_plus(
-                            json.dumps(
-                                dict(
-                                    redirect_uri=req.url,
-                                    tenant_id=ad_directory_id,
-                                )
-                            )
-                        ),
-                    )
-                    query_strings = []
-                    for k, v in query.items():
-                        query_strings.append(f'{k}={v}')
-                    model[
-                        'url'
-                    ] = f"https://login.microsoftonline.com/{ad_directory_id}/oauth2/v2.0/authorize?{'&'.join(query_strings)}"  # noqa: E501
+                model[
+                    'url'
+                ] = f"https://github.com/login/oauth/authorize?{'&'.join(query_strings)}"
+            elif provider_instance is not None:
+                resp = provider_instance.get_auth_url_response(
+                    redirect_uri=redirect_uri
+                )
+                if resp:
+                    model.update(resp)
 
         return self(model, user, **kwargs)
 

mage_ai/api/resources/PipelineResource.py

@@ -67,15 +67,15 @@ class PipelineResource(BaseResource):
             tags = new_tags
 
         pipeline_types = query.get('type[]', [])
-        if pipeline_types:
+        if pipeline_types and len(pipeline_types) == 1:
             pipeline_types = pipeline_types[0]
-        if pipeline_types:
+        if pipeline_types and isinstance(pipeline_types, str):
             pipeline_types = pipeline_types.split(',')
 
         pipeline_statuses = query.get('status[]', [])
-        if pipeline_statuses:
+        if pipeline_statuses and len(pipeline_statuses) == 1:
             pipeline_statuses = pipeline_statuses[0]
-        if pipeline_statuses:
+        if pipeline_statuses and isinstance(pipeline_statuses, str):
             pipeline_statuses = pipeline_statuses.split(',')
 
         from_history_days = query.get('from_history_days', [None])

mage_ai/api/resources/PipelineScheduleResource.py

@@ -3,6 +3,7 @@ import uuid
 from sqlalchemy.orm import selectinload
 
 from mage_ai.api.resources.DatabaseResource import DatabaseResource
+from mage_ai.data_preparation.models.pipeline import Pipeline
 from mage_ai.data_preparation.models.triggers import (
     Trigger,
     add_or_update_trigger_for_pipeline_and_persist,
@@ -146,6 +147,20 @@ class PipelineScheduleResource(DatabaseResource):
         if 'token' not in payload:
             payload['token'] = uuid.uuid4().hex
 
+        if pipeline.should_save_trigger_in_code_automatically():
+            def _callback(resource, kwargs=kwargs, user=user):
+                from mage_ai.api.resources.PipelineTriggerResource import (
+                    PipelineTriggerResource,
+                )
+
+                PipelineTriggerResource.create(
+                    dict(pipeline_schedule_id=resource.id),
+                    user,
+                    **kwargs,
+                )
+
+            self.on_create_callback = _callback
+
         return super().create(payload, user, **kwargs)
 
     @safe_db_query
@@ -257,22 +272,28 @@ class PipelineScheduleResource(DatabaseResource):
 
         resource = super().update(payload)
         updated_model = resource.model
-        trigger = Trigger(
-            name=updated_model.name,
-            pipeline_uuid=updated_model.pipeline_uuid,
-            schedule_interval=updated_model.schedule_interval,
-            schedule_type=updated_model.schedule_type,
-            settings=updated_model.settings,
-            sla=updated_model.sla,
-            start_time=updated_model.start_time,
-            status=updated_model.status,
-            variables=updated_model.variables,
-        )
-        add_or_update_trigger_for_pipeline_and_persist(
-            trigger,
-            updated_model.pipeline_uuid,
-            update_only_if_exists=True,
-        )
+
+        pipeline = Pipeline.get(updated_model.pipeline_uuid)
+        if pipeline:
+            trigger = Trigger(
+                name=updated_model.name,
+                pipeline_uuid=updated_model.pipeline_uuid,
+                schedule_interval=updated_model.schedule_interval,
+                schedule_type=updated_model.schedule_type,
+                settings=updated_model.settings,
+                sla=updated_model.sla,
+                start_time=updated_model.start_time,
+                status=updated_model.status,
+                variables=updated_model.variables,
+            )
+
+            update_only_if_exists = not pipeline.should_save_trigger_in_code_automatically()
+
+            add_or_update_trigger_for_pipeline_and_persist(
+                trigger,
+                pipeline.uuid,
+                update_only_if_exists=update_only_if_exists,
+            )
 
         return self
 

mage_ai/api/resources/ProjectResource.py

@@ -86,6 +86,9 @@ class ProjectResource(GenericResource):
         if 'emr_config' in payload:
             data['emr_config'] = payload['emr_config']
 
+        if 'pipelines' in payload:
+            data['pipelines'] = payload['pipelines']
+
         if 'spark_config' in payload:
             data['spark_config'] = payload['spark_config']
 
@@ -100,9 +103,8 @@ class ProjectResource(GenericResource):
         if should_log_project:
             await UsageStatisticLogger().project_impression()
 
-        if Project(repo_config=repo_config).is_feature_enabled(
-            FeatureUUID.DATA_INTEGRATION_IN_BATCH_PIPELINE,
-        ):
+        project = Project(repo_config=repo_config)
+        if project.is_feature_enabled(FeatureUUID.DATA_INTEGRATION_IN_BATCH_PIPELINE):
             await BlockActionObjectCache.initialize_cache(replace=True)
 
         return self

mage_ai/api/resources/SessionResource.py

@@ -4,9 +4,8 @@ from mage_ai.api.errors import ApiError
 from mage_ai.api.resources.BaseResource import BaseResource
 from mage_ai.authentication.ldap import new_ldap_connection
 from mage_ai.authentication.oauth2 import encode_token, generate_access_token
-from mage_ai.authentication.oauth.active_directory import get_user_info
-from mage_ai.authentication.oauth.constants import OAUTH_PROVIDER_ACTIVE_DIRECTORY
 from mage_ai.authentication.passwords import verify_password
+from mage_ai.authentication.providers.constants import NAME_TO_PROVIDER
 from mage_ai.orchestration.db import safe_db_query
 from mage_ai.orchestration.db.models.oauth import Role, User
 from mage_ai.settings import AUTHENTICATION_MODE, LDAP_DEFAULT_ACCESS
@@ -23,18 +22,34 @@ class SessionResource(BaseResource):
         token = payload.get('token')
         provider = payload.get('provider')
 
+        oauth_client = kwargs.get('oauth_client')
+
+        # Oauth sign in
         if token and provider:
-            if provider == OAUTH_PROVIDER_ACTIVE_DIRECTORY:
-                user_info = get_user_info(token)
-                principal_name = user_info.get('userPrincipalName')
-                user = User.query.filter(User.email == principal_name).first()
+            roles = []
+            provider_class = NAME_TO_PROVIDER.get(provider)
+            if provider_class is not None:
+                provider_instance = provider_class()
+                user_info = await provider_instance.get_user_info(access_token=token)
+                email = user_info.get('email')
+                username = user_info.get('username')
+                if 'user_roles' in user_info:
+                    roles = user_info.get('user_roles')
+
+            if not email:
+                error = ApiError.RESOURCE_NOT_FOUND
+                error.update({'message': 'Could not get email from oauth provider.'})
+                raise ApiError(error)
+            else:
+                user = User.query.filter(User.email == email).first()
                 if not user:  # noqa: E712
                     print('first user login, creating user.')
                     user = User.create(
-                        username=principal_name,
-                        email=principal_name,
+                        username=username,
+                        email=email,
+                        roles_new=roles,
                     )
-                oauth_token = generate_access_token(user, kwargs['oauth_client'])
+                oauth_token = generate_access_token(user, oauth_client)
                 return self(oauth_token, user, **kwargs)
 
         error = ApiError.RESOURCE_NOT_FOUND

mage_ai/api/resources/SparkApplicationResource.py

@@ -1,26 +1,26 @@
 from mage_ai.api.resources.GenericResource import GenericResource
-from mage_ai.services.spark.api.local import LocalAPI
+from mage_ai.api.resources.mixins.spark import SparkApplicationChild
 from mage_ai.services.spark.models.applications import Application
 from mage_ai.shared.hash import index_by
 
 
-class SparkApplicationResource(GenericResource):
+class SparkApplicationResource(GenericResource, SparkApplicationChild):
     @classmethod
     async def get_model(self, pk):
         return Application.load(id=pk)
 
     @classmethod
     async def collection(self, _query, _meta, user, **kwargs):
-        applications = await LocalAPI().applications()
+        applications = await self.build_api().applications()
         mapping = index_by(lambda x: x.id, applications)
 
         applications_cache = Application.get_applications_from_cache()
         if applications_cache:
             for application in applications_cache.values():
-                if application.id in mapping:
+                if application.calculated_id() in mapping:
                     continue
                 applications.append(application)
-                mapping[application.id] = application
+                mapping[application.calculated_id()] = application
 
         return self.build_result_set(
             applications,

mage_ai/api/resources/SparkEnvironmentResource.py

@@ -1,13 +1,12 @@
 from mage_ai.api.resources.GenericResource import GenericResource
 from mage_ai.api.resources.mixins.spark import SparkApplicationChild
-from mage_ai.services.spark.api.local import LocalAPI
 
 
 class SparkEnvironmentResource(GenericResource, SparkApplicationChild):
     @classmethod
     async def member(self, _pk, user, **kwargs):
         return self(
-            await LocalAPI().environment(),
+            await self.build_api().environment(),
             user,
             **kwargs,
         )

mage_ai/api/resources/SparkExecutorResource.py

@@ -1,6 +1,5 @@
 from mage_ai.api.resources.GenericResource import GenericResource
 from mage_ai.api.resources.mixins.spark import SparkApplicationChild
-from mage_ai.services.spark.api.local import LocalAPI
 from mage_ai.services.spark.models.executors import Executor
 
 
@@ -10,7 +9,7 @@ class SparkExecutorResource(GenericResource, SparkApplicationChild):
         application_id = await self.get_application_id(**kwargs)
 
         return self.build_result_set(
-            await LocalAPI().executors(application_id=application_id),
+            await self.build_api().executors(application_id=application_id),
             user,
             **kwargs,
         )

mage_ai/api/resources/SparkJobResource.py

@@ -1,13 +1,12 @@
 from mage_ai.api.resources.GenericResource import GenericResource
 from mage_ai.api.resources.mixins.spark import SparkApplicationChild
-from mage_ai.services.spark.api.local import LocalAPI
 
 
 class SparkJobResource(GenericResource, SparkApplicationChild):
     @classmethod
     async def collection(self, _query, _meta, user, **kwargs):
         return self.build_result_set(
-            await LocalAPI().jobs(),
+            await self.build_api().jobs(),
             user,
             **kwargs,
         )
@@ -16,16 +15,14 @@ class SparkJobResource(GenericResource, SparkApplicationChild):
     async def member(self, pk, user, **kwargs):
         query_arg = kwargs.get('query')
 
-        application_id = query_arg.get('application_id', [])
-        if application_id:
-            application_id = application_id[0]
+        application_id = self.application_calculated_id_from_query(query_arg)
 
         application_spark_ui_url = query_arg.get('application_spark_ui_url', [])
         if application_spark_ui_url:
             application_spark_ui_url = application_spark_ui_url[0]
 
         return self(
-            await LocalAPI().job(
+            await self.build_api().job(
                 job_id=pk,
                 application_id=application_id,
                 application_spark_ui_url=application_spark_ui_url,

mage_ai/api/resources/SparkSqlResource.py

@@ -2,7 +2,6 @@ from functools import reduce
 
 from mage_ai.api.resources.GenericResource import GenericResource
 from mage_ai.api.resources.mixins.spark import SparkApplicationChild
-from mage_ai.services.spark.api.local import LocalAPI
 
 
 class SparkSqlResource(GenericResource, SparkApplicationChild):
@@ -16,16 +15,14 @@ class SparkSqlResource(GenericResource, SparkApplicationChild):
             if length:
                 query['length'] = length
 
-        application_id = query_arg.get('application_id', [])
-        if application_id:
-            application_id = application_id[0]
+        application_id = self.application_calculated_id_from_query(query_arg)
 
         application_spark_ui_url = query_arg.get('application_spark_ui_url', [])
         if application_spark_ui_url:
             application_spark_ui_url = application_spark_ui_url[0]
 
         return self.build_result_set(
-            await LocalAPI().sqls(
+            await self.build_api().sqls(
                 application_id=application_id,
                 application_spark_ui_url=application_spark_ui_url,
                 query=query,
@@ -42,15 +39,13 @@ class SparkSqlResource(GenericResource, SparkApplicationChild):
         if include_jobs_and_stages:
             include_jobs_and_stages = include_jobs_and_stages[0]
 
-        application_id = query_arg.get('application_id', [])
-        if application_id:
-            application_id = application_id[0]
+        application_id = self.application_calculated_id_from_query(query_arg)
 
         application_spark_ui_url = query_arg.get('application_spark_ui_url', [])
         if application_spark_ui_url:
             application_spark_ui_url = application_spark_ui_url[0]
 
-        model = await LocalAPI().sql(
+        model = await self.build_api().sql(
             application_id=application_id,
             application_spark_ui_url=application_spark_ui_url,
             sql_id=pk,
@@ -59,14 +54,14 @@ class SparkSqlResource(GenericResource, SparkApplicationChild):
         if include_jobs_and_stages:
             job_ids = (model.failed_job_ids or []) + (
                 model.running_job_ids or []) + (model.success_job_ids or [])
-            model.jobs = [await LocalAPI().job(
+            model.jobs = [await self.build_api().job(
                 application_id=application_id,
                 application_spark_ui_url=application_spark_ui_url,
                 job_id=job_id,
             ) for job_id in job_ids]
 
             stage_ids = reduce(lambda acc, job: acc + (job.stage_ids or []), model.jobs, [])
-            model.stages = [await LocalAPI().stage(
+            model.stages = [await self.build_api().stage(
                 application_id=application_id,
                 application_spark_ui_url=application_spark_ui_url,
                 stage_id=stage_id,

mage_ai/api/resources/SparkStageAttemptResource.py

@@ -1,6 +1,5 @@
 from mage_ai.api.resources.GenericResource import GenericResource
 from mage_ai.api.resources.mixins.spark import SparkApplicationChild
-from mage_ai.services.spark.api.local import LocalAPI
 
 
 class SparkStageAttemptResource(GenericResource, SparkApplicationChild):
@@ -9,7 +8,7 @@ class SparkStageAttemptResource(GenericResource, SparkApplicationChild):
         parent_model = kwargs.get('parent_model')
 
         return self.build_result_set(
-            await LocalAPI().stage_attempts(
+            await self.build_api().stage_attempts(
                 stage_id=parent_model.id,
             ),
             user,
@@ -21,7 +20,7 @@ class SparkStageAttemptResource(GenericResource, SparkApplicationChild):
         parent_model = kwargs.get('parent_model')
 
         return self(
-            await LocalAPI().stage_attempt(
+            await self.build_api().stage_attempt(
                 attempt_id=pk,
                 stage_id=parent_model.id,
             ),

mage_ai/api/resources/SparkStageAttemptTaskResource.py

@@ -1,7 +1,6 @@
 from mage_ai.api.errors import ApiError
 from mage_ai.api.resources.GenericResource import GenericResource
 from mage_ai.api.resources.mixins.spark import SparkApplicationChild
-from mage_ai.services.spark.api.local import LocalAPI
 
 
 class SparkStageAttemptTaskResource(GenericResource, SparkApplicationChild):
@@ -18,7 +17,7 @@ class SparkStageAttemptTaskResource(GenericResource, SparkApplicationChild):
             raise error
 
         return self.build_result_set(
-            await LocalAPI().stage_attempt_tasks(
+            await self.build_api().stage_attempt_tasks(
                 attempt_id=attempt_id,
                 stage_id=parent_model.id,
             ),

mage_ai/api/resources/SparkStageAttemptTaskSummaryResource.py

@@ -1,7 +1,6 @@
 from mage_ai.api.errors import ApiError
 from mage_ai.api.resources.GenericResource import GenericResource
 from mage_ai.api.resources.mixins.spark import SparkApplicationChild
-from mage_ai.services.spark.api.local import LocalAPI
 
 
 class SparkStageAttemptTaskSummaryResource(GenericResource, SparkApplicationChild):
@@ -14,7 +13,7 @@ class SparkStageAttemptTaskSummaryResource(GenericResource, SparkApplicationChil
             raise error
 
         return self(
-            await LocalAPI().stage_attempt_task_summary(
+            await self.build_api().stage_attempt_task_summary(
                 attempt_id=pk,
                 stage_id=parent_model.id,
             ),

mage_ai/api/resources/SparkStageResource.py

@@ -1,6 +1,5 @@
 from mage_ai.api.resources.GenericResource import GenericResource
 from mage_ai.api.resources.mixins.spark import SparkApplicationChild
-from mage_ai.services.spark.api.local import LocalAPI
 from mage_ai.services.spark.models.stages import Stage
 
 
@@ -16,7 +15,7 @@ class SparkStageResource(GenericResource, SparkApplicationChild):
                 query['details'] = details
 
         return self.build_result_set(
-            await LocalAPI().stages(query=query),
+            await self.build_api().stages(query=query),
             user,
             **kwargs,
         )
@@ -42,15 +41,13 @@ class SparkStageResource(GenericResource, SparkApplicationChild):
             if quantiles:
                 query['withSummaries'] = with_summaries
 
-        application_id = query_arg.get('application_id', [])
-        if application_id:
-            application_id = application_id[0]
+        application_id = self.application_calculated_id_from_query(query_arg)
 
         application_spark_ui_url = query_arg.get('application_spark_ui_url', [])
         if application_spark_ui_url:
             application_spark_ui_url = application_spark_ui_url[0]
 
-        stage = await LocalAPI().stage(
+        stage = await self.build_api().stage(
             application_id=application_id,
             application_spark_ui_url=application_spark_ui_url,
             stage_id=pk,

mage_ai/api/resources/SparkThreadResource.py

@@ -1,6 +1,5 @@
 from mage_ai.api.resources.GenericResource import GenericResource
 from mage_ai.api.resources.mixins.spark import SparkApplicationChild
-from mage_ai.services.spark.api.local import LocalAPI
 
 
 class SparkThreadResource(GenericResource, SparkApplicationChild):
@@ -9,7 +8,7 @@ class SparkThreadResource(GenericResource, SparkApplicationChild):
         parent_model = kwargs.get('parent_model')
 
         return self.build_result_set(
-            await LocalAPI().threads(
+            await self.build_api().threads(
                 executor_id=parent_model.id,
             ),
             user,