<commit_before>#!/usr/bin/env python3

# Author: <NAME> (gabfl)
# Github: https://github.com/gabfl/mysql-batch-update
# Version: 1.0.2
# Compatible with python 2.7 & 3

import sys, time
import pymysql.cursors, argparse

# Parse arguments
parser = argparse.ArgumentParser()
parser.add_argument("-H", "--host", default="127.0.0.1",
                    help="MySQL server host")
parser.add_argument("-P", "--port", type=int, default=3306,
                    help="MySQL server port")
parser.add_argument("-U", "--user", required = True,
                    help="MySQL user")
parser.add_argument("-p", "--password", default='',
                    help="MySQL password")
parser.add_argument("-d", "--database", required = True,
                    help="MySQL database name")
parser.add_argument("-t", "--table", required = True,
                    help="MySQL table")
parser.add_argument("-id", "--primary_key", default='id',
                    help="Name of the primary key column")
parser.add_argument("-w", "--where", required = True,
                    help="Select WHERE clause")
parser.add_argument("-s", "--set",
                    help="Update SET clause")
parser.add_argument("-rbz", "--read_batch_size", type=int, default = 1000,
                    help="Select batch size")
parser.add_argument("-wbz", "--write_batch_size", type=int, default = 50,
                    help="Update/delete batch size")
parser.add_argument("-S", "--sleep", type=float, default = 0.00,
                    help="Sleep after each batch")
parser.add_argument("-a", "--action", default = 'update', choices=['update', 'delete'],
                    help="Action ('update' or 'delete')")
args = parser.parse_args()

# Make sure we have a SET clause for updates
if args.action == 'update' and args.set is None:
    print ("Error: argument -s/--set is required for updates.");
    sys.exit();

def updateBatch(ids):
    global confirmedWrite

    """Update a batch of rows"""

    # Leave if ids is empty
    if not ids or len(ids) == 0:
        return None;

    # Prepare update
    print('* Updating %i rows...' % len(ids))
    sql = "UPDATE " + args.table + " SET " + args.set + " WHERE {0} IN (".format(args.primary_key) + ', '.join([str(x) for x in ids]) + ")"
    print ("   query: " + sql)

    if confirmedWrite or query_yes_no("* Start updating?"):
        # Switch confirmedWrite skip the question for the next update
        confirmedWrite = True

        # Execute query
        runQuery(sql)
    else: # answered "no"
        print ("Error: Update declined.");
        sys.exit();

def deleteBatch(ids):
    global confirmedWrite

    """Delete a batch of rows"""

    # Leave if ids is empty
    if not ids or len(ids) == 0:
        return None;

    # Prepare delete
    print('* Deleting %i rows...' % len(ids))
    sql = "DELETE FROM " + args.table + " WHERE {0} IN (".format(args.primary_key) + ', '.join([str(x) for x in ids]) + ")"
    print ("   query: " + sql)

    if confirmedWrite or query_yes_no("* Start deleting?"):
        # Switch confirmedWrite skip the question for the next delete
        confirmedWrite = True

        # Execute query
        runQuery(sql)
    else: # answered "no"
        print ("Error: Delete declined.");
        sys.exit();

def runQuery(sql):
    """Execute a write query"""

    # Execute query
    with connection.cursor() as cursorUpd:
        cursorUpd.execute(sql)
        connection.commit()

    # Optional Sleep
    if args.sleep > 0:
        time.sleep(args.sleep)

