<commit_before>import filesegmenter
import hashlib
import os
import shutil
import swiftclient
import sys
import datetime
import socket
import time
import olrcdb
from multiprocessing import Process, Lock, Value


#Settings
AUTH_VERSION = 2
SWIFT_AUTH_URL = ''
USERNAME = ''
PASSWORD = ''
CONTAINER = ''
AUTH_TOKEN = ''
STORAGE_URL = ''
TEMP_DIRECTORY = 'temp'
FILE_LIMIT = 0.5*10**9  # Max file size in bytes that a file can be uploaded.
# Anything larger is segmented
SEGMENT_SIZE = 100*10**6
COUNT = 0
TOTAL = 0
FAILED_COUNT = 0
LIMIT = 1000
RANGE = 0  # protected variable
SLEEP = 100

REQUIRED_VARIABLES = [
    'OS_AUTH_URL',
    'OS_USERNAME',
    'OS_TENANT_NAME',
    'OS_PASSWORD',
    "MYSQL_HOST",
    "MYSQL_USER",
    "MYSQL_PASSWD",
    "MYSQL_DB"
]


def olrc_upload(path):
    ''' Given a path, upload it to the OLRC. Return False if upload
    unsuccessful.'''

    global FAILED_COUNT

    # Check file not already online.
    if not is_uploaded(path):

        if not (olrc_upload_file(path)):
            FAILED_COUNT += 1
            error_log = open('error.log', 'a')
            error_log.write("\rFailed: {0}\n".format(path))
            error_log.close()
            return False

    return True


def olrc_upload_segments(source_file, target_directory):
    ''' Break up the source_file into segments and upload them into
    target_directory'''

    segments = filesegmenter.split_file(
        source_file,
        'temp',
        SEGMENT_SIZE
    )

    sys.stdout.flush()
    sys.stdout.write("\rPartitioning file {0}".format(source_file))

    for segment in segments:

        # Files are within the temp directory locally. On the server the file
        # will live in the target_directory so we need to remove 'temp/' from
        # target_file.
        target_file = os.path.join(target_directory, segment.split('/', 1)[1])
        sys.stdout.flush()
        sys.stdout.write("\rUploading file {0}".format(segment))

        # MDCheck
        if not is_uploaded(segment, target_file):
            if not olrc_upload_file(segment, target_file):
                return False
        else:

            sys.stdout.flush()
            sys.stdout.write(
                "\rSkipping: {0}, already uploaded.".format(source_file)
            )

    # Create and upload readme file.

    readme = "The file in this directory has been segmented for convenient" \
        " upload and download. To assemble the file, run the following " \
        "command on your machine in the directory with all the segments: " \
        "\n\ncat * >> {0}".format(source_file.split('/')[-1])
    outFile = open("temp/readme.txt", "wt")
    outFile.write(readme)
    outFile.close()

    olrc_upload_file(
        "temp/readme.txt",
        os.path.join(target_directory, "readme.txt")
    )
    #Clean up temp files
    if os.path.isdir(TEMP_DIRECTORY):
        shutil.rmtree(TEMP_DIRECTORY)

    return True


def olrc_upload_file(path):
    '''Given String source_file, upload the file to the OLRC to target_file
     and return True if successful. '''

    global SLEEP
    try:
        opened_source_file = open(path, 'r')
    except IOError:
        print("Error opening: " + path)
        return False
    try:
        swiftclient.client.put_object(
            STORAGE_URL,
            AUTH_TOKEN,
            CONTAINER,
            path,
            opened_source_file)

    except Exception, e:

        olrc_connect()
        sys.stdout.flush()
        sys.stdout.write("\rError! {0}\n".format(e))
        sys.stdout.write(
            "Error! {0} Upload to OLRC failed."
            " Trying again in {1} seconds.\n".format(
                time.strftime("%Y-%m-%d %H:%M:%S", time.gmtime()),
                SLEEP
            )
        )
        time.sleep(SLEEP)
        SLEEP += SLEEP
        olrc_upload_file(path)

    return True


