save_changed_data.py

##########################################################################
#
# pgAdmin 4 - PostgreSQL Tools
#
# Copyright (C) 2013 - 2026, The pgAdmin Development Team
# This software is released under the PostgreSQL Licence
#
##########################################################################

from flask import render_template
from collections import OrderedDict

from pgadmin.tools.sqleditor.utils.constant_definition import TX_STATUS_IDLE
from pgadmin.utils.exception import ExecuteError

ignore_type_cast_list = ['character', 'character[]', 'bit', 'bit[]']


def save_changed_data(changed_data, columns_info, conn, command_obj,
                      client_primary_key, auto_commit=True):
    """
    This function is used to save the data into the database.
    Depending on condition it will either update or insert the
    new row into the database.

    Args:
        changed_data: Contains data to be saved
        command_obj: The transaction object (command_obj or trans_obj)
        conn: The connection object
        columns_info: session_obj['columns_info']
        client_primary_key: session_obj['client_primary_key']
        auto_commit: If the changes should be committed automatically.
    """
    status = False
    res = None
    query_results = []
    operations = ('added', 'updated', 'deleted')
    list_of_sql = {}
    _rowid = None

    pgadmin_alias = {
        col_name: col_info['pgadmin_alias']
        for col_name, col_info in columns_info.items()
    }

    is_savepoint = False
    # Start the transaction if the session is idle
    if conn.transaction_status() == TX_STATUS_IDLE:
        sql = 'BEGIN;'
    else:
        sql = 'SAVEPOINT save_data;'
        is_savepoint = True

    status, res = execute_void_wrapper(conn, sql, query_results)
    if not status:
        return status, res, query_results, None

    # Iterate total number of records to be updated/inserted
    for of_type in changed_data:
        # No need to go further if its not add/update/delete operation
        if of_type not in operations:
            continue
        # if no data to be save then continue
        if len(changed_data[of_type]) < 1:
            continue

        column_type = {}
        column_data = {}
        type_cast_required = {}
        for each_col in columns_info:
            if (
                columns_info[each_col]['not_null'] and
                not columns_info[each_col]['has_default_val']
            ):
                column_data[each_col] = None
                column_type[each_col] = \
                    columns_info[each_col]['type_name']
            else:
                column_type[each_col] = \
                    columns_info[each_col]['type_name']

            type_cast_required[each_col] = \
                True if column_type[each_col] not in ignore_type_cast_list \
                else False

        # For newly added rows
        if of_type == 'added':
            # Python dict does not honour the inserted item order
            # So to insert data in the order, we need to make ordered
            # list of added index We don't need this mechanism in
            # updated/deleted rows as it does not matter in
            # those operations
            added_index = OrderedDict(
                sorted(
                    changed_data['added_index'].items(),
                    key=lambda x: int(x[0])
                )
            )
            list_of_sql[of_type] = []

            # When new rows are added, only changed columns data is
            # sent from client side. But if column is not_null and has
            # no_default_value, set column to blank, instead
            # of not null which is set by default.
            column_data = {}
            pk_names, primary_keys = command_obj.get_primary_keys()

            for each_row in added_index:
                # Get the row index to match with the added rows
                # dict key
                tmp_row_index = added_index[each_row]
                data = changed_data[of_type][tmp_row_index]['data']
                # Remove our unique tracking key
                data.pop(client_primary_key, None)
                data.pop('is_row_copied', None)

                # Remove oid col
                if command_obj.has_oids():
                    data.pop('oid', None)

                # Remove generated columns (GENERATED ALWAYS AS) as they
                # cannot be inserted - PostgreSQL auto-computes their values.
                for col_name, col_info in columns_info.items():
                    if col_info.get('is_generated', False):
                        data.pop(col_name, None)

                # Update columns value with columns having
                # not_null=False and has no default value
                column_data.update(data)
                use_default = False
                if not column_data:
                    for each_col in columns_info:
                        if columns_info[each_col]['has_default_val']:
                            column_data[each_col] = 'set_default'
                            use_default = True

                sql = render_template(
                    "/".join([command_obj.sql_path, 'insert.sql']),
                    data_to_be_saved=column_data,
                    pgadmin_alias=pgadmin_alias,
                    primary_keys=None,
                    object_name=command_obj.object_name,
                    nsp_name=command_obj.nsp_name,
                    data_type=column_type,
                    pk_names=pk_names,
                    has_oids=command_obj.has_oids(),
                    type_cast_required=type_cast_required,
                    use_default=use_default
                )

                select_sql = render_template(
                    "/".join([command_obj.sql_path, 'select.sql']),
                    object_name=command_obj.object_name,
                    nsp_name=command_obj.nsp_name,
                    pgadmin_alias=pgadmin_alias,
                    primary_keys=primary_keys,
                    has_oids=command_obj.has_oids()
                )

                list_of_sql[of_type].append({
                    'sql': sql, 'data': data,
                    'client_row': tmp_row_index,
                    'select_sql': select_sql,
                    'row_id': data.get(client_primary_key)
                })
                # Reset column data
                column_data = {}

        # For updated rows
        elif of_type == 'updated':
            list_of_sql[of_type] = []

            # Check if table has generated columns. If yes, we need to
            # refetch row after UPDATE to get recalculated values for UI.
            has_generated_cols = any(
                col_info.get('is_generated', False)
                for col_info in columns_info.values()
            )

            # Get primary keys info (same as INSERT) - needed for RETURNING
            # clause and SELECT query to refetch updated row.
            pk_names, primary_keys = command_obj.get_primary_keys()

            for each_row in changed_data[of_type]:
                data = changed_data[of_type][each_row]['data']
                row_primary_keys = changed_data[of_type][each_row][
                    'primary_keys']

                # Remove generated columns (GENERATED ALWAYS AS) as they
                # cannot be updated - PostgreSQL auto-computes their values.
                for col_name, col_info in columns_info.items():
                    if col_info.get('is_generated', False):
                        data.pop(col_name, None)

                pk_escaped = {
                    pk: pk_val.replace('%', '%%') if hasattr(
                        pk_val, 'replace') else pk_val
                    for pk, pk_val in row_primary_keys.items()
                }

                # Pass pk_names and has_oids for RETURNING clause in
                # UPDATE statement.
                # This will help to fetch the updated row's.
                sql = render_template(
                    "/".join([command_obj.sql_path, 'update.sql']),
                    data_to_be_saved=data,
                    pgadmin_alias=pgadmin_alias,
                    primary_keys=pk_escaped,
                    object_name=command_obj.object_name,
                    nsp_name=command_obj.nsp_name,
                    data_type=column_type,
                    type_cast_required=type_cast_required,
                    pk_names=pk_names if has_generated_cols else None,
                    has_oids=command_obj.has_oids(),
                    conn=conn
                )

                # For tables with generated columns, add select_sql to
                # refetch updated row.
                if has_generated_cols:
                    select_sql = render_template(
                        "/".join([command_obj.sql_path, 'select.sql']),
                        object_name=command_obj.object_name,
                        nsp_name=command_obj.nsp_name,
                        pgadmin_alias=pgadmin_alias,
                        primary_keys=primary_keys,
                        has_oids=command_obj.has_oids()
                    )
                    list_of_sql[of_type].append({
                        'sql': sql,
                        'data': data,
                        'client_row': each_row,
                        'select_sql': select_sql,
                        'row_id': data.get(client_primary_key)
                    })
                else:
                    list_of_sql[of_type].append({
                        'sql': sql,
                        'data': data,
                        'row_id': data.get(client_primary_key)
                    })

        # For deleted rows
        elif of_type == 'deleted':
            delete_all = changed_data.get('delete_all', False)
            list_of_sql[of_type] = []
            _, primary_keys = command_obj.get_primary_keys()
            keys = list(primary_keys.keys())
            no_of_keys = len(keys)
            if delete_all:
                # for delete all, we'll need all the row primary key
                # values.
                status, result = \
                    conn.async_fetchmany_2darray(records=-1)

                if not status:
                    raise ExecuteError(result)

                columns_info = conn.get_column_info()

                rows_to_delete = [
                    dict([
                        (col['name'], r[col['pos']])
                        for col in filter(lambda c: c['name'] in keys,
                                          changed_data['columns'])
                    ])
                    for r in result
                ]
            else:
                rows_to_delete = [
                    r for _, r in changed_data[of_type].items()
                ]

            sql = render_template(
                "/".join([command_obj.sql_path, 'delete.sql']),
                data=rows_to_delete,
                primary_key_labels=keys,
                no_of_keys=no_of_keys,
                object_name=command_obj.object_name,
                nsp_name=command_obj.nsp_name,
                conn=conn
            )
            list_of_sql[of_type].append({'sql': sql, 'data': {}})

    def failure_handle(res, row_id):
        mogrified_sql = conn.mogrify(item['sql'], item['data'])
        mogrified_sql = mogrified_sql if mogrified_sql is not None \
            else item['sql']
        query_results.append({
            'status': False,
            'result': res,
            'sql': mogrified_sql,
            'rows_affected': 0,
            'row_added': None
        })

        if is_savepoint:
            sql = 'ROLLBACK TO SAVEPOINT save_data;'
            msg = 'A ROLLBACK was done for the save operation only. ' \
                  'The active transaction is not affected.'
        else:
            sql = 'ROLLBACK;'
            msg = 'A ROLLBACK was done for the save transaction.'

        rollback_status, rollback_result = \
            execute_void_wrapper(conn, sql, query_results)
        if not rollback_status:
            return rollback_status, rollback_result, query_results, None

        # If we roll backed every thing then update the
        # message for each sql query.
        for query in query_results:
            if query['status']:
                query['result'] = msg

        return False, res, query_results, row_id

    for opr, sqls in list_of_sql.items():
        for item in sqls:
            if item['sql']:
                item['data'] = {
                    pgadmin_alias[k] if k in pgadmin_alias else k: v
                    for k, v in item['data'].items()
                }

                row_added = None

                try:
                    # Fetch oids/primary keys
                    if 'select_sql' in item and item['select_sql']:
                        status, res = conn.execute_dict(
                            item['sql'], item['data'])
                    else:
                        status, res = conn.execute_void(
                            item['sql'], item['data'])
                except Exception:
                    failure_handle(res, item.get('row_id', 0))
                    raise

                if not status:
                    return failure_handle(res, item.get('row_id', 0))

                # Select added/updated row from the table
                if 'select_sql' in item:
                    params = {
                        pgadmin_alias[k] if k in pgadmin_alias else k: v
                        for k, v in res['rows'][0].items()
                    }
                    status, sel_res = conn.execute_dict(
                        item['select_sql'], params)

                    if not status:
                        return failure_handle(sel_res, item.get('row_id', 0))

                    if 'rows' in sel_res and len(sel_res['rows']) > 0:
                        row_added = {
                            item['client_row']: sel_res['rows'][0]}

                rows_affected = conn.rows_affected()
                mogrified_sql = conn.mogrify(item['sql'], item['data'])
                mogrified_sql = mogrified_sql if mogrified_sql is not None \
                    else item['sql']
                # store the result of each query in dictionary
                query_results.append({
                    'status': status,
                    'result': None if row_added else res,
                    'sql': mogrified_sql,
                    'rows_affected': rows_affected,
                    'row_added': row_added
                })

    # Commit the transaction if no error is found & autocommit is activated
    if auto_commit:
        sql = 'COMMIT;'
        status, res = execute_void_wrapper(conn, sql, query_results)
        if not status:
            return status, res, query_results, None

    return status, res, query_results, _rowid


def execute_void_wrapper(conn, sql, query_results):
    """
    Executes a sql query with no return and adds it to query_results
    :param sql: Sql query
    :param query_results: A list of query results in the save operation
    :return: status, result
    """
    status, res = conn.execute_void(sql)
    if status:
        query_results.append({
            'status': status,
            'result': res,
            'sql': sql, 'rows_affected': 0,
            'row_added': None
        })
    return status, res