mage-ai 0.9.69__py3-none-any.whl → 0.9.71__py3-none-any.whl

mage_ai/tests/services/k8s/test_job_manager.py

@@ -7,6 +7,11 @@ from kubernetes.client import (
     V1Container,
     V1EnvFromSource,
     V1EnvVar,
+    V1JobSpec,
+    V1NodeAffinity,
+    V1NodeSelector,
+    V1NodeSelectorRequirement,
+    V1NodeSelectorTerm,
     V1Pod,
     V1PodSpec,
     V1SecretEnvSource,
@@ -54,6 +59,10 @@ MOCK_POD_CONFIG = V1Pod(
         volumes=[],
     )
 )
+MOCK_JOB_CONFIG = V1JobSpec(template=MOCK_POD_CONFIG,
+                            active_deadline_seconds=120,
+                            backoff_limit=5,
+                            ttl_seconds_after_finished=86400)
 
 
 class TestK8sUtilities(TestCase):
@@ -223,7 +232,9 @@ class JobManagerTests(TestCase):
         )
         mock_client.V1JobSpec.assert_called_once_with(
             template=mock_v1_pod_template_spec,
+            active_deadline_seconds=None,
             backoff_limit=0,
+            ttl_seconds_after_finished=None
         )
         mock_client.V1Job.assert_called_once_with(
             api_version='batch/v1',
@@ -251,6 +262,7 @@ class JobManagerTests(TestCase):
             'env': [dict(name='spark_host', value='127.0.0.1')],
             'resources': dict(limits={'cpu': '1000m'}, requests={'cpu': '500m'}),
         }
+        k8s_config.job_config = MOCK_JOB_CONFIG
 
         # Call the method to create the job object
         command = "echo 'hello world'"
@@ -351,6 +363,11 @@ class JobManagerTests(TestCase):
                     },
                 ],
             },
+            'job': {
+                'active_deadline_seconds': 0,
+                'backoff_limit': 3,
+                'ttl_seconds_after_finished': 100
+            }
         }
         mock_getenv.return_value = 'pod_name'
 
@@ -377,13 +394,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 +410,7 @@ class JobManagerTests(TestCase):
                     )
                 )
             )
-        ))
+        )
 
         self.assertEqual(job.spec.template.spec.service_account_name, 'secretaccount')
         self.assertEqual(job.spec.template.spec.image_pull_secrets,
@@ -425,6 +441,11 @@ class JobManagerTests(TestCase):
             )
         )
 
+        # Job
+        self.assertEqual(job.spec.active_deadline_seconds, 0)
+        self.assertEqual(job.spec.backoff_limit, 3)
+        self.assertEqual(job.spec.ttl_seconds_after_finished, 100)
+
     @patch('mage_ai.services.k8s.job_manager.client')
     @patch('mage_ai.services.k8s.job_manager.config')
     def test_create_job(self, mock_config, mock_client):

mage_ai/tests/streaming/sinks/test_oracledb.py

@@ -0,0 +1,38 @@
+from unittest.mock import patch
+
+from mage_ai.streaming.sinks.oracledb import OracleDbSink
+from mage_ai.tests.base_test import TestCase
+
+
+class OracleDbTests(TestCase):
+    def test_init(self):
+        with patch.object(OracleDbSink, 'init_client') as mock_init_client:
+            OracleDbSink(dict(
+                connector_type='oracledb',
+                user='test',
+                password='123456',
+                host='oracle-test',
+                port=1521,
+                table='test_table',
+                service_name='test_service',
+                mode='thin',
+            ))
+            mock_init_client.assert_called_once()
+
+    def test_init_invalid_config(self):
+        with patch.object(OracleDbSink, 'init_client') as mock_init_client:
+            with self.assertRaises(Exception) as context:
+                OracleDbSink(dict(
+                    connector_type='oracledb',
+                    user='test',
+                    password='123456',
+                    host='oracle-test',
+                    port=1521,
+                    table='test_table',
+                    mode='thin',
+                ))
+            self.assertTrue(
+                '__init__() missing 1 required positional argument: \'service_name\''
+                in str(context.exception),
+            )
+            self.assertEqual(mock_init_client.call_count, 0)

mage_ai/tests/test_shared.py

@@ -1,3 +1,7 @@
+from scipy.sparse import csr_matrix
+
+from mage_ai.data.models.generator import DataGenerator
+from mage_ai.data.tabular.mocks import create_dataframe
 from mage_ai.data_cleaner.shared.utils import wrap_column_name
 from mage_ai.tests.base_test import TestCase
 
@@ -22,3 +26,60 @@ class SharedTests(TestCase):
         ]
         for name, expected_name in zip(column_names, expected_names):
             self.assertEqual(wrap_column_name(name), expected_name)
