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

from boto.exception import GSResponseError
from boto.exception import S3ResponseError
from boto.gs.key import Key as GSKey
from boto.s3.key import Key as S3Key
from multiprocessing.pool import ThreadPool
from os import path, rename
from time import time

import argparse
import boto
import socket
import sqlite3
import sys
import tempfile

# Seemingly unused but needed for GCS authentication
import gcs_oauth2_boto_plugin


def get_bucket(bucket_uri, validate=False):
    scheme = bucket_uri[:2]
    bucket_name = bucket_uri[5:]
    connection = {
        'gs': boto.connect_gs,
        's3': boto.connect_s3,
    }[scheme]()
    return connection.get_bucket(bucket_name, validate=validate)


def get_key(bucket_uri, path):
    scheme = bucket_uri[:2]
    constructor = {
        'gs': GSKey,
        's3': S3Key,
    }[scheme]
    return constructor(get_bucket(bucket_uri), path)


# Download an object from source bucket, then upload it to destination bucket
def transfer(source, destination, path):
    while True:
        try:
            # Roll over when hitting 10 MB
            f = tempfile.SpooledTemporaryFile(max_size=10*2**20)
            get_key(source, path).get_contents_to_file(f)
            get_key(destination, path).set_contents_from_file(f, rewind=True)
            return path
        except socket.error as e:
            print e
        except (GSResponseError, S3ResponseError) as e:
            if e.status == 404:
                # TODO: Log this as a warning somewhere
                return path
            else:
                raise


# Remove an object from destination bucket, ignoring "not found" errors
def remove(destination, path):
    bucket = get_bucket(destination)
    try:
        bucket.delete_key(path)
    except (GSResponseError, S3ResponseError) as e:
        if e.status != 404:
            raise


def finished(connection):
    cursor = connection.cursor()
    cursor.execute('''SELECT 1 FROM source WHERE NOT processed LIMIT 1''')
    return cursor.fetchone() is None


def process(source, destination, connection):

    print 'Skipping over up-to-date objects...'
    connection.execute('''
        UPDATE source SET processed = 1
        WHERE rowid IN (
            SELECT s.rowid FROM source s JOIN destination d
            ON s.path = d.path AND s.hash = d.hash
        )
    ''')

    print 'Uploading new/updated objects from source to destination...'
    while not finished(connection):
        sql_it = connection.execute('''SELECT path FROM source WHERE NOT processed LIMIT 1000''')
        pool = ThreadPool()
        pool_it = pool.imap_unordered(lambda row: transfer(source, destination, row[0]),
                                      list(sql_it))
        for path in pool_it:
            connection.execute('''UPDATE source SET processed = 1 WHERE path = ?''', (path,))
            print 'Finished: %s' % path

    print 'Deleting objects in destination that have been deleted in source...'
    for row in connection.execute('''
        SELECT d.rowid, d.path
        FROM destination d LEFT JOIN source s
        ON d.path = s.path WHERE s.path IS NULL<commit_msg>print full path for finished objects<commit_after>#!/usr/bin/env python

from boto.exception import GSResponseError
from boto.exception import S3ResponseError
from boto.gs.key import Key as GSKey
from boto.s3.key import Key as S3Key
from multiprocessing.pool import ThreadPool
from os import path, rename
from time import time

import argparse
import boto
import socket
import sqlite3
import sys
import tempfile

# Seemingly unused but needed for GCS authentication
import gcs_oauth2_boto_plugin


def get_bucket(bucket_uri, validate=False):
    scheme = bucket_uri[:2]
    bucket_name = bucket_uri[5:]
    connection = {
        'gs': boto.connect_gs,
        's3': boto.connect_s3,
    }[scheme]()
    return connection.get_bucket(bucket_name, validate=validate)


def get_key(bucket_uri, path):
    scheme = bucket_uri[:2]
    constructor = {
        'gs': GSKey,
        's3': S3Key,
    }[scheme]
    return constructor(get_bucket(bucket_uri), path)


# Download an object from source bucket, then upload it to destination bucket
def transfer(source, destination, path):
    while True:
        try:
            # Roll over when hitting 10 MB
            f = tempfile.SpooledTemporaryFile(max_size=10*2**20)
            get_key(source, path).get_contents_to_file(f)
            get_key(destination, path).set_contents_from_file(f, rewind=True)
            return path
        except socket.error as e:
            print e
        except (GSResponseError, S3ResponseError) as e:
            if e.status == 404:
                # TODO: Log this as a warning somewhere
                return path
            else:
                raise


# Remove an object from destination bucket, ignoring "not found" errors
def remove(destination, path):
    bucket = get_bucket(destination)
    try:
        bucket.delete_key(path)
    except (GSResponseError, S3ResponseError) as e:
        if e.status != 404:
            raise


def finished(connection):
    cursor = connection.cursor()
    cursor.execute('''SELECT 1 FROM source WHERE NOT processed LIMIT 1''')
    return cursor.fetchone() is None


def process(source, destination, connection):

    print 'Skipping over up-to-date objects...'
    connection.execute('''
        UPDATE source SET processed = 1
        WHERE rowid IN (
            SELECT s.rowid FROM source s JOIN destination d
            ON s.path = d.path AND s.hash = d.hash
        )
    ''')

    print 'Uploading new/updated objects from source to destination...'
    while not finished(connection):
        sql_it = connection.execute('''SELECT path FROM source WHERE NOT processed LIMIT 1000''')
        pool = ThreadPool()
        pool_it = pool.imap_unordered(lambda row: transfer(source, destination, row[0]),
                                      list(sql_it))
        for path in pool_it:
            connection.execute('''UPDATE source SET processed = 1 WHERE path = ?''', (path,))
            print 'Finished: %s/%s -> %s/%s' % (source, path, destination, path)

    print 'Deleting objects in destination that have been deleted in source...'
    for row in connection.execute('''
        SELECT d.rowid, d.path
        FROM destination d LEFT JOIN source s
        ON d.path = s.path WHERE s.path IS NULL
{'fraction_non_alphanumeric': 0.08648142366531376, 'fraction_numerical': 0.01123946300343428, 'mean_word_length': 2.6034870641169854, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 4, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2070660', 'n_tokens_mistral': 1857, 'n_tokens_neox': 1764, 'n_words': 604}