def olrc_connect():
    '''Connect to the OLRC with the global variables. Exit if connection
    fails.'''

    global SLEEP
    try:
        (connection_storage_url, auth_token) = swiftclient.client.get_auth(
            SWIFT_AUTH_URL, USERNAME, PASSWORD,
            auth_version=AUTH_VERSION)
        global AUTH_TOKEN
        AUTH_TOKEN = auth_token
        global STORAGE_URL
        STORAGE_URL = connection_storage_url
    except swiftclient.client.ClientException, e:
        print(e)
        sys.stdout.flush()
        sys.stdout.write(
            "\rError! {0} Connection to OLRC failed."
            " Trying again in {1} seconds.\n".format(
                time.strftime("%Y-%m-%d %H:%M:%S", time.gmtime()),
                SLEEP
            )
        )
        time.sleep(SLEEP)
        SLEEP += SLEEP

        olrc_connect()


def is_uploaded(file_name):
    '''Return True if String file is already on the server and its etag
    matches it's md5. Delete the file from the server if the
    md5 does not match.'''

    # Swift stat on filename.
    try:
        object_stat = swiftclient.client.head_object(
            STORAGE_URL,
            AUTH_TOKEN,
            CONTAINER,
            file_name
        )
        try:
            etag = object_stat['etag']
        except:
            # Return if no etag
            return False

        md5 = checksum_md5(file_name)
        match = etag == md5

        # Delete the file if the md5 does not match.
        if not match:
            try:
                swiftclient.client.delete_object(
                    STORAGE_URL,
                    AUTH_TOKEN,
                    CONTAINER,
                    file_name
                )
            except:
                pass

        return match
    except:
        return False


def checksum_md5(filename):
    md5 = hashlib.md5()
    with open(filename, 'rb') as f:
        for chunk in iter(lambda: f.read(8192), b''):
            md5.update(chunk)
    return md5.hexdigest()


def is_env_vars_set():
    '''Check all the required environment variables are set. Return false if
    any of them are undefined.'''

    global REQUIRED_VARIABLES
    for required_variable in REQUIRED_VARIABLES:
        if not os.environ.get(required_variable):
            return False

    return True


def set_env_vars():
    '''Set the global variables for swift client assuming they exist in the
    environment.'''

    global SWIFT_AUTH_URL
    SWIFT_AUTH_URL = os.environ.get("OS_AUTH_URL")

    global USERNAME
    USERNAME = os.environ.get("OS_TENANT_NAME") + \
        ":" + os.environ.get("OS_USERNAME")

    global PASSWORD
    PASSWORD = os.environ.get("OS_PASSWORD")

    return


def upload_table(lock, range, table_name, counter):
    '''
    Given a table_name, upload all the paths from the table where upload is 0.
    Upload within a LIMIT range at a time.
    '''
    global FAILED_COUNT, LIMIT, RANGE

    lock.acquire()
    while range.value <= TOTAL:

        # Access protected variable.
        # We want to make we only fetch within a unique range,
        # so RANGE is locked.
            # "grab" a LIMITs with of rows at a time.
        query = (
            "SELECT * FROM {0} WHERE uploaded=0"
            " AND id >= {1} AND id <{2}".format(
                table_name, range.value, range.value + LIMIT))
        # Let other processes know this range.value has been accounted for.
        range.value += LIMIT
        lock.release()

        # Connect to the database and fetch results.
        connect = olrcdb.DatabaseConnection()
        result = connect.execute_query(query)
        path_tuple = result.fetchone()
        # Loop until we run out of rows from the database.
        while (path_tuple):

            # if the upload is successful, update the database
            if olrc_upload(path_tuple[1]):
                lock.acquire()
                counter.value += 1
                lock.release()
                set_uploaded(path_tuple[0], table_name)

            print_status(counter, lock)

            path_tuple = result.fetchone()
        lock.acquire()
        #Executes on the last range.
    lock.release()


def get_total_to_upload(table_name):
    '''Given a table_name, get the total number of rows'''

    query = "SELECT COUNT(*) FROM {0}".format(table_name)
<commit_msg>Removed extra check for already uploaded files.<commit_after>import filesegmenter
import hashlib
import os
import shutil
import swiftclient
import sys
import datetime
import socket
import time
import olrcdb
from multiprocessing import Process, Lock, Value


