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

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,

mage_ai/api/resources/mixins/spark.py

@@ -1,19 +1,26 @@
+from typing import Dict
+
 from mage_ai.api.errors import ApiError
-from mage_ai.services.spark.api.local import LocalAPI
+from mage_ai.services.spark.api.base import BaseAPI
+from mage_ai.services.spark.api.service import API
 from mage_ai.services.spark.models.applications import Application
 
 
 class SparkApplicationChild:
+    @classmethod
+    def build_api(self) -> BaseAPI:
+        return API.build()
+
     @classmethod
     async def get_application_id(self, **kwargs) -> str:
         application_id = None
         parent_model = kwargs.get('parent_model')
         if parent_model and isinstance(parent_model, Application):
-            application_id = parent_model.id
+            application_id = parent_model.calculated_id()
         else:
-            models = await LocalAPI().applications()
+            models = await self.build_api().applications()
             if models and len(models) >= 1:
-                application_id = models[0].id
+                application_id = models[0].calculated_id()
 
         if not application_id:
             error = ApiError(ApiError.RESOURCE_NOT_FOUND)
@@ -22,3 +29,17 @@ class SparkApplicationChild:
             raise error
 
         return application_id
+
+    @classmethod
+    def application_calculated_id_from_query(self, query: Dict = None) -> str:
+        if not query:
+            return None
+
+        application_id = query.get('application_id', [])
+        if application_id:
+            application_id = application_id[0]
+
+        if application_id:
+            application_id = Application.load(id=application_id).calculated_id()
+
+        return application_id

mage_ai/authentication/oauth2.py

@@ -1,4 +1,3 @@
-import os
 import secrets
 from datetime import datetime, timedelta
 from typing import Dict
@@ -10,10 +9,9 @@ from mage_ai.orchestration.db.models.oauth import (
     Oauth2Application,
     User,
 )
-from mage_ai.settings import MAGE_ACCESS_TOKEN_EXPIRY_TIME
+from mage_ai.settings import JWT_SECRET, MAGE_ACCESS_TOKEN_EXPIRY_TIME
 
 JWT_ALGORITHM = 'HS256'
