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

mage_ai/io/bigquery.py

@@ -1,3 +1,4 @@
+import uuid
 from typing import Dict, List, Mapping, Union
 
 import numpy as np
@@ -14,6 +15,7 @@ from sqlglot import exp, parse_one
 
 from mage_ai.io.base import QUERY_ROW_LIMIT, BaseSQLDatabase, ExportWritePolicy
 from mage_ai.io.config import BaseConfigLoader, ConfigKey
+from mage_ai.io.constants import UNIQUE_CONFLICT_METHOD_UPDATE
 from mage_ai.io.export_utils import infer_dtypes
 from mage_ai.shared.custom_logger import DX_PRINTER
 from mage_ai.shared.environments import is_debug
@@ -200,6 +202,9 @@ WHERE TABLE_NAME = '{table_name}'
         overwrite_types: Dict = None,
         query_string: Union[str, None] = None,
         verbose: bool = True,
+        unique_conflict_method: str = None,
+        unique_constraints: List[str] = None,
+        write_disposition: str = None,
         **configuration_params,
     ) -> None:
         """
@@ -230,24 +235,23 @@ WHERE TABLE_NAME = '{table_name}'
         elif type(df) is list:
             df = DataFrame(df)
 
-        def __process(database: Union[str, None]):
-            if query_string:
-                parts = table_id.split('.')
-                if len(parts) == 2:
-                    schema, table_name = parts
-                elif len(parts) == 3:
-                    database, schema, table_name = parts
+        def __process(database: Union[str, None], write_disposition: str = None):
+            parts = table_id.split('.')
+            if len(parts) == 2:
+                schema, table_name = parts
+            elif len(parts) == 3:
+                database, schema, table_name = parts
 
-                df_existing = self.client.query(f"""
+            df_existing = self.client.query(f"""
 SELECT 1
 FROM `{database}.{schema}.__TABLES_SUMMARY__`
 WHERE table_id = '{table_name}'
 """).to_dataframe()
 
-                full_table_name = f'`{database}.{schema}.{table_name}`'
-
-                table_doesnt_exist = df_existing.empty
+            full_table_name = f'`{database}.{schema}.{table_name}`'
 
+            table_doesnt_exist = df_existing.empty
+            if query_string:
                 if ExportWritePolicy.FAIL == if_exists and not table_doesnt_exist:
                     raise ValueError(
                         f'Table \'{full_table_name}\' already exists in database.',
@@ -268,60 +272,129 @@ WHERE table_id = '{table_name}'
                 self.client.query(sql)
 
             else:
-                config = LoadJobConfig(**configuration_params)
-                if overwrite_types is not None:
-                    config.schema = [SchemaField(k, v) for k, v in overwrite_types.items()]
-                if 'write_disposition' not in configuration_params:
-                    if if_exists == 'replace':
-                        config.write_disposition = WriteDisposition.WRITE_TRUNCATE
-                    elif if_exists == 'append':
-                        config.write_disposition = WriteDisposition.WRITE_APPEND
-                    elif if_exists == 'fail':
-                        config.write_disposition = WriteDisposition.WRITE_EMPTY
-                    else:
-                        raise ValueError(
-                            f'Invalid policy specified for handling existence of '
-                            f'table: \'{if_exists}\''
+                if (
+                    if_exists == ExportWritePolicy.APPEND
+                    and not table_doesnt_exist
+                    and unique_constraints
+                    and unique_conflict_method
+                ):
+                    temp_table_id = f'{table_id}_{uuid.uuid4().hex}'
+
+                    try:
+                        self.__write_table(
+                            df,
+                            temp_table_id,
+                            overwrite_types=overwrite_types,
+                            **configuration_params,
+                        )
+
+                        parts = temp_table_id.split('.')
+                        if len(parts) == 2:
+                            temp_table_name = parts[1]
+                        elif len(parts) == 3:
+                            temp_table_name = parts[2]
+                        column_types = self.get_column_types(schema, temp_table_name)
+                        columns = list(column_types.keys())
+                        if not columns:
+                            columns = df.columns.str.replace(' ', '_')
+
+                        on_conditions = []
+                        for col in unique_constraints:
+                            on_conditions.append(
+                                f'((a.{col} IS NULL AND b.{col} IS NULL) OR a.{col} = b.{col})',
+                            )
+
+                        insert_columns = ', '.join([f'`{col}`' for col in columns])
+
+                        merge_commands = [
+                            f'MERGE INTO `{table_id}` AS a',
+                            f'USING (SELECT * FROM `{temp_table_id}`) AS b',
+                            f"ON {' AND '.join(on_conditions)}",
+                        ]
+
+                        if UNIQUE_CONFLICT_METHOD_UPDATE == unique_conflict_method:
+                            set_command = ', '.join(
+                                [f'a.`{col}` = b.`{col}`' for col in columns],
+                            )
+                            merge_commands.append(f'WHEN MATCHED THEN UPDATE SET {set_command}')
+
+                        merge_values = f"({', '.join([f'b.`{col}`' for col in columns])})"
+                        merge_commands.append(
+                            f'WHEN NOT MATCHED THEN INSERT ({insert_columns}) VALUES {merge_values}',  # noqa: E501
                         )
-                parts = table_id.split('.')
-                if len(parts) == 2:
-                    schema = parts[0]
-                    table_name = parts[1]
-                elif len(parts) == 3:
-                    schema = parts[1]
-                    table_name = parts[2]
-
-                self.client.create_dataset(dataset=schema, exists_ok=True)
-
-                column_types = self.get_column_types(schema, table_name)
-
-                if df is not None:
-                    df.fillna(value=np.NaN, inplace=True)
-                    for col in df.columns:
-                        col_type = column_types.get(col)
-                        if not col_type:
-                            continue
-
-                        null_rows = df[col].isnull()
-                        if col_type.startswith('ARRAY<STRUCT'):
-                            df.loc[null_rows, col] = df.loc[null_rows, col].apply(lambda x: [{}])
-                        elif col_type.startswith('ARRAY'):
-                            df.loc[null_rows, col] = df.loc[null_rows, col].apply(lambda x: [])
-                        elif col_type.startswith('STRUCT'):
-                            df.loc[null_rows, col] = df.loc[null_rows, col].apply(lambda x: {})
-
-                    # Clean column names
-                    if type(df) is DataFrame:
-                        df.columns = df.columns.str.replace(' ', '_')
-
-                    self.client.load_table_from_dataframe(df, table_id, job_config=config).result()
+
+                        merge_command = '\n'.join(merge_commands)
+
+                        self.client.query(merge_command).result()
+                    finally:
+                        self.client.query(f'DROP TABLE IF EXISTS {temp_table_id}').result()
+                else:
+                    if not write_disposition:
+                        if if_exists == ExportWritePolicy.APPEND:
+                            write_disposition = WriteDisposition.WRITE_APPEND
+                        elif if_exists == ExportWritePolicy.REPLACE:
+                            write_disposition = WriteDisposition.WRITE_TRUNCATE
+                        elif if_exists == ExportWritePolicy.FAIL:
+                            write_disposition = WriteDisposition.WRITE_EMPTY
+                    self.__write_table(
+                        df,
+                        table_id,
+                        overwrite_types=overwrite_types,
+                        write_disposition=write_disposition,
+                        **configuration_params,
+                    )
 
         if verbose:
             with self.printer.print_msg(f'Exporting data to table \'{table_id}\''):
-                __process(database=database)
+                __process(database=database, write_disposition=write_disposition)
         else:
             __process(database=database)
 
+    def __write_table(
+        self,
+        df: DataFrame,
+        table_id: str,
+        overwrite_types: Dict = None,
+        **configuration_params,
+    ):
+        config = LoadJobConfig(**configuration_params)
+        if overwrite_types is not None:
+            config.schema = [SchemaField(k, v) for k, v in overwrite_types.items()]
+        if not config.write_disposition:
+            config.write_disposition = WriteDisposition.WRITE_APPEND
+        parts = table_id.split('.')
+        if len(parts) == 2:
+            schema = parts[0]
+            table_name = parts[1]
+        elif len(parts) == 3:
+            schema = parts[1]
+            table_name = parts[2]
+
+        self.client.create_dataset(dataset=schema, exists_ok=True)
+
+        column_types = self.get_column_types(schema, table_name)
+
+        if df is not None:
+            df.fillna(value=np.NaN, inplace=True)
+            for col in df.columns:
+                col_type = column_types.get(col)
+                if not col_type:
+                    continue
+
+                null_rows = df[col].isnull()
+                if col_type.startswith('ARRAY<STRUCT'):
+                    df.loc[null_rows, col] = df.loc[null_rows, col].apply(lambda x: [{}])
+                elif col_type.startswith('ARRAY'):
+                    df.loc[null_rows, col] = df.loc[null_rows, col].apply(lambda x: [])
+                elif col_type.startswith('STRUCT'):
+                    df.loc[null_rows, col] = df.loc[null_rows, col].apply(lambda x: {})
+
+            # Clean column names
+            if type(df) is DataFrame:
+                df.columns = df.columns.str.replace(' ', '_')
+
+            return self.client.load_table_from_dataframe(df, table_id, job_config=config).result()
+
     def execute(self, query_string: str, **kwargs) -> None:
         """
         Sends query to the connected BigQuery warehouse.

mage_ai/io/mysql.py

@@ -8,6 +8,7 @@ from mysql.connector.cursor import MySQLCursor
 from pandas import DataFrame, Series
 
 from mage_ai.io.config import BaseConfigLoader, ConfigKey
+from mage_ai.io.constants import UNIQUE_CONFLICT_METHOD_UPDATE
 from mage_ai.io.export_utils import PandasTypes
 from mage_ai.io.sql import BaseSQL
 from mage_ai.shared.parsers import encode_complex
@@ -60,7 +61,7 @@ class MySQL(BaseSQL):
     ) -> str:
         if unique_constraints is None:
             unique_constraints = []
-        query = []
+        columns_and_types = []
         for cname in dtypes:
             if overwrite_types is not None and cname in overwrite_types.keys():
                 dtypes[cname] = overwrite_types[cname]
@@ -68,9 +69,22 @@ class MySQL(BaseSQL):
                 cleaned_col_name = clean_name(cname, case_sensitive=case_sensitive)
             else:
                 cleaned_col_name = cname
-            query.append(f'`{cleaned_col_name}` {dtypes[cname]} NULL')
-
-        return f'CREATE TABLE {table_name} (' + ','.join(query) + ');'
+            columns_and_types.append(f'`{cleaned_col_name}` {dtypes[cname]} NULL')
+
+        if unique_constraints:
+            unique_constraints = [
+                clean_name(col, case_sensitive=case_sensitive)
+                for col in unique_constraints
+            ]
+            index_name = '_'.join([
+                clean_name(table_name, case_sensitive=case_sensitive),
+            ] + unique_constraints)
+            index_name = f'unique{index_name}'[:64]
+            columns_and_types.append(
+                f"CONSTRAINT {index_name} Unique({', '.join(unique_constraints)})"
+            )
+
+        return f'CREATE TABLE {table_name} (' + ','.join(columns_and_types) + ');'
 
     def open(self) -> None:
         with self.printer.print_msg('Opening connection to MySQL database'):
@@ -94,6 +108,9 @@ class MySQL(BaseSQL):
         dtypes: List[str],
         full_table_name: str,
         buffer: Union[IO, None] = None,
+        case_sensitive: bool = False,
+        unique_constraints: List[str] = None,
+        unique_conflict_method: str = None,
         **kwargs,
     ) -> None:
         def serialize_obj(val):
@@ -133,9 +150,24 @@ class MySQL(BaseSQL):
         for _, row in df_.iterrows():
             values.append(tuple([str(val) if type(val) is pd.Timestamp else val for val in row]))
 
-        insert_columns = ', '.join([f'`{col}`'for col in columns])
+        cleaned_columns = [clean_name(col, case_sensitive=case_sensitive) for col in columns]
+        insert_columns = ', '.join([f'`{col}`'for col in cleaned_columns])
+
+        query = [
+            f'INSERT INTO {full_table_name} ({insert_columns})',
+            f'VALUES ({values_placeholder})',
+        ]
+
+        if unique_constraints and unique_conflict_method:
+            if UNIQUE_CONFLICT_METHOD_UPDATE == unique_conflict_method:
+                update_command = [f'{col} = new.{col}' for col in cleaned_columns]
+                query += [
+                    'AS new',
+                    f"ON DUPLICATE KEY UPDATE {', '.join(update_command)}",
+                ]
+
+        sql = '\n'.join(query)
 
-        sql = f'INSERT INTO {full_table_name} ({insert_columns}) VALUES ({values_placeholder})'
         cursor.executemany(sql, values)
 
     def get_type(self, column: Series, dtype: str) -> str:

mage_ai/io/snowflake.py

@@ -11,6 +11,7 @@ from mage_ai.data_preparation.models.block.sql.utils.shared import (
 )
 from mage_ai.io.base import QUERY_ROW_LIMIT, BaseSQLConnection, ExportWritePolicy
 from mage_ai.io.config import BaseConfigLoader, ConfigKey
+from mage_ai.io.constants import UNIQUE_CONFLICT_METHOD_UPDATE
 from mage_ai.shared.hash import merge_dict
 
 DEFAULT_LOGIN_TIMEOUT = 20
@@ -201,6 +202,8 @@ class Snowflake(BaseSQLConnection):
         schema: str = None,
         if_exists: str = 'append',
         query_string: Union[str, None] = None,
+        unique_conflict_method: str = None,
+        unique_constraints: List[str] = None,
         verbose: bool = True,
         **kwargs,
     ) -> None:
@@ -265,37 +268,28 @@ class Snowflake(BaseSQLConnection):
                         cur.execute(f'DROP TABLE "{schema}"."{table_name}"', timeout=self.timeout)
                         should_create_table = True
 
-                if query_string:
-                    cur.execute(f'USE DATABASE {database}', timeout=self.timeout)
-                    cur.execute(f'USE SCHEMA {schema}', timeout=self.timeout)
-
-                    if should_create_table:
-                        cur.execute(f"""
-CREATE TABLE IF NOT EXISTS "{database}"."{schema}"."{table_name}" AS
-{query_string}
-""", timeout=self.timeout)
-                    else:
-                        cur.execute(f"""
-INSERT INTO "{database}"."{schema}"."{table_name}"
-{query_string}
-""", timeout=self.timeout)
-
-                else:
-                    write_kwargs = merge_dict(
-                        dict(
-                            auto_create_table=should_create_table,
-                            database=database,
-                            schema=schema,
-                            # This param makes sure datetime column is written correctly
-                            use_logical_type=True,
-                        ),
-                        kwargs or dict(),
-                    )
-                    write_pandas(
-                        self.conn,
+                if unique_constraints and unique_conflict_method and df is not None:
+                    self.__upsert_df_into_table(
+                        table_name,
                         df,
+                        cur,
+                        database,
+                        schema,
+                        should_create_table=should_create_table,
+                        unique_conflict_method=unique_conflict_method,
+                        unique_constraints=unique_constraints,
+                        **kwargs,
+                    )
+                else:
+                    self.__write_table(
                         table_name,
-                        **write_kwargs,
+                        df,
+                        cur,
+                        database,
+                        schema,
+                        should_create_table=should_create_table,
+                        query_string=query_string,
+                        **kwargs,
                     )
 
         if verbose:
@@ -306,6 +300,135 @@ INSERT INTO "{database}"."{schema}"."{table_name}"
         else:
             __process()
 
+    def __upsert_df_into_table(
+        self,
+        table_name: str,
+        df: DataFrame,
+        cursor,
+        database: str,
+        schema: str,
+        should_create_table: bool = False,
+        unique_conflict_method: str = None,
+        unique_constraints: List[str] = None,
+        allow_reserved_words: bool = True,
+        auto_clean_name: bool = True,
+        case_sensitive: bool = True,
+        **kwargs
+    ):
+        write_kwargs = merge_dict(
+            dict(
+                auto_create_table=True,
+                database=database,
+                schema=schema,
+                # This param makes sure datetime column is written correctly
+                use_logical_type=True,
+            ),
+            kwargs or dict(),
+        )
+        # should_create_table is True when the table does not exist, so just create the
+        # table as normal.
+        if should_create_table:
+            write_pandas(
+                self.conn,
+                df,
+                table_name,
+                **write_kwargs,
+            )
+        else:
+            temp_table_name = f'temp_{table_name}'
+            write_pandas(
+                self.conn,
+                df,
+                temp_table_name,
+                table_type='temp',
+                **write_kwargs,
+            )
+
+            cleaned_unique_constraints = []
+            for col in unique_constraints:
+                cleaned_col = self._clean_column_name(
+                    col,
+                    allow_reserved_words=allow_reserved_words,
+                    auto_clean_name=auto_clean_name,
+                    case_sensitive=case_sensitive,
+                )
+                cleaned_unique_constraints.append(f'"{cleaned_col}"')
+
+            cleaned_columns = []
+            for col in df.columns:
+                cleaned_col = self._clean_column_name(
+                    col,
+                    allow_reserved_words=allow_reserved_words,
+                    auto_clean_name=auto_clean_name,
+                    case_sensitive=case_sensitive,
+                )
+                cleaned_columns.append(f'"{cleaned_col}"')
+
+            merge_commands = [
+                f'MERGE INTO "{database}"."{schema}"."{table_name}" AS a',
+                f'USING (SELECT * FROM "{database}"."{schema}"."{temp_table_name}") AS b',
+                f"ON {' AND '.join([f'a.{col} = b.{col}' for col in cleaned_unique_constraints])}",
+            ]
+
+            if unique_conflict_method == UNIQUE_CONFLICT_METHOD_UPDATE:
+                set_command = ', '.join(
+                    [f'a.{col} = b.{col}' for col in cleaned_columns],
+                )
+                merge_commands.append(f'WHEN MATCHED THEN UPDATE SET {set_command}')
+
+            insert_columns = ', '.join(cleaned_columns)
+            merge_values = f"({', '.join([f'b.{col}' for col in cleaned_columns])})"
+            merge_commands.append(
+                f"WHEN NOT MATCHED THEN INSERT ({insert_columns}) VALUES {merge_values}",
+            )
+            merge_command = '\n'.join(merge_commands)
+
+            cursor.execute(merge_command, timeout=self.timeout)
+
+    def __write_table(
+        self,
+        table_name: str,
+        df: DataFrame,
+        cursor,
+        database: str,
+        schema: str,
+        should_create_table: bool = False,
+        query_string: str = None,
+        **kwargs
+    ):
+        if query_string:
+            cursor.execute(f'USE DATABASE {database}', timeout=self.timeout)
+            cursor.execute(f'USE SCHEMA {schema}', timeout=self.timeout)
+
+            if should_create_table:
+                cursor.execute(f"""
+CREATE TABLE IF NOT EXISTS "{database}"."{schema}"."{table_name}" AS
+{query_string}
+""", timeout=self.timeout)
+            else:
+                cursor.execute(f"""
+INSERT INTO "{database}"."{schema}"."{table_name}"
+{query_string}
+""", timeout=self.timeout)
+
+        else:
+            write_kwargs = merge_dict(
+                dict(
+                    auto_create_table=should_create_table,
+                    database=database,
+                    schema=schema,
+                    # This param makes sure datetime column is written correctly
+                    use_logical_type=True,
+                ),
+                kwargs or dict(),
+            )
+            write_pandas(
+                self.conn,
+                df,
+                table_name,
+                **write_kwargs,
+            )
+
     @classmethod
     def with_config(
         cls,

mage_ai/orchestration/db/migrations/versions/42a14d6143f1_update_token_column_type.py

@@ -0,0 +1,54 @@
+"""Update token column type
+
+Revision ID: 42a14d6143f1
+Revises: b9a2d6d0a2c7
+Create Date: 2024-04-12 15:19:52.639580
+
+"""
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = '42a14d6143f1'
+down_revision = 'b9a2d6d0a2c7'
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table('oauth2_access_token', schema=None) as batch_op:
+        batch_op.alter_column(
+            'token',
+            existing_type=sa.String(length=255),
+            type_=sa.Text(),
+            existing_nullable=True,
+        )
+        batch_op.alter_column(
+            'refresh_token',
+            existing_type=sa.String(length=255),
+            type_=sa.Text(),
+            existing_nullable=True,
+        )
+
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table('oauth2_access_token', schema=None) as batch_op:
+        batch_op.alter_column(
+            'token',
+            existing_type=sa.Text(),
+            type_=sa.String(length=255),
+            existing_nullable=True,
+        )
+        batch_op.alter_column(
+            'refresh_token',
+            existing_type=sa.String(length=255),
+            type_=sa.Text(),
+            existing_nullable=True,
+        )
+
+    # ### end Alembic commands ###

mage_ai/orchestration/db/models/oauth.py

@@ -1,7 +1,7 @@
 import enum
 import re
 from datetime import datetime
-from typing import Dict, List, Union
+from typing import Callable, Dict, List, Union
 
 from sqlalchemy import (
     JSON,
@@ -12,6 +12,7 @@ from sqlalchemy import (
     ForeignKey,
     Integer,
     String,
+    Text,
     and_,
     asc,
     func,
@@ -324,10 +325,10 @@ class Role(BaseModel):
     @classmethod
     @safe_db_query
     def create_default_roles(
-        self,
+        cls,
         entity: Entity = None,
         entity_id: str = None,
-        prefix: str = None,
+        name_func: Callable[[str], str] = None,
     ) -> None:
         """
         Create default roles with associated permissions for a given entity and entity_id.
@@ -342,18 +343,18 @@ class Role(BaseModel):
             entity = Entity.GLOBAL
         permissions = Permission.create_default_permissions(entity=entity, entity_id=entity_id)
         mapping = {
-            self.DefaultRole.OWNER: Permission.Access.OWNER,
-            self.DefaultRole.ADMIN: Permission.Access.ADMIN,
-            self.DefaultRole.EDITOR: Permission.Access.EDITOR,
-            self.DefaultRole.VIEWER: Permission.Access.VIEWER,
+            cls.DefaultRole.OWNER: Permission.Access.OWNER,
+            cls.DefaultRole.ADMIN: Permission.Access.ADMIN,
+            cls.DefaultRole.EDITOR: Permission.Access.EDITOR,
+            cls.DefaultRole.VIEWER: Permission.Access.VIEWER,
         }
         for name, access in mapping.items():
             role_name = name
-            if prefix is not None:
-                role_name = f'{prefix}_{name}'
-            role = self.query.filter(self.name == role_name).first()
+            if name_func is not None:
+                role_name = name_func(name)
+            role = cls.query.filter(Role.name == role_name).first()
             if not role:
-                self.create(
+                cls.create(
                     name=role_name,
                     permissions=[
                         Permission.query.filter(
@@ -880,10 +881,10 @@ class Oauth2AccessToken(BaseModel):
     expires = Column(DateTime(timezone=True))
     oauth2_application = relationship(Oauth2Application, back_populates='oauth2_access_tokens')
     oauth2_application_id = Column(Integer, ForeignKey('oauth2_application.id'))
-    token = Column(String(255), index=True, unique=True)
+    token = Column(Text, index=True, unique=True)
     user = relationship(User, back_populates='oauth2_access_tokens')
     user_id = Column(Integer, ForeignKey('user.id'))
-    refresh_token = Column(String(255))
+    refresh_token = Column(Text)
 
     def is_valid(self) -> bool:
         return self.token and \