#Settings
AUTH_VERSION = 2
SWIFT_AUTH_URL = ''
USERNAME = ''
PASSWORD = ''
CONTAINER = ''
AUTH_TOKEN = ''
STORAGE_URL = ''
TEMP_DIRECTORY = 'temp'
FILE_LIMIT = 0.5*10**9  # Max file size in bytes that a file can be uploaded.
# Anything larger is segmented
SEGMENT_SIZE = 100*10**6
COUNT = 0
TOTAL = 0
FAILED_COUNT = 0
LIMIT = 1000
RANGE = 0  # protected variable
SLEEP = 100

REQUIRED_VARIABLES = [
    'OS_AUTH_URL',
    'OS_USERNAME',
    'OS_TENANT_NAME',
    'OS_PASSWORD',
    "MYSQL_HOST",
    "MYSQL_USER",
    "MYSQL_PASSWD",
    "MYSQL_DB"
]


#Currently not being used
def olrc_upload(path):
    ''' Given a path, upload it to the OLRC. Return False if upload
    unsuccessful.'''

    global FAILED_COUNT

    # Check file not already online.
    #if not is_uploaded(path):

    if not (olrc_upload_file(path)):
        FAILED_COUNT += 1
        error_log = open('error.log', 'a')
        error_log.write("\rFailed: {0}\n".format(path))
        error_log.close()
        return False

    return True


#Currently not being used
def olrc_upload_segments(source_file, target_directory):
    ''' Break up the source_file into segments and upload them into
    target_directory'''

    segments = filesegmenter.split_file(
        source_file,
        'temp',
        SEGMENT_SIZE
    )

    sys.stdout.flush()
    sys.stdout.write("\rPartitioning file {0}".format(source_file))

    for segment in segments:

        # Files are within the temp directory locally. On the server the file
        # will live in the target_directory so we need to remove 'temp/' from
        # target_file.
        target_file = os.path.join(target_directory, segment.split('/', 1)[1])
        sys.stdout.flush()
        sys.stdout.write("\rUploading file {0}".format(segment))

        # MDCheck
        if not is_uploaded(segment, target_file):
            if not olrc_upload_file(segment, target_file):
                return False
        else:

            sys.stdout.flush()
            sys.stdout.write(
                "\rSkipping: {0}, already uploaded.".format(source_file)
            )

    # Create and upload readme file.

    readme = "The file in this directory has been segmented for convenient" \
        " upload and download. To assemble the file, run the following " \
        "command on your machine in the directory with all the segments: " \
        "\n\ncat * >> {0}".format(source_file.split('/')[-1])
    outFile = open("temp/readme.txt", "wt")
    outFile.write(readme)
    outFile.close()

    olrc_upload_file(
        "temp/readme.txt",
        os.path.join(target_directory, "readme.txt")
    )
    #Clean up temp files
    if os.path.isdir(TEMP_DIRECTORY):
        shutil.rmtree(TEMP_DIRECTORY)

    return True


def olrc_upload_file(path):
    '''Given String source_file, upload the file to the OLRC to target_file
     and return True if successful. '''

    global SLEEP
    try:
        opened_source_file = open(path, 'r')
    except IOError:
        print("Error opening: " + path)
        return False
    try:
        swiftclient.client.put_object(
            STORAGE_URL,
            AUTH_TOKEN,
            CONTAINER,
            path,
            opened_source_file)

    except Exception, e:

        olrc_connect()
        sys.stdout.flush()
        sys.stdout.write("\rError! {0}\n".format(e))
        sys.stdout.write(
            "Error! {0} Upload to OLRC failed."
            " Trying again in {1} seconds.\n".format(
                time.strftime("%Y-%m-%d %H:%M:%S", time.gmtime()),
                SLEEP
            )
        )
        time.sleep(SLEEP)
        SLEEP += SLEEP
        olrc_upload_file(path)

    return True