def query_yes_no(question, default="yes"):
    """Ask a yes/no question via raw_input() and return their answer.

    "question" is a string that is presented to the user.
    "default" is the presumed answer if the user just hits <Enter>.
        It must be "yes" (the default), "no" or None (meaning
        an answer is required of the user).

    The "answer" return value is True for "yes" or False for "no".

    (thanks https://code.activestate.com/recipes/577058/)
    """
    valid = {"yes": True, "y": True, "ye": True,
             "no": False, "n": False}
    if default is None:
        prompt = " [y/n] "
    elif default == "yes":
        prompt = " [Y/n] "
    elif default == "no":
        prompt = " [y/N] "
    else:
        raise ValueError("invalid default answer: '%s'" % default)

    while True:
        sys.stdout.write(question + prompt)

        # Get user choice with python 2.7 retro-compatibility
        if sys.version_info >= (3,0):
            # Python 3
            # print ("python >= 3");
            choice = input().lower()
        else:
            # Python 2.7 retro-compatibility
            # print ("python 2.7");
            choice = raw_input().lower()

        if default is not None and choice == '':
            return valid[default]
        elif choice in valid:
            return valid[choice]
        else:
            sys.stdout.write("Please respond with 'yes' or 'no' "
                             "(or 'y' or 'n').\n")

# Connect to the database
try:
    connection = pymysql.connect(host = args.host,
                                 user = args.user,
                                 port = args.port,
                                 password = args.password,
                                 db = args.database,
                                 charset='utf8mb4',
                                 cursorclass=pymysql.cursors.DictCursor)
except:
    print ("Error: MySQL connection failed.");
    sys.exit();

try:
    with connection.cursor() as cursor:
        minId = 0
        confirmedWrite = False
        while 1: # Infinite loop, will be broken by sys.exit()
            # Get rows to modify
            print("* Selecting data...")
            sql = "SELECT {0} as id FROM ".format(args.primary_key) + args.table + " WHERE " + args.where + " AND {0} > %s ORDER BY {1} LIMIT %s".format(args.primary_key, args.primary_key)
            print ("   query: " + sql % (minId, args.read_batch_size))
            cursor.execute(sql, (minId, args.read_batch_size))

            # Row count
            count = cursor.rowcount

            # No more rows
            if count == 0:<commit_msg>New flag to allow processing without confirmation (for use in crons)<commit_after>#!/usr/bin/env python3

# Author: <NAME> (gabfl)
# Github: https://github.com/gabfl/mysql-batch-update
# Version: 1.0.2
# Compatible with python 2.7 & 3

import sys, time
import pymysql.cursors, argparse

# Parse arguments
parser = argparse.ArgumentParser()
parser.add_argument("-H", "--host", default="127.0.0.1",
                    help="MySQL server host")
parser.add_argument("-P", "--port", type=int, default=3306,
                    help="MySQL server port")
parser.add_argument("-U", "--user", required = True,
                    help="MySQL user")
parser.add_argument("-p", "--password", default='',
                    help="MySQL password")
parser.add_argument("-d", "--database", required = True,
                    help="MySQL database name")
parser.add_argument("-t", "--table", required = True,
                    help="MySQL table")
parser.add_argument("-id", "--primary_key", default='id',
                    help="Name of the primary key column")
parser.add_argument("-w", "--where", required = True,
                    help="Select WHERE clause")
parser.add_argument("-s", "--set",
                    help="Update SET clause")
parser.add_argument("-rbz", "--read_batch_size", type=int, default = 10000,
                    help="Select batch size")
parser.add_argument("-wbz", "--write_batch_size", type=int, default = 50,
                    help="Update/delete batch size")
parser.add_argument("-S", "--sleep", type=float, default = 0.00,
                    help="Sleep after each batch")
parser.add_argument("-a", "--action", default = 'update', choices=['update', 'delete'],
                    help="Action ('update' or 'delete')")
parser.add_argument("-n", "--no_confirm", action='store_true',
                    help="Don't ask for confirmation before to run the write queries")
args = parser.parse_args()

# Make sure we have a SET clause for updates
if args.action == 'update' and args.set is None:
    print ("Error: argument -s/--set is required for updates.");
    sys.exit();

def updateBatch(ids):
    global confirmedWrite

    """Update a batch of rows"""

    # Leave if ids is empty
    if not ids or len(ids) == 0:
        return None;

    # Prepare update
    print('* Updating %i rows...' % len(ids))
    sql = "UPDATE " + args.table + " SET " + args.set + " WHERE {0} IN (".format(args.primary_key) + ', '.join([str(x) for x in ids]) + ")"
    print ("   query: " + sql)

    if confirmedWrite or query_yes_no("* Start updating?"):
        # Switch confirmedWrite skip the question for the next update
        confirmedWrite = True

        # Execute query
        runQuery(sql)
    else: # answered "no"
        print ("Error: Update declined.");
        sys.exit();

