mage-ai 0.8.47__py3-none-any.whl → 0.8.49__py3-none-any.whl

mage_ai/api/resources/DataProviderResource.py

@@ -8,12 +8,24 @@ import yaml
 
 DATA_PROVIDERS = [
     DataSource.BIGQUERY,
+    DataSource.CLICKHOUSE,
     DataSource.MSSQL,
     DataSource.MYSQL,
     DataSource.POSTGRES,
     DataSource.REDSHIFT,
     DataSource.SNOWFLAKE,
+    DataSource.TRINO,
 ]
+DATA_PROVIDERS_NAME = {
+    DataSource.BIGQUERY: 'BigQuery',
+    DataSource.CLICKHOUSE: 'ClickHouse',
+    DataSource.MSSQL: 'Microsoft SQL Server',
+    DataSource.MYSQL: 'MySQL',
+    DataSource.POSTGRES: 'PostgreSQL',
+    DataSource.REDSHIFT: 'Redshift',
+    DataSource.SNOWFLAKE: 'Snowflake',
+    DataSource.TRINO: 'Trino',
+}
 
 
 class DataProviderResource(GenericResource):
@@ -27,7 +39,7 @@ class DataProviderResource(GenericResource):
                 print(exc)
 
         collection = [dict(
-            id=ds.title(),
+            id=DATA_PROVIDERS_NAME[ds.value],
             profiles=[p for p in profiles if p != 'version'],
             value=ds.value,
         ) for ds in DATA_PROVIDERS]

mage_ai/data_preparation/git/__init__.py

@@ -5,7 +5,6 @@ from mage_ai.orchestration.db.models.oauth import User
 from urllib.parse import urlparse
 import asyncio
 import base64
-import git
 import os
 import subprocess
 
@@ -15,6 +14,7 @@ REMOTE_NAME = 'mage-repo'
 
 class Git:
     def __init__(self, git_config: GitConfig):
+        import git
         self.remote_repo_link = git_config.remote_repo_link
         self.repo_path = git_config.repo_path or os.getcwd()
         os.makedirs(self.repo_path, exist_ok=True)

mage_ai/data_preparation/models/block/__init__.py

@@ -655,6 +655,7 @@ class Block:
         dynamic_block_uuid: str = None,
         dynamic_upstream_block_uuids: List[str] = None,
         run_settings: Dict = None,
+        **kwargs,
     ) -> Dict:
         try:
             if not run_all_blocks:
@@ -685,6 +686,7 @@ class Block:
                 dynamic_block_index=dynamic_block_index,
                 dynamic_upstream_block_uuids=dynamic_upstream_block_uuids,
                 run_settings=run_settings,
+                **kwargs,
             )
             block_output = self.post_process_output(output)
             variable_mapping = dict()
@@ -867,6 +869,7 @@ class Block:
         dynamic_block_index: int = None,
         dynamic_upstream_block_uuids: List[str] = None,
         run_settings: Dict = None,
+        **kwargs,
     ) -> Dict:
         # Add pipeline uuid and block uuid to global_vars
         global_vars = merge_dict(
@@ -922,6 +925,7 @@ class Block:
                 runtime_arguments=runtime_arguments,
                 upstream_block_uuids=upstream_block_uuids,
                 run_settings=run_settings,
+                **kwargs,
             )
 
         output_message = dict(output=outputs)
@@ -942,6 +946,7 @@ class Block:
         runtime_arguments: Dict = None,
         upstream_block_uuids: List[str] = None,
         run_settings: Dict = None,
+        **kwargs,
     ) -> List:
         decorated_functions = []
         test_functions = []

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