def olrc_connect():
    '''Connect to the OLRC with the global variables. Exit if connection
    fails.'''

    global SLEEP
    try:
        (connection_storage_url, auth_token) = swiftclient.client.get_auth(
            SWIFT_AUTH_URL, USERNAME, PASSWORD,
            auth_version=AUTH_VERSION)
        global AUTH_TOKEN
        AUTH_TOKEN = auth_token
        global STORAGE_URL
        STORAGE_URL = connection_storage_url
    except swiftclient.client.ClientException, e:
        print(e)
        sys.stdout.flush()
        sys.stdout.write(
            "\rError! {0} Connection to OLRC failed."
            " Trying again in {1} seconds.\n".format(
                time.strftime("%Y-%m-%d %H:%M:%S", time.gmtime()),
                SLEEP
            )
        )
        time.sleep(SLEEP)
        SLEEP += SLEEP

        olrc_connect()


#Currently not being used
def is_uploaded(file_name):
    '''Return True if String file is already on the server and its etag
    matches it's md5. Delete the file from the server if the
    md5 does not match.'''

    # Swift stat on filename.
    try:
        object_stat = swiftclient.client.head_object(
            STORAGE_URL,
            AUTH_TOKEN,
            CONTAINER,
            file_name
        )
        try:
            etag = object_stat['etag']
        except:
            # Return if no etag
            return False

        md5 = checksum_md5(file_name)
        match = etag == md5

        # Delete the file if the md5 does not match.
        if not match:
            try:
                swiftclient.client.delete_object(
                    STORAGE_URL,
                    AUTH_TOKEN,
                    CONTAINER,
                    file_name
                )
            except:
                pass

        return match
    except:
        return False


#Currently not being used
def checksum_md5(filename):
    md5 = hashlib.md5()
    with open(filename, 'rb') as f:
        for chunk in iter(lambda: f.read(8192), b''):
            md5.update(chunk)
    return md5.hexdigest()


def is_env_vars_set():
    '''Check all the required environment variables are set. Return false if
    any of them are undefined.'''

    global REQUIRED_VARIABLES
    for required_variable in REQUIRED_VARIABLES:
        if not os.environ.get(required_variable):
            return False

    return True


def set_env_vars():
    '''Set the global variables for swift client assuming they exist in the
    environment.'''

    global SWIFT_AUTH_URL
    SWIFT_AUTH_URL = os.environ.get("OS_AUTH_URL")

    global USERNAME
    USERNAME = os.environ.get("OS_TENANT_NAME") + \
        ":" + os.environ.get("OS_USERNAME")

    global PASSWORD
    PASSWORD = os.environ.get("OS_PASSWORD")

    return


def upload_table(lock, range, table_name, counter):
    '''
    Given a table_name, upload all the paths from the table where upload is 0.
    Upload within a LIMIT range at a time.
    '''
    global FAILED_COUNT, LIMIT, RANGE

    connect = olrcdb.DatabaseConnection()
    lock.acquire()
    while range.value <= TOTAL:

        # Access protected variable.
        # We want to make we only fetch within a unique range,
        # so RANGE is locked.
            # "grab" a LIMITs with of rows at a time.
        query = (
            "SELECT * FROM {0} WHERE uploaded=0"
            " AND id >= {1} AND id <{2}".format(
                table_name, range.value, range.value + LIMIT))
        # Let other processes know this range.value has been accounted for.
        range.value += LIMIT
        lock.release()

        # etch results.
        result = connect.execute_query(query)
        path_tuple = result.fetchone()
        # Loop until we run out of rows from the database.
        while (path_tuple):

            # if the upload is successful, update the database
            if olrc_upload_file(path_tuple[1]):
                lock.acquire()
                counter.value += 1
                lock.release()
                set_uploaded(path_tuple[0], table_name)
            else:

                FAILED_COUNT += 1
                error_log = open('error.log', 'a')
                error_log.write("\rFailed: {0}\n".format(path))
                error_log.close()

            print_status(counter, lock)

            path_tuple = result.fetchone()
        lock.acquire()
        #Executes on the last range.
    lock.release()


def get_total_to_upload(table_name):
    '''Given a table_name, get the total number of rows'''

    query = "SELECT COUNT(*) FROM {0}".format(table_name)

{'fraction_non_alphanumeric': 0.09430231861482236, 'fraction_numerical': 0.007429153435983464, 'mean_word_length': 2.061067302402347, 'pattern_counts': {'":': 2, '<': 7, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2073222', 'n_tokens_mistral': 5159, 'n_tokens_neox': 4954, 'n_words': 1418}