def deleteBatch(ids):
    global confirmedWrite

    """Delete a batch of rows"""

    # Leave if ids is empty
    if not ids or len(ids) == 0:
        return None;

    # Prepare delete
    print('* Deleting %i rows...' % len(ids))
    sql = "DELETE FROM " + args.table + " WHERE {0} IN (".format(args.primary_key) + ', '.join([str(x) for x in ids]) + ")"
    print ("   query: " + sql)

    if confirmedWrite or query_yes_no("* Start deleting?"):
        # Switch confirmedWrite skip the question for the next delete
        confirmedWrite = True

        # Execute query
        runQuery(sql)
    else: # answered "no"
        print ("Error: Delete declined.");
        sys.exit();

def runQuery(sql):
    """Execute a write query"""

    # Execute query
    with connection.cursor() as cursorUpd:
        cursorUpd.execute(sql)
        connection.commit()

    # Optional Sleep
    if args.sleep > 0:
        time.sleep(args.sleep)

def query_yes_no(question, default="yes"):
    """Ask a yes/no question via raw_input() and return their answer.

    "question" is a string that is presented to the user.
    "default" is the presumed answer if the user just hits <Enter>.
        It must be "yes" (the default), "no" or None (meaning
        an answer is required of the user).

    The "answer" return value is True for "yes" or False for "no".

    (thanks https://code.activestate.com/recipes/577058/)
    """
    valid = {"yes": True, "y": True, "ye": True,
             "no": False, "n": False}
    if default is None:
        prompt = " [y/n] "
    elif default == "yes":
        prompt = " [Y/n] "
    elif default == "no":
        prompt = " [y/N] "
    else:
        raise ValueError("invalid default answer: '%s'" % default)

    while True:
        sys.stdout.write(question + prompt)

        # Get user choice with python 2.7 retro-compatibility
        if sys.version_info >= (3,0):
            # Python 3
            # print ("python >= 3");
            choice = input().lower()
        else:
            # Python 2.7 retro-compatibility
            # print ("python 2.7");
            choice = raw_input().lower()

        if default is not None and choice == '':
            return valid[default]
        elif choice in valid:
            return valid[choice]
        else:
            sys.stdout.write("Please respond with 'yes' or 'no' "
                             "(or 'y' or 'n').\n")

# Connect to the database
try:
    connection = pymysql.connect(host = args.host,
                                 user = args.user,
                                 port = args.port,
                                 password = <PASSWORD>,
                                 db = args.database,
                                 charset='utf8mb4',
                                 cursorclass=pymysql.cursors.DictCursor)
except:
    print ("Error: MySQL connection failed.");
    sys.exit();

try:
    # confirmedWrite default value
    confirmedWrite = False
    if args.no_confirm:
        confirmedWrite = True

    with connection.cursor() as cursor:
        # Default vars
        minId = 0

        while 1: # Infinite loop, will be broken by sys.exit()
            # Get rows to modify
            print("* Selecting data...")
            sql = "SELECT {0} as id FROM ".format(args.primary_key) + args.table + " WHERE " + args.where + " AND {0} > %s ORDER BY {1} LIMIT %s".format(args.primary_key, args.primary_key)
            print ("   query: " + sql % (minId, args.read_batch_size))
            cursor.execute(sql, (minId, args.read_batch_size))

            # Row count
            count = cursor.rowcount

            # No more rows
            if count == 0:
{'fraction_non_alphanumeric': 0.13652524294856602, 'fraction_numerical': 0.008769850675515525, 'mean_word_length': 2.067862336403296, 'pattern_counts': {'":': 14, '<': 8, '<?xml version=': 0, '>': 16, 'https://': 4, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 2, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 12, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '453891', 'n_tokens_mistral': 3696, 'n_tokens_neox': 3570, 'n_words': 1137}