@@ -1,6 +1,7 @@
 from mage_ai.data_preparation.models.block import Block
 from mage_ai.data_preparation.models.block.sql import (
     bigquery,
+    clickhouse,
     mssql,
     mysql,
     postgres,
@@ -14,8 +15,12 @@ from mage_ai.data_preparation.models.block.sql.utils.shared import (
 )
 from mage_ai.data_preparation.models.constants import BlockType
 from mage_ai.data_preparation.repo_manager import get_repo_path
-from mage_ai.io.base import DataSource, ExportWritePolicy
-from mage_ai.io.config import ConfigFileLoader
+from mage_ai.io.base import (
+    DataSource,
+    ExportWritePolicy,
+    QUERY_ROW_LIMIT,
+)
+from mage_ai.io.config import ConfigFileLoader, ConfigKey
 from os import path
 from time import sleep
 from typing import Any, Dict, List
@@ -37,6 +42,7 @@ def execute_sql_code(
     global_vars: Dict = None,
     config_file_loader: Any = None,
     configuration: Dict = None,
+    test_execution: bool = False,
 ) -> List[Any]:
     configuration = configuration if configuration else block.configuration
     use_raw_sql = configuration.get('use_raw_sql')
@@ -66,6 +72,12 @@ def execute_sql_code(
         verbose=BlockType.DATA_EXPORTER == block.type,
     )
 
+    limit = int(configuration.get('limit') or QUERY_ROW_LIMIT)
+    if test_execution:
+        limit = min(limit, QUERY_ROW_LIMIT)
+    else:
+        limit = QUERY_ROW_LIMIT
+
     if DataSource.BIGQUERY.value == data_provider:
         from mage_ai.io.bigquery import BigQuery
 
@@ -111,12 +123,54 @@ def execute_sql_code(
                     try:
                         result = loader.load(
                             f'SELECT * FROM {database}.{schema}.{table_name}',
+                            limit=limit,
                             verbose=False,
                         )
                         return [result]
                     except Exception as err:
                         if '404' not in str(err):
                             raise err
+    elif DataSource.CLICKHOUSE.value == data_provider:
+        from mage_ai.io.clickhouse import ClickHouse
+
+        loader = ClickHouse.with_config(config_file_loader)
+        clickhouse.create_upstream_block_tables(
+            loader,
+            block,
+            configuration=configuration,
+            execution_partition=execution_partition,
+        )
+
+        query_string = clickhouse.interpolate_input_data(block, query)
+        query_string = interpolate_vars(
+            query_string, global_vars=global_vars)
+
+        database = database or 'default'
+
+        if use_raw_sql:
+            return execute_raw_sql(
+                loader,
+                block,
+                query_string,
+                configuration=configuration,
+                should_query=should_query,
+            )
+        else:
+            loader.export(
+                None,
+                table_name=table_name,
+                database=database,
+                query_string=query_string,
+                **kwargs_shared,
+            )
+
+            if should_query:
+                return [
+                    loader.load(
+                        f'SELECT * FROM {database}.{table_name}',
+                        verbose=False,
+                    ),
+                ]
     elif DataSource.MSSQL.value == data_provider:
         from mage_ai.io.mssql import MSSQL
 
@@ -155,6 +209,7 @@ def execute_sql_code(
                     return [
                         loader.load(
                             f'SELECT * FROM {table_name}',
+                            limit=limit,
                             verbose=False,
                         ),
                     ]
@@ -193,6 +248,7 @@ def execute_sql_code(
                     return [
                         loader.load(
                             f'SELECT * FROM {table_name}',
+                            limit=limit,
                             verbose=False,
                         ),
                     ]
@@ -231,6 +287,7 @@ def execute_sql_code(
                     return [
                         loader.load(
                             f'SELECT * FROM {schema}.{table_name}',
+                            limit=limit,
                             verbose=False,
                         ),
                     ]
@@ -269,6 +326,7 @@ def execute_sql_code(
                     return [
                             loader.load(
                                 f'SELECT * FROM {schema}.{table_name}',
+                                limit=limit,
                                 verbose=False,
                             ),
                         ]
@@ -316,6 +374,7 @@ def execute_sql_code(
                             database=database,
                             schema=schema,
                             table_name=table_name,
+                            limit=limit,
                             verbose=False,
                         ),
                     ]
@@ -344,21 +403,20 @@ def execute_sql_code(
             else:
                 loader.export(
                     None,
-                    table_name,
-                    database,
                     schema,
+                    table_name,
                     if_exists=export_write_policy,
                     query_string=query_string,
                     verbose=BlockType.DATA_EXPORTER == block.type,
                 )
 
                 if should_query:
+                    catalog = config_file_loader[ConfigKey.TRINO_CATALOG]
+
                     return [
                         loader.load(
-                            f'SELECT * FROM "{database}"."{schema}"."{table_name}"',
-                            database=database,
-                            schema=schema,
-                            table_name=table_name,
+                            f'SELECT * FROM "{catalog}"."{schema}"."{table_name}"',
+                            limit=limit,
                             verbose=False,
                         ),
                     ]
@@ -449,9 +507,11 @@ class SQLBlock(Block):
         global_vars=None,
         **kwargs,
     ) -> List:
+        test_execution = kwargs.get('test_execution')
         return execute_sql_code(
             self,
             custom_code or self.content,
             execution_partition=execution_partition,
             global_vars=global_vars,
+            test_execution=test_execution,
         )

mage_ai/data_preparation/models/block/sql/clickhouse.py

@@ -0,0 +1,38 @@
+from mage_ai.data_preparation.models.block.sql.utils.shared import (
+    create_upstream_block_tables as create_upstream_block_tables_orig,
+    interpolate_input,
+)
+from mage_ai.io.config import ConfigKey
+from typing import Dict
+
+
+def create_upstream_block_tables(
+    loader,
+    block,
+    cascade_on_drop: bool = False,
+    configuration: Dict = None,
+    execution_partition: str = None,
+    cache_upstream_dbt_models: bool = False,
+):
+    create_upstream_block_tables_orig(
+        loader,
+        block,
+        cascade_on_drop,
+        configuration,
+        execution_partition,
+        cache_upstream_dbt_models,
+        cache_keys=[
+            ConfigKey.CLICKHOUSE_DATABASE,
+            ConfigKey.CLICKHOUSE_HOST,
+            ConfigKey.CLICKHOUSE_PORT,
+        ],
+        no_schema=True,
+    )
+
+
+def interpolate_input_data(block, query):
+    return interpolate_input(
+        block,
+        query,
+        lambda db, schema, tn: tn,
+    )

mage_ai/data_preparation/models/block/utils.py

@@ -163,8 +163,10 @@ def create_block_runs_from_dynamic_block(
                         skip_creating_downstream = True
                         break
 
-                    down_uuids = [down.uuid for down in dynamic_ancestor.downstream_blocks]
-                    down_uuids_as_ancestors = [i for i in down_uuids if i in ancestors_uuids]
+                    down_uuids_as_ancestors = []
+                    for down in dynamic_ancestor.downstream_blocks:
+                        if down.uuid in ancestors_uuids and not should_reduce_output(down):
+                            down_uuids_as_ancestors.append(down.uuid)
                     skip_creating_downstream = len(down_uuids_as_ancestors) >= 2
 
                 # Only create downstream block runs if it doesn’t have dynamically created upstream

mage_ai/data_preparation/models/file.py

@@ -20,7 +20,7 @@ INACCESSIBLE_DIRS = frozenset(['__pycache__'])
 MAX_DEPTH = 30
 MAX_NUMBER_OF_FILE_VERSIONS = int(os.getenv('MAX_NUMBER_OF_FILE_VERSIONS', 100) or 100)
 
-PIPELINES_FOLDER_PREFIX = 'pipelines/'
+PIPELINES_FOLDER_PREFIX = f'pipelines{os.sep}'
 
 
 class File:
@@ -118,8 +118,9 @@ class File:
         filename: str,
     ) -> Tuple[str, str]:
         return os.path.join(
-            f'{repo_path}/{FILE_VERSIONS_DIR}',
-            dir_path.replace(repo_path, f'{repo_path}/{FILE_VERSIONS_DIR}'),
+            repo_path,
+            FILE_VERSIONS_DIR,
+            dir_path.replace(repo_path, os.path.join(repo_path, FILE_VERSIONS_DIR)),
             filename,
         )
 

mage_ai/data_preparation/models/pipeline.py

@@ -21,7 +21,7 @@ from mage_ai.data_preparation.templates.utils import copy_template_directory
 from mage_ai.data_preparation.variable_manager import VariableManager
 from mage_ai.orchestration.db import db_connection, safe_db_query
 from mage_ai.shared.array import find
-from mage_ai.shared.hash import extract, ignore_keys, merge_dict
+from mage_ai.shared.hash import extract, index_by, ignore_keys, merge_dict
 from mage_ai.shared.io import safe_write, safe_write_async
 from mage_ai.shared.strings import format_enum
 from mage_ai.shared.utils import clean_name
@@ -674,6 +674,12 @@ class Pipeline:
                 )
             should_save = True
 
+        blocks = data.get('blocks', [])
+
+        if blocks:
+            if not should_save and self.__update_block_order(blocks):
+                should_save = True
+
         if should_save:
             self.save()
 
@@ -682,8 +688,8 @@ class Pipeline:
 
             arr = []
 
-            if 'blocks' in data:
-                arr.append(('blocks', data['blocks'], self.blocks_by_uuid))
+            if blocks:
+                arr.append(('blocks', blocks, self.blocks_by_uuid))
             if 'widgets' in data:
                 arr.append(('widgets', data['widgets'], self.widgets_by_uuid))
 
@@ -778,6 +784,33 @@ class Pipeline:
                 if should_save_async:
                     await self.save_async(widget=widget)
 
+    def __update_block_order(self, blocks: List[Dict]) -> bool:
+        uuids_new = [b['uuid'] for b in blocks]
+        uuids_old = [b['uuid'] for b in self.block_configs]
+
+        min_length = min(len(uuids_new), len(uuids_old))
+
+        # If there are no blocks or the order has changed
+        if min_length == 0 or uuids_new[:min_length] == uuids_old[:min_length]:
+            return False
+
+        block_configs_by_uuids = index_by(lambda x: x['uuid'], self.block_configs)
+
+        block_configs = []
+        blocks_by_uuid = {}
+
+        for block_uuid in uuids_new:
+            if block_uuid in block_configs_by_uuids:
+                block_configs.append(block_configs_by_uuids[block_uuid])
+
+            if block_uuid in self.blocks_by_uuid:
+                blocks_by_uuid[block_uuid] = self.blocks_by_uuid[block_uuid]
+
+        self.block_configs = block_configs
+        self.blocks_by_uuid = blocks_by_uuid
+
+        return True
+
     def __add_block_to_mapping(
         self,
         blocks_by_uuid,

mage_ai/data_preparation/templates/repo/io_config.yaml

@@ -10,6 +10,12 @@ default:
   AZURE_CLIENT_SECRET: "{{ env_var('AZURE_CLIENT_SECRET') }}"
   AZURE_STORAGE_ACCOUNT_NAME: "{{ env_var('AZURE_STORAGE_ACCOUNT_NAME') }}"
   AZURE_TENANT_ID: "{{ env_var('AZURE_TENANT_ID') }}"
+  CLICKHOUSE_DATABASE: default
+  CLICKHOUSE_HOST: host.docker.internal
+  CLICKHOUSE_INTERFACE: http
+  CLICKHOUSE_PASSWORD: null
+  CLICKHOUSE_PORT: 8123
+  CLICKHOUSE_USERNAME: null
   GOOGLE_SERVICE_ACC_KEY:
     type: service_account
     project_id: project-id
@@ -46,3 +52,9 @@ default:
   SNOWFLAKE_DEFAULT_DB: optional_default_database
   SNOWFLAKE_DEFAULT_SCHEMA: optional_default_schema
   SNOWFLAKE_ROLE: role
+  TRINO_CATALOG: catalog
+  TRINO_HOST: '127.0.0.1'
+  TRINO_PASSWORD: pasword
+  TRINO_PORT: 8080
+  TRINO_SCHEMA: schema
+  TRINO_USER: user

mage_ai/io/base.py

@@ -14,6 +14,7 @@ QUERY_ROW_LIMIT = 10_000_000
 class DataSource(str, Enum):
     API = 'api'
     BIGQUERY = 'bigquery'
+    CLICKHOUSE = 'clickhouse'
     FILE = 'file'
     GOOGLE_CLOUD_STORAGE = 'google_cloud_storage'
     KAFKA = 'kafka'

mage_ai/io/clickhouse.py

@@ -0,0 +1,237 @@
+from mage_ai.io.base import BaseSQLDatabase, ExportWritePolicy, QUERY_ROW_LIMIT
+from mage_ai.io.config import BaseConfigLoader, ConfigKey
+from mage_ai.io.export_utils import (
+    clean_df_for_export,
+    infer_dtypes,
+)
+from pandas import DataFrame
+from typing import Dict, List, Union
+import clickhouse_connect
+
+
+class ClickHouse(BaseSQLDatabase):
+    """
+    Handles data transfer betwee a ClickHouse data warehouse and the Mage app.
+    """
+    def __init__(self, **kwargs) -> None:
+        """
+        Initializes settings for connecting to a ClickHouse warehouse.
+
+        To authenticate (and authorize) access to a ClickHouse warehouse,
+        credentials, i.e., username and password, must be provided.
+
+        All keyword arguments will be passed to the ClickHouse client.
+        """
+        if kwargs.get('verbose') is not None:
+            kwargs.pop('verbose')
+        super().__init__(verbose=kwargs.get('verbose', True))
+        with self.printer.print_msg('Connecting to ClickHouse'):
+            self.client = clickhouse_connect.get_client(**kwargs)
+
+    @classmethod
+    def with_config(cls, config: BaseConfigLoader) -> 'ClickHouse':
+        """
+        Initializes ClickHouse client from configuration loader
+
+        Args:
+            config (BaseConfigLoader): Configuration loader object
+        """
+        if ConfigKey.CLICKHOUSE_HOST not in config:
+            raise ValueError(
+                'No valid configuration settings found for ClickHouse. '
+                'You must specify host.'
+            )
+        return cls(
+            database=config[ConfigKey.CLICKHOUSE_DATABASE],
+            host=config[ConfigKey.CLICKHOUSE_HOST],
+            interface=config[ConfigKey.CLICKHOUSE_INTERFACE],
+            password=config[ConfigKey.CLICKHOUSE_PASSWORD],
+            port=config[ConfigKey.CLICKHOUSE_PORT],
+            username=config[ConfigKey.CLICKHOUSE_USERNAME],
+        )
+
+    def execute(self, command_string: str, **kwargs) -> None:
+        """
+        Sends command to the connected ClickHouse warehouse.
+
+        Args:
+            command_string (str): Command to execute on the ClickHouse warehouse.
+            **kwargs: Additional arguments to pass to command, such as configurations
+        """
+        with self.printer.print_msg(f'Executing query \'{command_string}\''):
+            command_string = self._clean_query(command_string)
+            self.client.command(command_string, **kwargs)
+
+    def execute_query(
+        self,
+        query: str,
+        parameters: Dict = None,
+        **kwargs,
+    ) -> DataFrame:
+        """
+        Sends query to the connected ClickHouse warehouse.
+
+        Args:
+            query (str): Query to execute on the ClickHouse warehouse.
+            **kwargs: Additional arguments to pass to query, such as query configurations
+        """
+        query = self._clean_query(query)
+        with self.printer.print_msg(f'Executing query \'{query}\''):
+            result = self.client.query_df(query, parameters=parameters)
+
+        return result
+
+    def execute_queries(
+        self,
+        queries: List[str],
+        query_variables: List[Dict] = None,
+        fetch_query_at_indexes: List[bool] = None,
+        **kwargs,
+    ) -> List:
+        results = []
+
+        for idx, query in enumerate(queries):
+            parameters = query_variables[idx] \
+                        if query_variables and idx < len(query_variables) \
+                        else {}
+            query = self._clean_query(query)
+
+            if fetch_query_at_indexes and idx < len(fetch_query_at_indexes) and \
+                    fetch_query_at_indexes[idx]:
+                result = self.client.query_df(query, parameters=parameters)
+            else:
+                result = self.client.command(query, parameters=parameters)
+
+            results.append(result)
+
+        return results
+
+    def load(
+        self,
+        query_string: str,
+        limit: int = QUERY_ROW_LIMIT,
+        display_query: Union[str, None] = None,
+        verbose: bool = True,
+        **kwargs,
+    ) -> DataFrame:
+        """
+        Loads data from ClickHouse into a Pandas data frame based on the query given.
+        This will fail if the query returns no data from the database. When a select query
+        is provided, this function will load at maximum 10,000,000 rows of data. To operate on more
+        data, consider performing data transformations in warehouse.
+
+        Args:
+            query_string (str): Query to fetch a table or subset of a table.
+            limit (int, Optional): The number of rows to limit the loaded dataframe to. Defaults to
+                                    10,000,000.
+            **kwargs: Additional arguments to pass to query, such as query configurations
+
+        Returns:
+            DataFrame: Data frame associated with the given query.
+        """
+
+        print_message = 'Loading data'
+        if verbose:
+            print_message += ' with query'
+
+            if display_query:
+                for line in display_query.split('\n'):
+                    print_message += f'\n{line}'
+            else:
+                print_message += f'\n{query_string}'
+
+        query_string = self._clean_query(query_string)
+
+        with self.printer.print_msg(print_message):
+            return self.client.query_df(
+                self._enforce_limit(query_string, limit), **kwargs
+            )
+
+    def export(
+        self,
+        df: DataFrame,
+        table_name: str,
+        database: str = 'default',
+        if_exists: str = 'append',
+        index: bool = False,
+        query_string: Union[str, None] = None,
+        create_table_statement: Union[str, None] = None,
+        verbose: bool = True,
+        **kwargs,
+    ) -> None:
+        """
+        Exports a Pandas data frame to a ClickHouse warehouse based on the table name.
+        If table doesn't exist, the table is automatically created.
+
+        Args:
+            df (DataFrame): Data frame to export
+            table_name (str): Name of the table to export data to (excluding database).
+            If this table exists, the table schema must match the data frame schema.
+            If this table doesn't exist, query_string must be specified to create the new table.
+            database (str): Name of the database in which the table is located.
+            if_exists (str, optional): Specifies export policy if table exists. Either
+                - `'fail'`: throw an error.
+                - `'replace'`: drops existing table and creates new table of same name.
+                - `'append'`: appends data frame to existing table. In this case the schema must
+                                match the original table.
+            Defaults to `'append'`.
+            **kwargs: Additional arguments to pass to writer
+        """
+
+        if type(df) is dict:
+            df = DataFrame([df])
+        elif type(df) is list:
+            df = DataFrame(df)
+
+        if not query_string:
+            if index:
+                df = df.reset_index()
+
+            dtypes = infer_dtypes(df)
+            df = clean_df_for_export(df, self.clean, dtypes)
+
+        def __process(database: Union[str, None]):
+
+            df_existing = self.client.query_df(f"""
+EXISTS TABLE {database}.{table_name}
+""")
+
+            table_exists = not df_existing.empty and df_existing.iloc[0, 0] == 1
+            should_create_table = not table_exists
+
+            if table_exists:
+                if ExportWritePolicy.FAIL == if_exists:
+                    raise ValueError(
+                        f'Table \'{table_name}\' already'
+                        ' exists in database {database}.',
+                    )
+                elif ExportWritePolicy.REPLACE == if_exists:
+                    self.client.command(
+                        f'DROP TABLE IF EXISTS {database}.{table_name}')
+                    should_create_table = True
+
+            if query_string:
+                self.client.command(f'USE {database}')
+
+                if should_create_table:
+                    self.client.command(f"""
+CREATE TABLE IF NOT EXISTS {database}.{table_name} ENGINE = Memory AS
+{query_string}
+""")
+                else:
+                    self.client.command(f"""
+INSERT INTO {database}.{table_name}
+{query_string}
+""")
+            else:
+                if should_create_table:
+                    self.client.command(create_table_statement)
+
+                self.client.insert_df(f'{database}.{table_name}', df)
+
+        if verbose:
+            with self.printer.print_msg(
+                    f'Exporting data to table \'{database}.{table_name}\''):
+                __process(database=database)
+        else:
+            __process(database=database)