+
+
+def build_pandas(*args, **kwargs):
+    return create_dataframe(n_rows=1_000, use_pandas=True)
+
+
+def build_pandas_series(*args, **kwargs):
+    df = create_dataframe(n_rows=2_000, use_pandas=True)
+    return df[df.columns[0]]
+
+
+def build_polars(*args, **kwargs):
+    return create_dataframe(n_rows=3_000)
+
+
+def build_polars_series(*args, **kwargs):
+    df = create_dataframe(n_rows=4_000)
+    return df[df.columns[0]]
+
+
+def build_iterable(*args, **kwargs):
+    return [i for i in range(5_000)]
+
+
+def build_matrix_sparse(*args, **kwargs):
+    row_count = 6000
+
+    data = []
+    rows = []
+    cols = []
+
+    for i in range(int(row_count / 2)):
+        # Replicate the 3x3 pattern for each block of 2 rows
+        data.extend([1, 2, 1, 2, 3, 2])  # Replicated entries for simplicity
+        rows.extend([
+            i * 2,
+            i * 2,
+            i * 2 + 1,
+            i * 2 + 1,
+            i * 2 + 1,
+            i * 2 + 1,
+        ])  # Repeat every two rows
+        cols.extend([0, 2, 2, 0, 1, 2])  # Columns indices are consistent with your pattern
+
+    # Create the CSR matrix with specified shape
+    return csr_matrix((data, (rows, cols)), shape=(row_count, 3))
+
+
+def build_list_complex(*args, **kwargs):
+    return [DataGenerator([i for i in range(100)]) for i in range(100)] + [
+        build_pandas(),
+        build_pandas_series(),
+        build_polars(),
+        build_polars_series(),
+        build_iterable(),
+        build_matrix_sparse(),
+    ]

mage_ai/usage_statistics/logger.py

@@ -45,8 +45,13 @@ from mage_ai.usage_statistics.utils import build_event_data_for_chart
 
 
 class UsageStatisticLogger():
-    def __init__(self, project=None):
-        self.project = project or Project()
+    def __init__(
+        self,
+        context_data: Dict = None,
+        project=None,
+        repo_path: str = None,
+    ):
+        self.project = project or Project(context_data=context_data, repo_path=repo_path)
 
     async def block_create(
         self,

mage_ai/utils/code.py

@@ -2,12 +2,14 @@ import importlib
 import re
 from typing import Tuple
 
+from jupyter_client import KernelClient
+
 from mage_ai.autocomplete.utils import extract_all_imports
 from mage_ai.settings.repo import get_repo_path
 
 
-def reload_all_repo_modules(content: str) -> None:
-    parts = get_repo_path().split('/')
+def reload_all_repo_modules(content: str, client: KernelClient) -> None:
+    parts = get_repo_path(root_project=True).split('/')
     project_name = parts[-1]
 
     for line in extract_all_imports(content):
@@ -21,6 +23,12 @@ def reload_all_repo_modules(content: str) -> None:
             m1, m2 = matches[0]
             try:
                 importlib.reload(importlib.import_module(m1 or m2))
+                client.execute(
+                    f"""
+import importlib
+importlib.reload(importlib.import_module("{m1}" or "{m2}"))
+                """
+                )
             except Exception:
                 pass
 
@@ -35,34 +43,40 @@ def extract_decorated_function(code: str, decorated_function_name: str) -> Tuple
     number_of_lines = len(lines)
 
     for idx, line in enumerate(lines):
-        if line and \
-                span_current_start is not None and \
-                span_current_def is not None:
+        if line and span_current_start is not None and span_current_def is not None:
 
             if line.startswith('@') or line.startswith('def '):
-                spans.append((
-                    span_current_start,
-                    idx - 1,
-                ))
+                spans.append(
+                    (
+                        span_current_start,
+                        idx - 1,
+                    )
+                )
                 span_current_start = None
                 span_current_def = None
 
-        if line and \
-                span_current_start is not None and \
-                span_current_def is None and \
-                line.startswith('def '):
+        if (
+            line
+            and span_current_start is not None
+            and span_current_def is None
+            and line.startswith('def ')
+        ):
 
             span_current_def = idx
 
-        if line and \
-                span_current_start is None and \
-                line.startswith(f'@{decorated_function_name}'):
+        if (
+            line
+            and span_current_start is None
+            and line.startswith(f'@{decorated_function_name}')
+        ):
 
             span_current_start = idx
 
-        if span_current_start is not None and \
-                span_current_def is not None and \
-                idx == number_of_lines - 1:
+        if (
+            span_current_start is not None
+            and span_current_def is not None
+            and idx == number_of_lines - 1
+        ):
 
             spans.append((span_current_start, idx))
 

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.69.dist-info → mage_ai-0.9.71.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: mage-ai
-Version: 0.9.69
+Version: 0.9.71
 Summary: Mage is a tool for building and deploying data pipelines.
 Home-page: https://github.com/mage-ai/mage-ai
 Author: Mage
@@ -14,6 +14,7 @@ License-File: LICENSE
 Requires-Dist: Faker ==4.14.0
 Requires-Dist: GitPython ==3.1.41
 Requires-Dist: Jinja2 ==3.1.3
+Requires-Dist: Pillow ==10.3.0
 Requires-Dist: PyGithub ==1.59.0
 Requires-Dist: PyJWT ==2.6.0
 Requires-Dist: aiofiles ==22.1.0
@@ -31,18 +32,20 @@ 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
+Requires-Dist: memory-profiler
 Requires-Dist: newrelic ==8.8.0
 Requires-Dist: numpy >=1.22.0
 Requires-Dist: pandas >=1.3.0
 Requires-Dist: polars <0.19.2,>=0.18.0
 Requires-Dist: protobuf ~=4.21.12
+Requires-Dist: psutil ==5.9.8
 Requires-Dist: pyarrow ==14.0.1
 Requires-Dist: python-dateutil ==2.8.2
-Requires-Dist: pytz ==2022.2.1
+Requires-Dist: pytz >=2022.2.1
 Requires-Dist: pyyaml ~=6.0
 Requires-Dist: redis ~=5.0.1
 Requires-Dist: requests ~=2.31.0