-JWT_SECRET = os.getenv('JWT_SECRET', 'materia')
 
 
 def generate_access_token(

mage_ai/authentication/permissions/constants.py

@@ -17,6 +17,9 @@ class EntityName(str, Enum):
     Chart = 'Chart'
     ClientPage = 'ClientPage'
     Cluster = 'Cluster'
+    ComputeCluster = 'ComputeCluster'
+    ComputeConnection = 'ComputeConnection'
+    ComputeService = 'ComputeService'
     CustomTemplate = 'CustomTemplate'
     DataProvider = 'DataProvider'
     Database = 'Database'
@@ -33,6 +36,7 @@ class EntityName(str, Enum):
     GitCustomBranch = 'GitCustomBranch'
     GitFile = 'GitFile'
     GlobalDataProduct = 'GlobalDataProduct'
+    GlobalHook = 'GlobalHook'
     IntegrationDestination = 'IntegrationDestination'
     IntegrationSource = 'IntegrationSource'
     IntegrationSourceStream = 'IntegrationSourceStream'

mage_ai/authentication/providers/active_directory.py

@@ -0,0 +1,136 @@
+import json
+import uuid
+from typing import Awaitable, Dict
+from urllib.parse import quote_plus
+
+import aiohttp
+
+from mage_ai.authentication.oauth.constants import (
+    ACTIVE_DIRECTORY_CLIENT_ID as ACTIVE_DIRECTORY_MAGE_CLIENT_ID,
+)
+from mage_ai.authentication.oauth.constants import OAUTH_PROVIDER_ACTIVE_DIRECTORY
+from mage_ai.authentication.providers.oauth import OauthProvider
+from mage_ai.authentication.providers.sso import SsoProvider
+from mage_ai.authentication.providers.utils import get_base_url
+from mage_ai.settings import ACTIVE_DIRECTORY_DIRECTORY_ID
+from mage_ai.settings.sso import (
+    ACTIVE_DIRECTORY_CLIENT_ID,
+    ACTIVE_DIRECTORY_CLIENT_SECRET,
+)
+
+
+class ADProvider(SsoProvider, OauthProvider):
+    provider = OAUTH_PROVIDER_ACTIVE_DIRECTORY
+
+    def __init__(self):
+        self.__validate()
+
+    def __validate(self):
+        if not ACTIVE_DIRECTORY_DIRECTORY_ID:
+            raise Exception(
+                'AD directory id is empty. '
+                'Make sure the ACTIVE_DIRECTORY_DIRECTORY_ID environment variable is set.')
+        if ACTIVE_DIRECTORY_CLIENT_ID and not ACTIVE_DIRECTORY_CLIENT_SECRET:
+            raise Exception(
+                'AD client secret is empty. '
+                'Make sure the ACTIVE_DIRECTORY_CLIENT_SECRET environment variable is set.')
+
+    def get_auth_url_response(self, redirect_uri: str = None, **kwargs) -> Dict:
+        """
+        For active directory, the user can set up the Oauth provider in two different ways. They
+        can either just set the ACTIVE_DIRECTORY_DIRECTORY_ID which will use the Mage application we
+        set up in Azure. They can additionally set the ACTIVE_DIRECTORY_CLIENT_ID and
+        ACTIVE_DIRECTORY_CLIENT_SECRET which will use their own Mage application in Azure.
+        """
+        ad_directory_id = ACTIVE_DIRECTORY_DIRECTORY_ID
+        if ACTIVE_DIRECTORY_CLIENT_ID:
+            base_url = get_base_url(redirect_uri)
+            redirect_uri_query = dict(
+                provider=self.provider,
+                redirect_uri=redirect_uri,
+            )
+            query = dict(
+                client_id=ACTIVE_DIRECTORY_CLIENT_ID,
+                redirect_uri=quote_plus(
+                    f'{base_url}/oauth',
+                ),
+                response_type='code',
+                scope='User.Read',
+                state=uuid.uuid4().hex,
+            )
+            query_strings = []
+            for k, v in query.items():
+                query_strings.append(f'{k}={v}')
+
+            return dict(
+                url=f"https://login.microsoftonline.com/{ad_directory_id}/oauth2/v2.0/authorize?{'&'.join(query_strings)}",  # noqa: E501
+                redirect_query_params=redirect_uri_query,
+            )
+        else:
+            from requests.models import PreparedRequest
+
+            req = PreparedRequest()
+            req.prepare_url(redirect_uri, dict(provider=self.provider))
+            query = dict(
+                client_id=ACTIVE_DIRECTORY_MAGE_CLIENT_ID,
+                redirect_uri=f'https://api.mage.ai/v1/oauth/{self.provider}',
+                response_type='code',
+                scope='User.Read',
+                state=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}')
+
+            return dict(
+                url=f"https://login.microsoftonline.com/{ad_directory_id}/oauth2/v2.0/authorize?{'&'.join(query_strings)}",  # noqa: E501
+            )
+
+    async def get_access_token_response(self, code: str, **kwargs) -> Awaitable[Dict]:
+        base_url = get_base_url(kwargs.get('redirect_uri'))
+
+        async with aiohttp.ClientSession() as session:
+            async with session.post(
+                'https://login.microsoftonline.com/organizations/oauth2/v2.0/token',
+                headers={
+                    'Accept': 'application/json',
+                },
+                data=dict(
+                    client_id=ACTIVE_DIRECTORY_CLIENT_ID,
+                    client_secret=ACTIVE_DIRECTORY_CLIENT_SECRET,
+                    code=code,
+                    grant_type='authorization_code',
+                    redirect_uri=f'{base_url}/oauth',
+                    tenant=ACTIVE_DIRECTORY_DIRECTORY_ID,
+                ),
+                timeout=20,
+            ) as response:
+                data = await response.json()
+
+        return data
+
+    async def get_user_info(self, access_token: str = None, **kwargs) -> Awaitable[Dict]:
+        if access_token is None:
+            raise Exception('Access token is required to fetch user info.')
+        async with aiohttp.ClientSession() as session:
+            async with session.get(
+                'https://graph.microsoft.com/v1.0/me',
+                headers={
+                    'Content-Type': 'application\\json',
+                    'Authorization': f'Bearer {access_token}',
+                },
+                timeout=10,
+            ) as response:
+                userinfo_resp = await response.json()
+
+        return dict(
+            email=userinfo_resp.get('userPrincipalName'),
+            username=userinfo_resp.get('userPrincipalName'),
+        )

mage_ai/authentication/providers/constants.py

@@ -1,13 +1,16 @@
 from mage_ai.authentication.oauth.constants import (
+    OAUTH_PROVIDER_ACTIVE_DIRECTORY,
     OAUTH_PROVIDER_GHE,
     OAUTH_PROVIDER_GOOGLE,
     OAUTH_PROVIDER_OKTA,
 )
+from mage_ai.authentication.providers.active_directory import ADProvider
 from mage_ai.authentication.providers.ghe import GHEProvider
 from mage_ai.authentication.providers.google import GoogleProvider
 from mage_ai.authentication.providers.okta import OktaProvider
 
 NAME_TO_PROVIDER = {
+    OAUTH_PROVIDER_ACTIVE_DIRECTORY: ADProvider,
     OAUTH_PROVIDER_GHE: GHEProvider,
     OAUTH_PROVIDER_GOOGLE: GoogleProvider,
     OAUTH_PROVIDER_OKTA: OktaProvider,

mage_ai/authentication/providers/ghe.py

@@ -1,7 +1,7 @@
 import os
 import urllib.parse
 import uuid
-from typing import Dict
+from typing import Awaitable, Dict
 
 import aiohttp
 
@@ -39,30 +39,29 @@ class GHEProvider(OauthProvider):
                 'Make sure the GHE_CLIENT_SECRET environment variable is set.')
 
     def get_auth_url_response(self, redirect_uri: str = None, **kwargs) -> Dict:
-        if self.hostname:
-            base_url = get_base_url(redirect_uri)
-            redirect_uri_query = dict(
-                provider=self.provider,
-                redirect_uri=redirect_uri,
-            )
-            query = dict(
-                client_id=os.getenv(GHE_CLIENT_ID_ENV_VAR),
-                redirect_uri=urllib.parse.quote_plus(
-                    f'{base_url}/oauth',
-                ),
-                scope='repo',
-                state=uuid.uuid4().hex,
-            )
-            query_strings = []
-            for k, v in query.items():
-                query_strings.append(f'{k}={v}')
+        base_url = get_base_url(redirect_uri)
+        redirect_uri_query = dict(
+            provider=self.provider,
+            redirect_uri=redirect_uri,
+        )
+        query = dict(
+            client_id=self.client_id,
+            redirect_uri=urllib.parse.quote_plus(
+                f'{base_url}/oauth',
+            ),
+            scope='repo',
+            state=uuid.uuid4().hex,
+        )
+        query_strings = []
+        for k, v in query.items():
+            query_strings.append(f'{k}={v}')
 
-            return dict(
-                url=f"{self.hostname}/login/oauth/authorize?{'&'.join(query_strings)}",
-                redirect_query_params=redirect_uri_query,
-            )
+        return dict(
+            url=f"{self.hostname}/login/oauth/authorize?{'&'.join(query_strings)}",
+            redirect_query_params=redirect_uri_query,
+        )
 
-    async def get_access_token_response(self, code: str, **kwargs) -> Dict:
+    async def get_access_token_response(self, code: str, **kwargs) -> Awaitable[Dict]:
         data = dict()
         async with aiohttp.ClientSession() as session:
             async with session.post(
@@ -71,8 +70,8 @@ class GHEProvider(OauthProvider):
                     'Accept': 'application/json',
                 },
                 data=dict(
-                    client_id=os.getenv(GHE_CLIENT_ID_ENV_VAR),
-                    client_secret=os.getenv(GHE_CLIENT_SECRET_ENV_VAR),
+                    client_id=self.client_id,
+                    client_secret=self.client_secret,
                     code=code,
                 ),
                 timeout=20,

mage_ai/authentication/providers/google.py

@@ -1,6 +1,6 @@
 import urllib.parse
 import uuid
-from typing import Dict
+from typing import Awaitable, Dict
 
 import aiohttp
 
@@ -25,32 +25,31 @@ class GoogleProvider(SsoProvider, OauthProvider):
                 'Make sure the GOOGLE_CLIENT_SECRET environment variable is set.')
 
     def get_auth_url_response(self, redirect_uri: str = None, **kwargs) -> Dict:
-        if GOOGLE_CLIENT_ID:
-            base_url = get_base_url(redirect_uri)
-            redirect_uri_query = dict(
-                provider=self.provider,
-                redirect_uri=redirect_uri,
-            )
-            query = dict(
-                client_id=GOOGLE_CLIENT_ID,
-                prompt='select_account',
-                redirect_uri=urllib.parse.quote_plus(
-                    f'{base_url}/oauth',
-                ),
-                response_type='code',
-                scope='https://www.googleapis.com/auth/userinfo.email https://www.googleapis.com/auth/userinfo.profile',  # noqa: E501
-                state=uuid.uuid4().hex,
-            )
-            query_strings = []
-            for k, v in query.items():
-                query_strings.append(f'{k}={v}')
+        base_url = get_base_url(redirect_uri)
+        redirect_uri_query = dict(
+            provider=self.provider,
+            redirect_uri=redirect_uri,
+        )
+        query = dict(
+            client_id=GOOGLE_CLIENT_ID,
+            prompt='select_account',
+            redirect_uri=urllib.parse.quote_plus(
+                f'{base_url}/oauth',
+            ),
+            response_type='code',
+            scope='https://www.googleapis.com/auth/userinfo.email https://www.googleapis.com/auth/userinfo.profile',  # noqa: E501
+            state=uuid.uuid4().hex,
+        )
+        query_strings = []
+        for k, v in query.items():
+            query_strings.append(f'{k}={v}')
 
-            return dict(
-                url=f"https://accounts.google.com/o/oauth2/v2/auth?{'&'.join(query_strings)}",
-                redirect_query_params=redirect_uri_query,
-            )
+        return dict(
+            url=f"https://accounts.google.com/o/oauth2/v2/auth?{'&'.join(query_strings)}",
+            redirect_query_params=redirect_uri_query,
+        )
 
-    async def get_access_token_response(self, code: str, **kwargs) -> Dict:
+    async def get_access_token_response(self, code: str, **kwargs) -> Awaitable[Dict]:
         base_url = get_base_url(kwargs.get('redirect_uri'))
 
         data = dict()
@@ -70,7 +69,7 @@ class GoogleProvider(SsoProvider, OauthProvider):
 
         return data
 
-    async def get_user_info(self, access_token: str = None, **kwargs) -> Dict:
+    async def get_user_info(self, access_token: str = None, **kwargs) -> Awaitable[Dict]:
         if access_token is None:
             raise Exception('Access token is required to fetch user info.')
         async with aiohttp.ClientSession() as session: