QFQ

import mysql.connector

import subprocess

import sys

import os

import getpass

import argparse

exit_flag = False

finish_flag = False

def get_global_exit():

    global exit_flag

    return exit_flag

def get_global_finish():

    global finish_flag

    return finish_flag

def restart_script():

    restart_choice = ""

    global exit_flag

    restart_choice = input("Do you want to restart the script from the beginning? (y/N): ").strip().lower()

    if restart_choice == "y":

        print("Restarting the script...")

        subprocess.run(['python3', 'table_data_transfer.py'])  # Restart the script

    else:

        print("Exiting script.")

        exit_flag = True

        sys.exit()

def truncate_and_import(tables, host, username, password, database, directory):

    connection = mysql.connector.connect(host=host,

                                 user=username,

                                 password=password,

                                 database=database)

    cursor = connection.cursor()

    for table in tables:

        print('Truncating target table: {}'.format(table))

        cursor.execute('TRUNCATE TABLE {}'.format(table))

        filename = '{}.sql'.format(table)

        filepath = os.path.join(directory, filename)

        print('Importing data to target table: {}'.format(table))

        with open(filepath) as sql_file:

            sql_query = sql_file.read()

            subprocess.run(['mysql', '-h', host, '-u', username, '-p'+password, database], input=sql_query, universal_newlines=True)

    cursor.close()

    connection.commit()

    connection.close()

def dump(tables, host, username, password, database, directory):

    try:

        for table in tables:

            try:

                filename = '{}.sql'.format(table)

                filepath = os.path.join(directory, filename)

                print('Dumping data from table: {}'.format(table))

                with open(filepath, 'w') as f:

                    process = subprocess.Popen(['mysqldump', '-h', host, '-u', username, '-p'+password, database, table], stdout=f, stderr=subprocess.PIPE, universal_newlines=True)

                    _, stderr_output = process.communicate()

                if stderr_output and "Access denied for user" in stderr_output:

                    print(f"Error: Access denied. Check your MySQL credentials. Password or db name ({database}) is wrong.")

                    restart_script()

                elif stderr_output and f'Couldn\'t find table: "{table}"' in stderr_output:

                    print(f"Error: Table '{database}.{table}' doesn't exist.")

                    restart_script()

                else:

                    print("Source dump completed successfully.")

            except mysql.connector.Error as e:

                if e.errno == mysql.connector.errorcode.ER_NO_SUCH_TABLE:

                    print("Error: Table '{}' doesn't exist.".format(table))

                else:

                    print("An mysql error occurred during source data dump:", e)

                restart_script()

    except subprocess.CalledProcessError as e:

        print("An error occurred during source data dump:", e)

        restart_script()

def dump_target(tables, host, username, password, database, directory):

    try:

        for table in tables:

            try:

                print('Dumping target data from table: {}'.format(table))

                file_number = 0

                filename = 'target_{}.sql'.format(table)

                filepath = os.path.join(directory, filename)

                while os.path.exists(filepath):

                    file_number += 1

                    filename = 'target_{}_{}.sql'.format(table, file_number)

                    filepath = os.path.join(directory, filename)

                with open(filepath, 'w') as f:

                    process = subprocess.Popen(['mysqldump', '-h', host, '-u', username, '-p'+password, database, table], stdout=f, stderr=subprocess.PIPE, universal_newlines=True)

                    _, stderr_output = process.communicate()

                if stderr_output and "Access denied for user" in stderr_output:

                    print(f"Error: Access denied. Check your MySQL credentials. Password or db name ({database}) is wrong.")

                    restart_script()

                elif stderr_output and f'Couldn\'t find table: "{table}"' in stderr_output:

                    print(f"Error: Table '{database}.{table}' doesn't exist.")

                    restart_script()

                else:

                    print("Target dump completed successfully.")

            except mysql.connector.Error as e:

                if e.errno == mysql.connector.errorcode.ER_NO_SUCH_TABLE:

                    print("Error: Table '{}' doesn't exist.".format(table))

                else:

                    print("An error occurred during data dump:", e)

                restart_script()

    except subprocess.CalledProcessError as e:

        print("An error occurred during target data dump:", e)

        restart_script()

def get_table_names():

    table_names = input("Enter comma-separated table names for data transfer: ").strip()

    table_names = [name.strip() for name in table_names.split(",")]

    clean_table_names = ",".join(table_names)

    return table_names

if __name__ == "__main__":

    directory = '/var/www/transfered_tables'    

    # Create the directory if it doesn't exist

    if not os.path.exists(directory):

        os.makedirs(directory)

    print("SQL dump destination: " + directory)

    # Dump from server 1

    dump_host = 'wdb16.math.uzh.ch'

    dump_user = 'enured'

    dump_database = input("Source DB name: ")

    dump_password = getpass.getpass('Enter database password for ' + dump_host + '.' + dump_database + ', user: ' + dump_user)

    # Import to server 2

    import_host = 'wdb22.math.uzh.ch'

    import_user = 'import_user'

    import_database = input("Target DB name: ")

    import_password = getpass.getpass('Enter database password for ' + import_host + '.' + import_database + ', user: ' + import_user)

    tables = get_table_names() # list of table names

    finish_flag = get_global_finish()

    exit_flag = get_global_exit()

    if not exit_flag:

        print("----------------------\nStarting dump source\n----------------------")

        dump(tables, dump_host, dump_user, dump_password, dump_database, directory)

        if not exit_flag:

            print("----------------------\nStarting dump target\n----------------------")

            dump_target(tables, import_host, import_user, import_password, import_database, directory)

            if not exit_flag and not finish_flag:

                print("----------------------\nStart truncate and import\n-----------------------")

                truncate_and_import(tables, import_host, import_user, import_password, import_database, directory)

                exit_flag = True

    dump_database = ""

    import_database = ""