<commit_before>import sys, os
import gridfs
import uuid
from pymongo import MongoClient
import re
import bson
import random
import StringIO

def cleanFileName(fileName):
    fileNameSplit = fileName.split('/')
    nameLength = len(fileNameSplit)
    if "revision" in fileNameSplit:
        return fileNameSplit[nameLength - 2] +"/" +  fileNameSplit[nameLength - 1]
    else:
        return fileNameSplit[nameLength - 1]

def genFileDir(fileName, dirLevels):
    directory = ''
    minChunkLen = 4
    nameChunkLen = len(fileName) / dirLevels
    if nameChunkLen < minChunkLen:
        nameChunkLen = minChunkLen
    for i in range(dirLevels):
        chunkStart = (i * nameChunkLen) % len(fileName)
        fileNameHash = hash(fileName[chunkStart : chunkStart + nameChunkLen] + str(random.random()))
        directory += str(fileNameHash & 255) + '/'
    return directory

if len(sys.argv) < 4:
    print("Not enough arguments.")
    print("moveModelsToFS.py <mongoURL> <mongoPort> <userName> <password> <localFolder>")
    sys.exit(0)

mongoURL = sys.argv[1]
mongoPort = sys.argv[2]
userName = sys.argv[3]
password = sys.argv[4]
localFolder = sys.argv[5]

compression = None

if len(sys.argv) > 6:
    compression = sys.argv[6]

if not os.path.exists(localFolder):
    print("LocalFolder " + localFolder + " does not exist.")
    sys.exit(0)

connString = "mongodb://"+ userName + ":" + password +"@"+mongoURL + ":" + mongoPort + "/"

##### Enable dry run to not commit to the database #####
dryRun = True

##### Number of directory levels to use #####
dirLevels = 2

##### Connect to the Database #####
db = MongoClient(connString)
for database in db.database_names():
    if database != "admin" and database != "local":
        db = MongoClient(connString)[database]
        print("--database:" + database)

##### Get a model ID and find entries #####
        for setting in db.settings.find(no_cursor_timeout=True):
            modelId = setting.get('_id')
            print("\t--model: " +  modelId)
            for colPrefix in [".history",".stash.json_mpc", ".stash.unity3d" ]:
                colName = modelId + colPrefix
                print("\t\t--stash: " +  colName)
                fs = gridfs.GridFS(db, colName)
                for entry in fs.find({"filename":{"$not": re.compile("unityAssets.json$")}}):
                    fileId = str(uuid.uuid4())
                    directory = genFileDir(fileId, dirLevels)
                    fileLink = directory + fileId
                    filePath = localFolder + '/' + fileLink
##### Create Reference BSON #####
                    bsonData = {}
                    bsonData['_id'] = cleanFileName(entry.filename)
                    bsonData['link'] = fileLink
                    bsonData['type'] = "fs"
                    bsonData['size'] = entry.length
                    if compression:
                        bsonData['compression'] = compression # snappy / none

                    if dryRun:
                        print( "\t\t Writing: " + str(bsonData))
                        print("\t\t Path:" + str(filePath))
                    else:
##### Upload to FS and insert BSON #####
                        if not os.path.exists(os.path.dirname(filePath)):
                            os.makedirs(os.path.dirname(filePath))
                        file = open(filePath,'wb')
                        file.write(StringIO.StringIO(entry.read()).getvalue())
                        file.close()
                        targetCol = colName + ".ref"
                        db[targetCol].save(bsonData)
<commit_msg>ISSUE #1487 introduce overwrite flag (#1488)<commit_after>import sys, os
import gridfs
import uuid
from pymongo import MongoClient
import re
import bson
import random
import StringIO

def cleanFileName(fileName):
    fileNameSplit = fileName.split('/')
    nameLength = len(fileNameSplit)
    if "revision" in fileNameSplit:
        return fileNameSplit[nameLength - 2] +"/" +  fileNameSplit[nameLength - 1]
    else:
        return fileNameSplit[nameLength - 1]

def genFileDir(fileName, dirLevels):
    directory = ''
    minChunkLen = 4
    nameChunkLen = len(fileName) / dirLevels
    if nameChunkLen < minChunkLen:
        nameChunkLen = minChunkLen
    for i in range(dirLevels):
        chunkStart = (i * nameChunkLen) % len(fileName)
        fileNameHash = hash(fileName[chunkStart : chunkStart + nameChunkLen] + str(random.random()))
        directory += str(fileNameHash & 255) + '/'
    return directory

if len(sys.argv) <= 5:
    print("Not enough arguments.")
    print("moveModelsToFS.py <mongoURL> <mongoPort> <userName> <password> <localFolder>")
    sys.exit(0)

mongoURL = sys.argv[1]
mongoPort = sys.argv[2]
userName = sys.argv[3]
password = sys.argv[4]
localFolder = sys.argv[5]

compression = None

if len(sys.argv) > 6:
    compression = sys.argv[6]

if not os.path.exists(localFolder):
    print("LocalFolder " + localFolder + " does not exist.")
    sys.exit(0)

connString = "mongodb://"+ userName + ":" + password +"@"+mongoURL + ":" + mongoPort + "/"

##### Enable dry run to not commit to the database #####
dryRun = True
overwrite = True #if there is already an entry for the filename: True = Overwrite regardless, False = Use existing entry

##### Number of directory levels to use #####
dirLevels = 2

##### Connect to the Database #####
db = MongoClient(connString)
for database in db.database_names():
    if database != "admin" and database != "local":
        db = MongoClient(connString)[database]
        print("--database:" + database)

##### Get a model ID and find entries #####
        for setting in db.settings.find(no_cursor_timeout=True):
            modelId = setting.get('_id')
            print("\t--model: " +  modelId)
            for colPrefix in [".history",".stash.json_mpc", ".stash.unity3d" ]:
                colName = modelId + colPrefix
                targetCol = colName + ".ref"
                print("\t\t--stash: " +  colName)
                fs = gridfs.GridFS(db, colName)
                for entry in fs.find({"filename":{"$not": re.compile("unityAssets.json$")}}):
                    fileId = str(uuid.uuid4())
                    directory = genFileDir(fileId, dirLevels)
                    fileLink = directory + fileId
                    filePath = localFolder + '/' + fileLink
##### Create Reference BSON #####
                    if not overwrite:
                        if db[targetCol].find_one({"_id" : cleanFileName(entry.filename), "type": "fs"}) != None:
                            print("\t\t Found entry for " + cleanFileName(entry.filename) + ", skipping...");
                            continue

                    bsonData = {}
                    bsonData['_id'] = cleanFileName(entry.filename)
                    bsonData['link'] = fileLink
                    bsonData['type'] = "fs"
                    bsonData['size'] = entry.length
                    if compression:
                        bsonData['compression'] = compression # snappy / none

                    if dryRun:
                        print( "\t\t Writing: " + str(bsonData))
                        print("\t\t Path:" + str(filePath))
                    else:
##### Upload to FS and insert BSON #####
                        if not os.path.exists(os.path.dirname(filePath)):
                            os.makedirs(os.path.dirname(filePath))
                        file = open(filePath,'wb')
                        file.write(StringIO.StringIO(entry.read()).getvalue())
                        file.close()
                        db[targetCol].save(bsonData)

{'fraction_non_alphanumeric': 0.1312852829197847, 'fraction_numerical': 0.006301693580149665, 'mean_word_length': 2.09926769731489, 'pattern_counts': {'":': 11, '<': 17, '<?xml version=': 0, '>': 15, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2756948', 'n_tokens_mistral': 2228, 'n_tokens_neox': 2150, 'n_words': 527}