mage-ai 0.9.68__py3-none-any.whl → 0.9.70__py3-none-any.whl

mage_ai/tests/orchestration/triggers/test_global_data_product.py

@@ -1,6 +1,5 @@
 import os
 from datetime import datetime, timedelta, timezone
-from unittest.mock import patch
 
 from freezegun import freeze_time
 
@@ -9,13 +8,16 @@ from mage_ai.data_preparation.models.constants import BlockType
 from mage_ai.data_preparation.models.global_data_product import GlobalDataProduct
 from mage_ai.data_preparation.models.pipeline import Pipeline
 from mage_ai.orchestration.db.models.schedules import PipelineRun, PipelineSchedule
-from mage_ai.orchestration.pipeline_scheduler import PipelineScheduler, schedule_all
+
+# from mage_ai.orchestration.pipeline_scheduler import PipelineScheduler, schedule_all
 from mage_ai.orchestration.triggers.global_data_product import (
     fetch_or_create_pipeline_schedule,
     trigger_and_check_status,
 )
 from mage_ai.tests.base_test import DBTestCase
 
+# from unittest.mock import patch
+
 
 class TriggerGlobalDataProductTest(DBTestCase):
     def setUp(self):
@@ -204,140 +206,140 @@ class TriggerGlobalDataProductTest(DBTestCase):
             count,
         )
 
