mage-ai 0.9.67__py3-none-any.whl → 0.9.69__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/export_utils.py

@@ -108,6 +108,7 @@ def gen_table_creation_query(
     case_sensitive: bool = False,
     unique_constraints: List[str] = None,
     overwrite_types: Dict = None,
+    skip_semicolon_at_end: bool = False,
 ) -> str:
     """
     Generates a database table creation query from a data frame.
@@ -157,4 +158,6 @@ def gen_table_creation_query(
         query.append(
             f"CONSTRAINT {index_name} UNIQUE ({', '.join(unique_constraints_escaped)})",
         )
+    if skip_semicolon_at_end:
+        return f'CREATE TABLE {full_table_name} (' + ','.join(query) + ')'
     return f'CREATE TABLE {full_table_name} (' + ','.join(query) + ');'

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/oracledb.py

@@ -1,13 +1,17 @@
 import warnings
-from typing import Union
+from typing import IO, Any, List, Union
 
+import numpy as np
 import oracledb
-from pandas import DataFrame, read_sql
+import simplejson
+from pandas import DataFrame, Series, read_sql
 
-from mage_ai.io.base import QUERY_ROW_LIMIT
+from mage_ai.io.base import QUERY_ROW_LIMIT, ExportWritePolicy
 from mage_ai.io.config import BaseConfigLoader, ConfigKey
+from mage_ai.io.export_utils import PandasTypes
 from mage_ai.io.sql import BaseSQL
 from mage_ai.server.logger import Logger
+from mage_ai.shared.parsers import encode_complex
 
 logger = Logger().new_server_logger(__name__)
 
@@ -122,3 +126,134 @@ SELECT *
 FROM subquery
 FETCH FIRST {limit} ROWS ONLY
                 """
+
+    def table_exists(self, schema_name: str, table_name: str) -> bool:
+        with self.conn.cursor() as cur:
+            try:
+                cur.execute(f"select * from {table_name} where rownum=1")
+            except Exception as exc:
+                logger.info(f"Table not existing: {table_name}. Exception: {exc}")
+                return False
+        return True
+
+    def get_type(self, column: Series, dtype: str) -> str:
+        if dtype in (
+            PandasTypes.MIXED,
+            PandasTypes.UNKNOWN_ARRAY,
+            PandasTypes.COMPLEX,
+        ):
+            return 'CHAR(255)'
+        elif dtype in (PandasTypes.DATETIME, PandasTypes.DATETIME64):
+            try:
+                if column.dt.tz:
+                    return 'TIMESTAMP'
+            except AttributeError:
+                pass
+            return 'TIMESTAMP'
+        elif dtype == PandasTypes.TIME:
+            try:
+                if column.dt.tz:
+                    return 'TIMESTAMP'
+            except AttributeError:
+                pass
+            return 'TIMESTAMP'
+        elif dtype == PandasTypes.DATE:
+            return 'DATE'
+        elif dtype == PandasTypes.STRING:
+            return 'CHAR(255)'
+        elif dtype == PandasTypes.CATEGORICAL:
+            return 'CHAR(255)'
+        elif dtype == PandasTypes.BYTES:
+            return 'CHAR(255)'
+        elif dtype in (PandasTypes.FLOATING, PandasTypes.DECIMAL, PandasTypes.MIXED_INTEGER_FLOAT):
+            return 'NUMBER'
+        elif dtype == PandasTypes.INTEGER:
+            max_int, min_int = column.max(), column.min()
+            if np.int16(max_int) == max_int and np.int16(min_int) == min_int:
+                return 'NUMBER'
+            elif np.int32(max_int) == max_int and np.int32(min_int) == min_int:
+                return 'NUMBER'
+            else:
+                return 'NUMBER'
+        elif dtype == PandasTypes.BOOLEAN:
+            return 'CHAR(52)'
+        elif dtype in (PandasTypes.TIMEDELTA, PandasTypes.TIMEDELTA64, PandasTypes.PERIOD):
+            return 'NUMBER'
+        elif dtype == PandasTypes.EMPTY:
+            return 'CHAR(255)'
+        else:
+            print(f'Invalid datatype provided: {dtype}')
+
+        return 'CHAR(255)'
+
+    def export(
+        self,
+        df: DataFrame,
+        table_name: str = None,
+        if_exists: ExportWritePolicy = ExportWritePolicy.REPLACE,
+        **kwargs,
+    ) -> None:
+        super().export(
+            df,
+            **kwargs,
+            table_name=table_name,
+            if_exists=if_exists,
+            # Oracle cursor execute will automatically add a semicolon at the end of the query.
+            skip_semicolon_at_end=True
+        )
+
+    def upload_dataframe(
+        self,
+        cursor: Any,
+        df: DataFrame,
+        db_dtypes: List[str],
+        dtypes: List[str],
+        full_table_name: str,
+        buffer: Union[IO, None] = None,
+        **kwargs,
+    ) -> None:
+        def serialize_obj(val):
+            if type(val) is dict or type(val) is np.ndarray:
+                return simplejson.dumps(
+                    val,
+                    default=encode_complex,
+                    ignore_nan=True,
+                )
+            elif type(val) is list and len(val) >= 1 and type(val[0]) is dict:
+                return simplejson.dumps(
+                    val,
+                    default=encode_complex,
+                    ignore_nan=True,
+                )
+            return val
+
+        # Create values
+        df_ = df.copy()
+        columns = df_.columns
+        for col in columns:
+            dtype = df_[col].dtype
+            if dtype == PandasTypes.OBJECT:
+                df_[col] = df_[col].apply(lambda x: serialize_obj(x))
+            elif dtype in (
+                PandasTypes.MIXED,
+                PandasTypes.UNKNOWN_ARRAY,
+                PandasTypes.COMPLEX,
+            ):
+                df_[col] = df_[col].astype('string')
+
+            # Remove extraneous surrounding double quotes
+            # that get added while performing conversion to string.
+            df_[col] = df_[col].apply(lambda x: x.strip('"') if x and isinstance(x, str) else x)
+        df_.replace({np.NaN: None}, inplace=True)
+        values = []
+        for i in range(0, len(df_)):
+            values.append(tuple(df_.fillna('').values[i]))
+
+        # Create values placeholder
+        colmn_names = df.columns.tolist()
+        values_placeholder = ""
+        for i in range(0, len(colmn_names)):
+            values_placeholder += f':{str(i + 1)},'
+
+        insert_sql = f'INSERT INTO {full_table_name} VALUES({values_placeholder.rstrip(",")})'
+        cursor.executemany(insert_sql, values)