mage-ai 0.8.11__py3-none-any.whl → 0.8.13__py3-none-any.whl

mage_ai/api/policies/BackfillPolicy.py

@@ -22,7 +22,10 @@ BackfillPolicy.allow_actions([
     OauthScope.CLIENT_PRIVATE,
 ], condition=lambda policy: policy.has_at_least_editor_role())
 
-BackfillPolicy.allow_read([] + BackfillPresenter.default_attributes, scopes=[
+BackfillPolicy.allow_read([
+    'pipeline_run_dates',
+    'total_run_count',
+] + BackfillPresenter.default_attributes, scopes=[
     OauthScope.CLIENT_PRIVATE,
 ], on_action=[
     constants.DETAIL,

mage_ai/api/policies/BlockPolicy.py

@@ -38,6 +38,7 @@ BlockPolicy.allow_read(BlockPresenter.default_attributes + [], scopes=[
 ], condition=lambda policy: policy.has_at_least_editor_role())
 
 BlockPolicy.allow_read([
+    'bookmarks',
     'content',
     'outputs',
 ] + BlockPresenter.default_attributes, scopes=[
@@ -68,9 +69,11 @@ BlockPolicy.allow_write([
 
 BlockPolicy.allow_write([
     'all_upstream_blocks_executed',
+    'bookmark_values',
     'color',
     'configuration',
     'content',
+    'destination_table',
     'downstream_blocks',
     'executor_config',
     'executor_type',
@@ -80,6 +83,7 @@ BlockPolicy.allow_write([
     'name',
     'outputs',
     'status',
+    'tap_stream_id',
     'type',
     'upstream_blocks',
     'uuid',
@@ -88,3 +92,12 @@ BlockPolicy.allow_write([
 ], on_action=[
     constants.UPDATE,
 ], condition=lambda policy: policy.has_at_least_editor_role())
+
+BlockPolicy.allow_query([
+    'destination_table',
+    'state_stream',
+], scopes=[
+    OauthScope.CLIENT_PRIVATE,
+], on_action=[
+    constants.DETAIL,
+], condition=lambda policy: policy.has_at_least_viewer_role())

mage_ai/api/policies/SyncPolicy.py

@@ -0,0 +1,30 @@
+from mage_ai.api.oauth_scope import OauthScope
+from mage_ai.api.operations import constants
+from mage_ai.api.policies.BasePolicy import BasePolicy
+from mage_ai.api.presenters.SyncPresenter import SyncPresenter
+
+
+class SyncPolicy(BasePolicy):
+    pass
+
+
+SyncPolicy.allow_read(SyncPresenter.default_attributes, scopes=[
+    OauthScope.CLIENT_PRIVATE,
+], on_action=[
+    constants.CREATE,
+    constants.UPDATE,
+], condition=lambda policy: policy.has_at_least_viewer_role())
+
+SyncPolicy.allow_write(SyncPresenter.default_attributes, scopes=[
+    OauthScope.CLIENT_PRIVATE,
+], on_action=[
+    constants.CREATE,
+], condition=lambda policy: policy.has_at_least_editor_role())
+
+SyncPolicy.allow_write([
+    'action_type',
+], scopes=[
+    OauthScope.CLIENT_PRIVATE,
+], on_action=[
+    constants.UPDATE,
+], condition=lambda policy: policy.has_at_least_editor_role())

mage_ai/api/presenters/BlockPresenter.py

@@ -30,13 +30,24 @@ class BlockPresenter(BasePresenter):
         if display_format in [constants.CREATE, constants.UPDATE]:
             return self.model.to_dict(include_content=True)
         elif display_format in [constants.DETAIL, 'dbt']:
-            include_outputs = kwargs.get('query', {}).get('include_outputs', [True])
+            query = kwargs.get('query', {})
+
+            include_outputs = query.get('include_outputs', [True])
             if include_outputs:
                 include_outputs = include_outputs[0]
 
+            state_stream = query.get('state_stream', [None])
+            if state_stream:
+                state_stream = state_stream[0]
+            destination_table = query.get('destination_table', [None])
+            if destination_table:
+                destination_table = destination_table[0]
+
             data = self.model.to_dict(
+                destination_table=destination_table,
                 include_content=True,
                 include_outputs=include_outputs,
+                state_stream=state_stream,
             )
 
             if 'dbt' == display_format:

mage_ai/api/presenters/SyncPresenter.py

@@ -0,0 +1,11 @@
+from mage_ai.api.presenters.BasePresenter import BasePresenter
+
+
+class SyncPresenter(BasePresenter):
+    default_attributes = [
+        'branch',
+        'remote_repo_link',
+        'repo_path',
+        'sync_on_pipeline_run',
+        'type',
+    ]

mage_ai/api/resources/BlockResource.py

@@ -78,4 +78,11 @@ class BlockResource(GenericResource):
 
     @safe_db_query
     def update(self, payload, **kwargs):
-        self.model.update(payload)
+        query = kwargs.get('query', {})
+        update_state = query.get('update_state', [False])
+        if update_state:
+            update_state = update_state[0]
+        self.model.update(
+            payload,
+            update_state=update_state,
+        )

mage_ai/api/resources/SyncResource.py

@@ -0,0 +1,43 @@
+from mage_ai.api.resources.GenericResource import GenericResource
+from mage_ai.data_preparation.preferences import get_preferences
+from mage_ai.data_preparation.sync import SyncType, SyncConfig
+from mage_ai.data_preparation.sync.git_sync import GitSync
+
+
+class SyncResource(GenericResource):
+    @classmethod
+    def collection(self, query, meta, user, **kwargs):
+        preferences = get_preferences()
+        sync_config = preferences.sync_config
+        return self.build_result_set(
+            [sync_config],
+            user,
+            **kwargs,
+        )
+
+    @classmethod
+    def create(self, payload, user, **kwargs):
+        sync_config = SyncConfig.load(config=payload)
+        get_preferences().update_preferences(
+            dict(sync_config=sync_config.to_dict())
+        )
+
+        if sync_config.type == SyncType.GIT:
+            # set up Git repo
+            GitSync(sync_config)
+
+        return self(get_preferences().sync_config, user, **kwargs)
+
+    @classmethod
+    def member(self, pk, user, **kwargs):
+        return self(get_preferences().sync_config, user, **kwargs)
+
+    def update(self, payload, **kwargs):
+        sync_config = SyncConfig.load(config=self.model)
+        action_type = payload.get('action_type')
+        if action_type == 'sync_data':
+            if sync_config.type == SyncType.GIT:
+                sync = GitSync(sync_config)
+                sync.sync_data()
+
+        return self

mage_ai/data_preparation/executors/pipeline_executor.py

@@ -1,6 +1,7 @@
 from mage_ai.data_preparation.logging.logger import DictLogger
 from mage_ai.data_preparation.logging.logger_manager_factory import LoggerManagerFactory
 from mage_ai.data_preparation.models.pipeline import Pipeline
+from mage_ai.shared.hash import merge_dict
 from typing import Dict
 import asyncio
 
@@ -33,3 +34,11 @@ class PipelineExecutor:
             update_status=update_status,
         ))
         self.logger_manager.output_logs_to_destination()
+
+    def _build_tags(self, **kwargs):
+        default_tags = dict(
+            pipeline_uuid=self.pipeline.uuid,
+        )
+        if kwargs.get('pipeline_run_id'):
+            default_tags['pipeline_run_id'] = kwargs.get('pipeline_run_id')
+        return merge_dict(kwargs.get('tags', {}), default_tags)

mage_ai/data_preparation/executors/streaming_pipeline_executor.py

@@ -69,10 +69,12 @@ class StreamingPipelineExecutor(PipelineExecutor):
         # TODOs:
         # 1. Support multiple sources and sinks
         # 2. Support flink pipeline
+
+        tags = self._build_tags(**kwargs)
         if build_block_output_stdout:
             stdout = build_block_output_stdout(self.pipeline.uuid)
         else:
-            stdout = StreamToLogger(self.logger)
+            stdout = StreamToLogger(self.logger, logging_tags=tags)
         try:
             with redirect_stdout(stdout):
                 with redirect_stderr(stdout):

mage_ai/data_preparation/models/block/__init__.py

@@ -1263,6 +1263,7 @@ df = get_variable('{self.pipeline.uuid}', '{block_uuid}', 'df')
         include_outputs=False,
         sample_count=None,
         check_if_file_exists: bool = False,
+        **kwargs,
     ):
         data = self.to_dict_base()
         if include_content:
@@ -1280,6 +1281,7 @@ df = get_variable('{self.pipeline.uuid}', '{block_uuid}', 'df')
                         'Delete the current block to remove it from the pipeline or write code ' +
                         f'and save the pipeline to create a new file at {file_path}.',
                     )
+
         return data
 
     async def to_dict_async(
@@ -1314,7 +1316,7 @@ df = get_variable('{self.pipeline.uuid}', '{block_uuid}', 'df')
 
         return data
 
-    def update(self, data):
+    def update(self, data, **kwargs):
         if 'name' in data and data['name'] != self.name:
             self.__update_name(data['name'])
         if (
@@ -1338,6 +1340,7 @@ df = get_variable('{self.pipeline.uuid}', '{block_uuid}', 'df')
             if self.has_callback:
                 CallbackBlock.create(self.uuid)
             self.__update_pipeline_block()
+
         return self
 
     def update_upstream_blocks(self, upstream_blocks: List[Any]) -> None:
@@ -1505,8 +1508,6 @@ df = get_variable('{self.pipeline.uuid}', '{block_uuid}', 'df')
             raise Exception(f'Failed to pass tests for block {self.uuid}')
 
     def analyze_outputs(self, variable_mapping, shape_only: bool = False):
-        from mage_ai.data_cleaner.data_cleaner import clean as clean_data
-
         if self.pipeline is None:
             return
         for uuid, data in variable_mapping.items():
@@ -1532,6 +1533,7 @@ df = get_variable('{self.pipeline.uuid}', '{block_uuid}', 'df')
                 else:
                     data_for_analysis = data.reset_index(drop=True)
                 try:
+                    from mage_ai.data_cleaner.data_cleaner import clean as clean_data
                     analysis = clean_data(
                         data_for_analysis,
                         df_original=data,

mage_ai/data_preparation/models/block/integration/__init__.py

@@ -314,6 +314,65 @@ class SourceBlock(IntegrationBlock):
 
 
 class DestinationBlock(IntegrationBlock):
+    def to_dict(
+        self,
+        include_content=False,
+        include_outputs=False,
+        sample_count=None,
+        check_if_file_exists: bool = False,
+        destination_table: str = None,
+        state_stream: str = None,
+    ):
+        data = {}
+        if state_stream and destination_table:
+            from mage_ai.data_preparation.models.pipelines.integration_pipeline \
+                import IntegrationPipeline
+            integration_pipeline = IntegrationPipeline(self.pipeline.uuid)
+            destination_state_file_path = integration_pipeline.destination_state_file_path(
+                destination_table=destination_table,
+                stream=state_stream,
+            )
+            if os.path.isfile(destination_state_file_path):
+                with open(destination_state_file_path, 'r') as f:
+                    text = f.read()
+                    d = json.loads(text) if text else {}
+                    bookmark_values = d.get('bookmarks', {}).get(state_stream)
+                    data['bookmarks'] = bookmark_values
+
+        return merge_dict(
+            super().to_dict(
+                include_content,
+                include_outputs,
+                sample_count,
+                check_if_file_exists,
+            ),
+            data,
+        )
+
+    def update(self, data, update_state=False):
+        if update_state:
+            from mage_ai.data_preparation.models.pipelines.integration_pipeline \
+                import IntegrationPipeline
+            from mage_integrations.destinations.utils \
+                import update_destination_state_bookmarks
+
+            integration_pipeline = IntegrationPipeline(self.pipeline.uuid)
+            tap_stream_id = data.get('tap_stream_id')
+            destination_table = data.get('destination_table')
+            bookmark_values = data.get('bookmark_values', {})
+            if tap_stream_id and destination_table:
+                destination_state_file_path = integration_pipeline.destination_state_file_path(
+                    destination_table=destination_table,
+                    stream=tap_stream_id,
+                )
+                update_destination_state_bookmarks(
+                    destination_state_file_path,
+                    tap_stream_id,
+                    bookmark_values=bookmark_values
+                )
+
+        return super().update(data)
+
     def output_variables(self, execution_partition: str = None) -> List[str]:
         return []
 

mage_ai/data_preparation/models/constants.py

@@ -2,6 +2,7 @@ from enum import Enum
 
 PIPELINES_FOLDER = 'pipelines'
 PIPELINE_CONFIG_FILE = 'metadata.yaml'
+PREFERENCES_FILE = '.preferences.yaml'
 DATA_INTEGRATION_CATALOG_FILE = 'data_integration_catalog.json'
 
 

mage_ai/data_preparation/preferences.py

@@ -0,0 +1,43 @@
+from typing import Dict
+from mage_ai.data_preparation.models.constants import PREFERENCES_FILE
+from mage_ai.data_preparation.repo_manager import get_repo_path
+import os
+import traceback
+import yaml
+
+
+class Preferences:
+    def __init__(self, repo_path: str = None, config_dict: Dict = None):
+        self.repo_path = repo_path or get_repo_path()
+        self.preferences_file_path = \
+            os.path.join(self.repo_path, PREFERENCES_FILE)
+
+        try:
+            if not config_dict:
+                if os.path.exists(self.preferences_file_path):
+                    with open(self.preferences_file_path) as f:
+                        preferences = yaml.full_load(f.read()) or {}
+                else:
+                    preferences = dict()
+            else:
+                preferences = config_dict
+
+            self.sync_config = preferences.get('sync_config', dict())
+        except Exception:
+            traceback.print_exc()
+            pass
+
+    def update_preferences(self, updates: Dict):
+        preferences = self.to_dict()
+        preferences.update(updates)
+        with open(self.preferences_file_path, 'w') as f:
+            yaml.dump(preferences, f)
+
+    def to_dict(self) -> Dict:
+        return dict(
+            sync_config=self.sync_config
+        )
+
+
+def get_preferences(repo_path=None) -> Preferences:
+    return Preferences(repo_path=repo_path)

mage_ai/data_preparation/sync/__init__.py

@@ -0,0 +1,17 @@
+from dataclasses import dataclass
+from enum import Enum
+from mage_ai.shared.config import BaseConfig
+import os
+
+
+class SyncType(str, Enum):
+    GIT = 'git'
+
+
+@dataclass
+class SyncConfig(BaseConfig):
+    type: SyncType
+    remote_repo_link: str
+    repo_path: str = os.getcwd()
+    branch: str = 'main'
+    sync_on_pipeline_run: bool = False

mage_ai/data_preparation/sync/base_sync.py

@@ -0,0 +1,7 @@
+from abc import ABC, abstractmethod
+
+
+class BaseSync(ABC):
+    @abstractmethod
+    def sync_data(self) -> None:
+        pass

mage_ai/data_preparation/sync/git_sync.py

@@ -0,0 +1,26 @@
+from mage_ai.data_preparation.sync import SyncConfig
+from mage_ai.data_preparation.sync.base_sync import BaseSync
+from mage_ai.shared.logger import VerboseFunctionExec
+import git
+
+
+class GitSync(BaseSync):
+    def __init__(self, sync_config: SyncConfig):
+        self.remote_repo_link = sync_config.remote_repo_link
+        self.repo_path = sync_config.repo_path
+        self.branch = sync_config.branch
+        try:
+            self.repo = git.Repo(self.repo_path)
+        except git.exc.InvalidGitRepositoryError:
+            self.repo = git.Repo.init(self.repo_path)
+            self.repo.create_remote('origin', self.remote_repo_link)
+
+        self.origin = self.repo.remotes.origin
+
+    def sync_data(self):
+        with VerboseFunctionExec(
+            f'Syncing data with remote repo {self.remote_repo_link}',
+            verbose=True,
+        ):
+            self.origin.fetch(kill_after_timeout=60)
+            self.repo.git.reset('--hard', f'origin/{self.branch}')

mage_ai/data_preparation/templates/data_loaders/streaming/kafka.yaml

@@ -18,3 +18,8 @@ consumer_group: unique_consumer_group
 #   mechanism: "PLAIN"
 #   username: username
 #   password: password
+
+# Uncomment the config below to use protobuf schema to deserialize message
+# serde_config:
+#   serialization_method: PROTOBUF
+#   schema_classpath: "path.to.schema.SchemaClass"

mage_ai/io/export_utils.py

@@ -36,6 +36,7 @@ class PandasTypes(str, Enum):
     MIXED = 'mixed'
     MIXED_INTEGER = 'mixed-integer'
     MIXED_INTEGER_FLOAT = 'mixed-integer-float'
+    OBJECT = 'object'
     PERIOD = 'period'
     STRING = 'string'
     TIME = 'time'

mage_ai/io/postgres.py

@@ -4,8 +4,10 @@ from mage_ai.io.sql import BaseSQL
 from pandas import DataFrame, Series
 from psycopg2 import connect, _psycopg
 from sshtunnel import SSHTunnelForwarder
-import numpy as np
 from typing import Union, IO
+import json
+import numpy as np
+import pandas as pd
 
 
 class Postgres(BaseSQL):
@@ -138,6 +140,29 @@ class Postgres(BaseSQL):
             PandasTypes.UNKNOWN_ARRAY,
             PandasTypes.COMPLEX,
         ):
+            series = column[column.notnull()]
+            values = series.values
+
+            column_type = None
+
+            if len(values) >= 1:
+                value = values[0]
+                column_type = 'JSONB'
+
+                if type(value) is list:
+                    if len(value) >= 1:
+                        item = value[0]
+                        item_series = pd.Series(data=item)
+                        item_dtype = item_series.dtype
+                        if PandasTypes.OBJECT != item_dtype:
+                            item_type = self.get_type(item_series, item_dtype)
+                            column_type = f'{item_type}[]'
+                    else:
+                        column_type = 'text[]'
+
+            if column_type:
+                return column_type
+
             raise BadConversionError(
                 f'Cannot convert column \'{column.name}\' with data type \'{dtype}\' to '
                 'a PostgreSQL datatype.'
@@ -180,6 +205,8 @@ class Postgres(BaseSQL):
             return 'bigint'
         elif dtype == PandasTypes.EMPTY:
             return 'text'
+        elif PandasTypes.OBJECT == dtype:
+            return 'JSONB'
         else:
             print(f'Invalid datatype provided: {dtype}')
 
@@ -192,9 +219,24 @@ class Postgres(BaseSQL):
         full_table_name: str,
         buffer: Union[IO, None] = None
     ) -> None:
-        columns_names = ', '.join(df.columns)
-        df.to_csv(buffer, index=False, header=False, na_rep='')
+        df_ = df.copy()
+        columns = df_.columns
+
+        for col in columns:
+            print(col, df_[col].dtype)
+            if PandasTypes.OBJECT == df_[col].dtype:
+                df_[col] = df_[col].apply(lambda x: json.dumps(x))
+
+        df_.to_csv(
+            buffer,
+            header=False,
+            index=False,
+            na_rep='',
+        )
+
         buffer.seek(0)
+
+        columns_names = ', '.join(columns)
         cursor.copy_expert(f"""
 COPY {full_table_name} FROM STDIN (
     FORMAT csv
@@ -202,4 +244,4 @@ COPY {full_table_name} FROM STDIN (
     , NULL \'\'
     , FORCE_NULL({columns_names})
 );
-""", buffer)
+    """, buffer)

mage_ai/orchestration/db/models.py

@@ -76,6 +76,11 @@ class BaseModel(Base):
     def full_clean(self, **kwargs) -> None:
         pass
 
+    @classmethod
+    @safe_db_query
+    def get(self, uuid):
+        return self.query.get(uuid)
+
     def save(self, commit=True) -> None:
         self.session.add(self)
         if commit:
@@ -667,6 +672,15 @@ class Backfill(BaseModel):
 
     pipeline_runs = relationship('PipelineRun', back_populates='backfill')
 
+    @classmethod
+    @safe_db_query
+    def filter(self, pipeline_schedule_ids: List = None):
+        if pipeline_schedule_ids is not None:
+            return Backfill.query.filter(
+                Backfill.pipeline_schedule_id.in_(pipeline_schedule_ids),
+            )
+        return []
+
 
 class Secret(BaseModel):
     name = Column(String(255), unique=True)

mage_ai/orchestration/pipeline_scheduler.py

@@ -14,7 +14,10 @@ from mage_ai.data_preparation.models.block.utils import (
 from mage_ai.data_preparation.models.constants import PipelineType
 from mage_ai.data_preparation.models.pipeline import Pipeline
 from mage_ai.data_preparation.models.pipelines.integration_pipeline import IntegrationPipeline
+from mage_ai.data_preparation.preferences import get_preferences
 from mage_ai.data_preparation.repo_manager import get_repo_config, get_repo_path
+from mage_ai.data_preparation.sync import SyncConfig
+from mage_ai.data_preparation.sync.git_sync import GitSync
 from mage_ai.data_preparation.variable_manager import get_global_variables
 from mage_ai.orchestration.db.models import (
     Backfill,
@@ -58,6 +61,12 @@ class PipelineScheduler:
         )
 
     def start(self, should_schedule: bool = True) -> None:
+        if get_preferences().sync_config:
+            sync_config = SyncConfig.load(config=get_preferences().sync_config)
+            if sync_config.sync_on_pipeline_run:
+                sync = GitSync(sync_config)
+                sync.sync_data()
+
         if self.pipeline_run.status == PipelineRun.PipelineRunStatus.RUNNING:
             return
         self.pipeline_run.update(status=PipelineRun.PipelineRunStatus.RUNNING)
@@ -126,7 +135,7 @@ class PipelineScheduler:
                 )
                 self.logger_manager.output_logs_to_destination()
 
-                schedule = PipelineSchedule.query.get(
+                schedule = PipelineSchedule.get(
                     self.pipeline_run.pipeline_schedule_id,
                 )
 
@@ -735,9 +744,8 @@ def schedule_all():
     active_pipeline_schedules = \
         list(PipelineSchedule.active_schedules(pipeline_uuids=repo_pipelines))
 
-    backfills = Backfill.query.filter(
-        Backfill.pipeline_schedule_id.in_([ps.id for ps in active_pipeline_schedules]),
-    )
+    backfills = Backfill.filter(pipeline_schedule_ids=[ps.id for ps in active_pipeline_schedules])
+
     backfills_by_pipeline_schedule_id = index_by(
         lambda backfill: backfill.pipeline_schedule_id,
         backfills,

mage_ai/server/constants.py

@@ -12,4 +12,4 @@ DATAFRAME_OUTPUT_SAMPLE_COUNT = 10
 # Dockerfile depends on it because it runs ./scripts/install_mage.sh and uses
 # the last line to determine the version to install.
 VERSION = \
-'0.8.11'
+'0.8.13'