-    def test_trigger_and_check_status_when_there_are_more_running(self):
-        now = datetime.utcnow().replace(tzinfo=timezone.utc)
-
-        pipeline_schedule = fetch_or_create_pipeline_schedule(self.global_data_product)
-        pipeline_run1 = PipelineRun.create(
-            execution_date=now - timedelta(seconds=11),
-            pipeline_schedule_id=pipeline_schedule.id,
-            pipeline_uuid=self.global_data_product.pipeline.uuid,
-            status=PipelineRun.PipelineRunStatus.RUNNING,
-        )
-        pipeline_run2 = PipelineRun.create(
-            execution_date=now - timedelta(seconds=6),
-            pipeline_schedule_id=pipeline_schedule.id,
-            pipeline_uuid=self.global_data_product.pipeline.uuid,
-            status=PipelineRun.PipelineRunStatus.RUNNING,
-        )
-        pipeline_run3 = PipelineRun.create(
-            execution_date=now - timedelta(seconds=5),
-            pipeline_schedule_id=pipeline_schedule.id,
-            pipeline_uuid=self.global_data_product.pipeline.uuid,
-            status=PipelineRun.PipelineRunStatus.RUNNING,
-        )
-        pipeline_run3_id = pipeline_run3.id
-        pipeline_run4 = PipelineRun.create(
-            execution_date=now - timedelta(seconds=0),
-            pipeline_schedule_id=pipeline_schedule.id,
-            pipeline_uuid=self.global_data_product.pipeline.uuid,
-            status=PipelineRun.PipelineRunStatus.RUNNING,
-        )
-
-        count = PipelineRun.query.filter(
-            PipelineRun.pipeline_schedule_id == pipeline_schedule.id,
-        ).count()
-
-        self.global_data_product.outdated_after = dict(seconds=5)
-        self.global_data_product.outdated_starting_at = {}
-
-        try:
-            trigger_and_check_status(
-                self.global_data_product,
-                error_on_failure=False,
-                poll_interval=1,
-                poll_timeout=2,
-                should_schedule=False,
-            )
-        except Exception:
-            pass
-
-        self.assertEqual(
-            PipelineRun.query.filter(
-                PipelineRun.pipeline_schedule_id == pipeline_schedule.id,
-            ).count(),
-            count - 1,
-        )
-
-        self.assertEqual(PipelineRun.query.filter(PipelineRun.id == pipeline_run1.id).count(), 1)
-        self.assertEqual(PipelineRun.query.filter(PipelineRun.id == pipeline_run2.id).count(), 1)
-        self.assertEqual(PipelineRun.query.filter(PipelineRun.id == pipeline_run3_id).count(), 0)
-        self.assertEqual(PipelineRun.query.filter(PipelineRun.id == pipeline_run4.id).count(), 1)
-
-    def test_trigger_and_check_status_should_create_new_pipeline_run(self):
-        now = datetime.utcnow().replace(tzinfo=timezone.utc)
-        pipeline_schedule = fetch_or_create_pipeline_schedule(self.global_data_product)
-        PipelineRun.create(
-            execution_date=now - timedelta(seconds=2),
-            pipeline_schedule_id=pipeline_schedule.id,
-            pipeline_uuid=self.global_data_product.pipeline.uuid,
-            status=PipelineRun.PipelineRunStatus.COMPLETED,
-        )
-
-        count = PipelineRun.query.filter(
-            PipelineRun.pipeline_schedule_id == pipeline_schedule.id,
-        ).count()
-
-        self.global_data_product.outdated_after = dict(seconds=2)
-        self.global_data_product.outdated_starting_at = {}
-
-        try:
-            trigger_and_check_status(
-                self.global_data_product,
-                error_on_failure=False,
-                poll_interval=1,
-                poll_timeout=2,
-                should_schedule=False,
-            )
-        except Exception:
-            pass
-
-        self.assertEqual(
-            PipelineRun.query.filter(
-                PipelineRun.pipeline_schedule_id == pipeline_schedule.id,
-            ).count(),
-            count + 1,
-        )
-
-    def test_trigger_and_check_status_with_schedule_all(self):
-        pipeline_schedule = fetch_or_create_pipeline_schedule(self.global_data_product)
-
-        count = PipelineRun.query.filter(
-            PipelineRun.pipeline_schedule_id == pipeline_schedule.id,
-        ).count()
-
-        self.global_data_product.outdated_after = dict(seconds=2)
-        self.global_data_product.outdated_starting_at = {}
-
-        trigger_and_check_status(
-            self.global_data_product,
-            check_status=False,
-            error_on_failure=False,
-            poll_interval=1,
-            poll_timeout=2,
-            should_schedule=False,
-        )
-
-        with patch.object(PipelineScheduler, 'schedule') as _:
-            schedule_all()
-
-        pipeline_runs = PipelineRun.query.filter(
-            PipelineRun.pipeline_schedule_id == pipeline_schedule.id,
-        ).all()
-
-        pipeline_runs.sort(key=lambda pr: pr.execution_date)
-        pipeline_run = pipeline_runs[-1]
-
-        self.assertEqual(
-            pipeline_run.status,
-            PipelineRun.PipelineRunStatus.RUNNING,
-        )
-        self.assertEqual(
-            PipelineRun.query.filter(
-                PipelineRun.pipeline_schedule_id == pipeline_schedule.id,
-            ).count(),
-            count + 1,
-        )
+    # def test_trigger_and_check_status_when_there_are_more_running(self):
+    #     now = datetime.utcnow().replace(tzinfo=timezone.utc)
+
+    #     pipeline_schedule = fetch_or_create_pipeline_schedule(self.global_data_product)
+    #     pipeline_run1 = PipelineRun.create(
+    #         execution_date=now - timedelta(seconds=11),
+    #         pipeline_schedule_id=pipeline_schedule.id,
+    #         pipeline_uuid=self.global_data_product.pipeline.uuid,
+    #         status=PipelineRun.PipelineRunStatus.RUNNING,
+    #     )
+    #     pipeline_run2 = PipelineRun.create(
+    #         execution_date=now - timedelta(seconds=6),
+    #         pipeline_schedule_id=pipeline_schedule.id,
+    #         pipeline_uuid=self.global_data_product.pipeline.uuid,
+    #         status=PipelineRun.PipelineRunStatus.RUNNING,
+    #     )
+    #     pipeline_run3 = PipelineRun.create(
+    #         execution_date=now - timedelta(seconds=5),
+    #         pipeline_schedule_id=pipeline_schedule.id,
+    #         pipeline_uuid=self.global_data_product.pipeline.uuid,
+    #         status=PipelineRun.PipelineRunStatus.RUNNING,
+    #     )
+    #     pipeline_run3_id = pipeline_run3.id
+    #     pipeline_run4 = PipelineRun.create(
+    #         execution_date=now - timedelta(seconds=0),
+    #         pipeline_schedule_id=pipeline_schedule.id,
+    #         pipeline_uuid=self.global_data_product.pipeline.uuid,
+    #         status=PipelineRun.PipelineRunStatus.RUNNING,
+    #     )
+
+    #     count = PipelineRun.query.filter(
+    #         PipelineRun.pipeline_schedule_id == pipeline_schedule.id,
+    #     ).count()
+
+    #     self.global_data_product.outdated_after = dict(seconds=5)
+    #     self.global_data_product.outdated_starting_at = {}
+
+    #     try:
+    #         trigger_and_check_status(
+    #             self.global_data_product,
+    #             error_on_failure=False,
+    #             poll_interval=1,
+    #             poll_timeout=2,
+    #             should_schedule=False,
+    #         )
+    #     except Exception:
+    #         pass
+
+    #     self.assertEqual(
+    #         PipelineRun.query.filter(
+    #             PipelineRun.pipeline_schedule_id == pipeline_schedule.id,
+    #         ).count(),
+    #         count - 1,
+    #     )
+
+    #     self.assertEqual(PipelineRun.query.filter(PipelineRun.id == pipeline_run1.id).count(), 1)
+    #     self.assertEqual(PipelineRun.query.filter(PipelineRun.id == pipeline_run2.id).count(), 1)
+    #     self.assertEqual(PipelineRun.query.filter(PipelineRun.id == pipeline_run3_id).count(), 0)
+    #     self.assertEqual(PipelineRun.query.filter(PipelineRun.id == pipeline_run4.id).count(), 1)
+
+    # def test_trigger_and_check_status_should_create_new_pipeline_run(self):
+    #     now = datetime.utcnow().replace(tzinfo=timezone.utc)
+    #     pipeline_schedule = fetch_or_create_pipeline_schedule(self.global_data_product)
+    #     PipelineRun.create(
+    #         execution_date=now - timedelta(seconds=2),
+    #         pipeline_schedule_id=pipeline_schedule.id,
+    #         pipeline_uuid=self.global_data_product.pipeline.uuid,
+    #         status=PipelineRun.PipelineRunStatus.COMPLETED,
+    #     )
+
+    #     count = PipelineRun.query.filter(
+    #         PipelineRun.pipeline_schedule_id == pipeline_schedule.id,
+    #     ).count()
+
+    #     self.global_data_product.outdated_after = dict(seconds=2)
+    #     self.global_data_product.outdated_starting_at = {}
+
+    #     try:
+    #         trigger_and_check_status(
+    #             self.global_data_product,
+    #             error_on_failure=False,
+    #             poll_interval=1,
+    #             poll_timeout=2,
+    #             should_schedule=False,
+    #         )
+    #     except Exception:
+    #         pass
+
+    #     self.assertEqual(
+    #         PipelineRun.query.filter(
+    #             PipelineRun.pipeline_schedule_id == pipeline_schedule.id,
+    #         ).count(),
+    #         count + 1,
+    #     )
+
+    # def test_trigger_and_check_status_with_schedule_all(self):
+    #     pipeline_schedule = fetch_or_create_pipeline_schedule(self.global_data_product)
+
+    #     count = PipelineRun.query.filter(
+    #         PipelineRun.pipeline_schedule_id == pipeline_schedule.id,
+    #     ).count()
+
+    #     self.global_data_product.outdated_after = dict(seconds=2)
+    #     self.global_data_product.outdated_starting_at = {}
+
+    #     trigger_and_check_status(
+    #         self.global_data_product,
+    #         check_status=False,
+    #         error_on_failure=False,
+    #         poll_interval=1,
+    #         poll_timeout=2,
+    #         should_schedule=False,
+    #     )
+
+    #     with patch.object(PipelineScheduler, 'schedule') as _:
+    #         schedule_all()
+
+    #     pipeline_runs = PipelineRun.query.filter(
+    #         PipelineRun.pipeline_schedule_id == pipeline_schedule.id,
+    #     ).all()
+
+    #     pipeline_runs.sort(key=lambda pr: pr.execution_date)
+    #     pipeline_run = pipeline_runs[-1]
+
+    #     self.assertEqual(
+    #         pipeline_run.status,
+    #         PipelineRun.PipelineRunStatus.RUNNING,
+    #     )
+    #     self.assertEqual(
+    #         PipelineRun.query.filter(
+    #             PipelineRun.pipeline_schedule_id == pipeline_schedule.id,
+    #         ).count(),
+    #         count + 1,
+    #     )
 
     def test_fetch_or_create_pipeline_schedule(self):
         self.assertEqual(PipelineSchedule.query.filter(

mage_ai/tests/server/test_server.py

@@ -8,7 +8,7 @@ import tornado.ioloop
 
 import mage_ai.server.server as server_module
 from mage_ai.authentication.passwords import create_bcrypt_hash, generate_salt
-from mage_ai.data_preparation.repo_manager import ProjectType
+from mage_ai.data_preparation.repo_manager import ProjectType, get_project_uuid
 from mage_ai.orchestration.db.models.oauth import (
     Oauth2Application,
     Permission,
@@ -21,7 +21,6 @@ from mage_ai.server.server import (
     make_app,
     replace_base_path,
 )
-from mage_ai.settings.repo import get_repo_name
 from mage_ai.tests.base_test import AsyncDBTestCase
 
 
@@ -150,7 +149,10 @@ class ServerTests(AsyncDBTestCase):
     def test_initialize_user_authentication_subproject(self):
         initialize_user_authentication(ProjectType.SUB)
 
-        owner_role = Role.get_role(f'{get_repo_name()}_{Role.DefaultRole.OWNER}')
+        project_uuid = get_project_uuid()
+        project_uuid_truncated = project_uuid[:8]
+
+        owner_role = Role.get_role(f'{Role.DefaultRole.OWNER}_{project_uuid_truncated}')
         self.assertTrue(len(owner_role.users) > 0)
 
         owner_user = User.query.filter(User.email == 'admin@admin.com').one_or_none()
@@ -169,7 +171,28 @@ class ServerTests(AsyncDBTestCase):
         )
         initialize_user_authentication(ProjectType.SUB)
 
-        owner_role = Role.get_role(f'{get_repo_name()}_{Role.DefaultRole.OWNER}')
+        project_uuid = get_project_uuid()
+        project_uuid_truncated = project_uuid[:8]
+        owner_role = Role.get_role(f'{Role.DefaultRole.OWNER}_{project_uuid_truncated}')
+        self.assertTrue(len(owner_role.users) == 0)
+
+        owner_user = User.query.filter(User.email == 'admin@admin.com').one_or_none()
+        self.assertIsNotNone(owner_user)
+
+    def test_initialize_user_authentication_subproject_admin_exists(self):
+        password_salt = generate_salt()
+        owner_role = Role.get_role(Role.DefaultRole.OWNER)
+        User.create(
+            email='admin@admin.com',
+            password_hash=create_bcrypt_hash('admin', password_salt),
+            password_salt=password_salt,
+            username='admin',
+        )
+        initialize_user_authentication(ProjectType.SUB)
+
+        project_uuid = get_project_uuid()
+        project_uuid_truncated = project_uuid[:8]
+        owner_role = Role.get_role(f'{Role.DefaultRole.OWNER}_{project_uuid_truncated}')
         self.assertTrue(len(owner_role.users) == 0)
 
         owner_user = User.query.filter(User.email == 'admin@admin.com').one_or_none()

mage_ai/tests/services/k8s/test_job_manager.py

@@ -7,6 +7,10 @@ from kubernetes.client import (
     V1Container,
     V1EnvFromSource,
     V1EnvVar,
+    V1NodeAffinity,
+    V1NodeSelector,
+    V1NodeSelectorRequirement,
+    V1NodeSelectorTerm,
     V1Pod,
     V1PodSpec,
     V1SecretEnvSource,
@@ -377,13 +381,12 @@ class JobManagerTests(TestCase):
 
         # Pod
         self.assertEqual(job.spec.template.spec.affinity, V1Affinity(
-            **dict(
-                node_affinity=dict(
-                    required_during_scheduling_ignored_during_execution=dict(
+                node_affinity=V1NodeAffinity(
+                    required_during_scheduling_ignored_during_execution=V1NodeSelector(
                         node_selector_terms=[
-                            dict(
+                            V1NodeSelectorTerm(
                                 match_expressions=[
-                                    dict(
+                                    V1NodeSelectorRequirement(
                                         key='kubernetes.io/os',
                                         operator='In',
                                         values=['linux']
@@ -394,7 +397,7 @@ class JobManagerTests(TestCase):
                     )
                 )
             )
-        ))
+        )
 
         self.assertEqual(job.spec.template.spec.service_account_name, 'secretaccount')
         self.assertEqual(job.spec.template.spec.image_pull_secrets,

mage_ai/version_control/branch/utils.py

@@ -1,12 +1,13 @@
 from typing import List
 
-from mage_ai.data_preparation.git.api import push_raw
 from mage_ai.shared.strings import capitalize_remove_underscore_lower
 
 
 async def push(branch) -> List[str]:
     from git.exc import GitCommandError
 
+    from mage_ai.data_preparation.git.api import push_raw
+
     project = branch.project
     user = project.user
     provider = user.provider

mage_ai/version_control/models.py

@@ -10,11 +10,12 @@ from dataclasses import dataclass, field
 from typing import Any, Callable, Dict, List
 
 from mage_ai.authentication.oauth.constants import ProviderName
-from mage_ai.data_preparation.git.api import get_access_token_for_user, get_user
+from mage_ai.data_preparation.git.api import get_user
 from mage_ai.data_preparation.git.utils import (
     execute_on_remote_branch,
     get_access_token,
     get_auth_type_from_url,
+    get_oauth_access_token_for_user,
     get_provider_from_remote_url,
 )
 from mage_ai.data_preparation.preferences import Preferences, get_preferences
@@ -544,7 +545,7 @@ class ProjectUser(BaseVersionControl):
     @property
     def access_token(self) -> str:
         if self.user and self.provider:
-            access_token = get_access_token_for_user(self.user, provider=self.provider.name)
+            access_token = get_oauth_access_token_for_user(self.user, provider=self.provider.name)
             if access_token:
                 return access_token.token
 

{mage_ai-0.9.68.dist-info → mage_ai-0.9.70.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: mage-ai
-Version: 0.9.68
+Version: 0.9.70
 Summary: Mage is a tool for building and deploying data pipelines.
 Home-page: https://github.com/mage-ai/mage-ai
 Author: Mage
@@ -24,14 +24,14 @@ Requires-Dist: cryptography ==41.0.6
 Requires-Dist: dask >=2022.2.0
 Requires-Dist: datadog ==0.44.0
 Requires-Dist: freezegun ==1.2.2
-Requires-Dist: great-expectations ==0.15.50
+Requires-Dist: great-expectations ==0.18.12
 Requires-Dist: httpx ==0.23.1
 Requires-Dist: inflection ==0.5.1
 Requires-Dist: ipykernel ==6.15.0
 Requires-Dist: ipython ==8.10.0
 Requires-Dist: itsdangerous ~=1.1.0
 Requires-Dist: joblib >=1.1.0
-Requires-Dist: jupyter-server-proxy ==3.2.1
+Requires-Dist: jupyter-server-proxy ==3.2.3
 Requires-Dist: jupyter-server ==1.23.5
 Requires-Dist: jupyter-client ==7.4.4
 Requires-Dist: ldap3 ==2.9.1
@@ -57,7 +57,7 @@ Requires-Dist: terminado ==0.17.1
 Requires-Dist: thefuzz[speedup] ==0.19.0
 Requires-Dist: tornado ==6.3.3
 Requires-Dist: typer[all] ==0.9.0
-Requires-Dist: typing-extensions ==4.5.0
+Requires-Dist: typing-extensions ==4.10.0
 Requires-Dist: watchdog ==4.0.0
 Provides-Extra: ai
 Requires-Dist: astor >=0.8.1 ; extra == 'ai'
@@ -98,7 +98,7 @@ Requires-Dist: google-cloud-iam ~=2.13.0 ; extra == 'all'
 Requires-Dist: google-cloud-pubsub ~=2.19.0 ; extra == 'all'
 Requires-Dist: google-cloud-run ~=0.10.1 ; extra == 'all'
 Requires-Dist: google-cloud-storage ~=2.5.0 ; extra == 'all'
-Requires-Dist: great-expectations ==0.15.50 ; extra == 'all'
+Requires-Dist: great-expectations ==0.18.12 ; extra == 'all'
 Requires-Dist: gspread ==5.7.2 ; extra == 'all'
 Requires-Dist: influxdb-client ==1.36.1 ; extra == 'all'
 Requires-Dist: kafka-python ==2.0.2 ; extra == 'all'