<commit_before>#!/usr/bin/env python
from __future__ import print_function

import signal

import sys, os, datetime, time, json, traceback, csv
import multiprocessing, multiprocessing.queues, subprocess, re, ctypes, codecs
from optparse import OptionParser
from ._backup import *
import rethinkdb as r

# Used because of API differences in the csv module, taken from
# http://python3porting.com/problems.html
PY3 = sys.version > '3'

#json parameters
json_read_chunk_size = 32 * 1024
json_max_buffer_size = 128 * 1024 * 1024
max_nesting_depth = 100
try:
    import cPickle as pickle
except ImportError:
    import pickle
try:
    from itertools import imap
except ImportError:
    imap = map
try:
    xrange
except NameError:
    xrange = range
try:
    from multiprocessing import SimpleQueue
except ImportError:
    from multiprocessing.queues import SimpleQueue

info = "'rethinkdb import` loads data into a RethinkDB cluster"
usage = "\
  rethinkdb import -d DIR [-c HOST:PORT] [-a AUTH_KEY] [--force]\n\
      [-i (DB | DB.TABLE)] [--clients NUM]\n\
      [--shards NUM_SHARDS] [--replicas NUM_REPLICAS]\n\
  rethinkdb import -f FILE --table DB.TABLE [-c HOST:PORT] [-a AUTH_KEY]\n\
      [--force] [--clients NUM] [--format (csv | json)] [--pkey PRIMARY_KEY]\n\
      [--shards NUM_SHARDS] [--replicas NUM_REPLICAS]\n\
      [--delimiter CHARACTER] [--custom-header FIELD,FIELD... [--no-header]]"

def print_import_help():
    print(info)
    print(usage)
    print("")
    print("  -h [ --help ]                    print this help")
    print("  -c [ --connect ] HOST:PORT       host and client port of a rethinkdb node to connect")
    print("                                   to (defaults to localhost:28015)")
    print("  -a [ --auth ] AUTH_KEY           authorization key for rethinkdb clients")
    print("  --clients NUM_CLIENTS            the number of client connections to use (defaults")
    print("                                   to 8)")
    print("  --hard-durability                use hard durability writes (slower, but less memory")
    print("                                   consumption on the server)")
    print("  --force                          import data even if a table already exists, and")
    print("                                   overwrite duplicate primary keys")
    print("  --fields                         limit which fields to use when importing one table")
    print("")
    print("Import directory:")
    print("  -d [ --directory ] DIR           the directory to import data from")
    print("  -i [ --import ] (DB | DB.TABLE)  limit restore to the given database or table (may")
    print("                                   be specified multiple times)")
    print("  --no-secondary-indexes           do not create secondary indexes for the imported tables")
    print("")
    print("Import file:")
    print("  -f [ --file ] FILE               the file to import data from")
    print("  --table DB.TABLE                 the table to import the data into")
    print("  --format (csv | json)            the format of the file (defaults to json and accepts")
    print("                                   newline delimited json)")
    print("  --pkey PRIMARY_KEY               the field to use as the primary key in the table")
    print("")
    print("Import CSV format:")
    print("  --delimiter CHARACTER            character separating fields, or '\\t' for tab")
    print("  --no-header                      do not read in a header of field names")
    print("  --custom-header FIELD,FIELD...   header to use (overriding file header), must be")
    print("                                   specified if --no-header")
    print("")
    print("Import JSON format:")
    print("  --max-document-size              the maximum size in bytes that a single JSON document")
    print("                                   can have (defaults to 134217728).")
    print("  --max-nesting-depth              the maximum nesting depth of the JSON documents")
    print("")
    print("EXAMPLES:")
    print("")
    print("rethinkdb import -d rdb_export -c mnemosyne:39500 --clients 128")
    print("  Import data into a cluster running on host 'mnemosyne' with a client port at 39500,")
    print("  using 128 client connections and the named export directory.")
    print("")
    print("rethinkdb import -f site_history.csv --format csv --table test.history --pkey count")
    print("  Import data into a local cluster and the table 'history' in the 'test' database,")
    print("  using the named CSV file, and using the 'count' field as the primary key.")
    print("")
    print("rethinkdb import -d rdb_export -c hades -a hunter2 -i test")
    print("  Import data into a cluster running on host 'hades' which requires authorization,")
    print("  using only the database 'test' from the named export directory.")
    print("")
    print("rethinkdb import -f subscriber_info.json --fields id,name,hashtag --force")
    print("  Import data into a local cluster using the named JSON file, and only the fields")
    print("  'id', 'name', and 'hashtag', overwriting any existing rows with the same primary key.")
    print("")
    print("rethinkdb import -f user_data.csv --delimiter ';' --no-header --custom-header id,name,number")
    print("  Import data into a local cluster using the named CSV file with no header and instead")
    print("  use the fields 'id', 'name', and 'number', the delimiter is a semicolon (rather than")
    print("  a comma).")

def parse_options():
    parser = OptionParser(add_help_option=False, usage=usage)
    parser.add_option("-c", "--connect", dest="host", metavar="HOST:PORT", default="localhost:28015", type="string")
    parser.add_option("-a", "--auth", dest="auth_key", metavar="AUTHKEY", default="", type="string")
    parser.add_option("--fields", dest="fields", metavar="FIELD,FIELD...", default=None, type="string")
    parser.add_option("--clients", dest="clients", metavar="NUM_CLIENTS", default=8, type="int")
    parser.add_option("--hard-durability", dest="hard", action="store_true", default=False)
    parser.add_option("--force", dest="force", action="store_true", default=False)
    parser.add_option("--debug", dest="debug", action="store_true", default=False)
    parser.add_option("--max-document-size", dest="max_document_size",  default=0,type="int")
    parser.add_option("--max-nesting-depth", dest="max_nesting_depth", default=0, type="int")

    # Replication settings
    parser.add_option("--shards", dest="shards", metavar="NUM_SHARDS", default=0, type="int")
    parser.add_option("--replicas", dest="replicas", metavar="NUM_REPLICAS", default=0, type="int")

    # Directory import options
    parser.add_option("-d", "--directory", dest="directory", metavar="DIRECTORY", default=None, type="string")
    parser.add_option("-i", "--import", dest="tables", metavar="DB | DB.TABLE", default=[], action="append", type="string")
    parser.add_option("--no-secondary-indexes", dest="create_sindexes", action="store_false", default=True)

    # File import options
    parser.add_option("-f", "--file", dest="import_file", metavar="FILE", default=None, type="string")
    parser.add_option("--format", dest="import_format", metavar="json | csv", default=None, type="string")
    parser.add_option("--table", dest="import_table", metavar="DB.TABLE", default=None, type="string")
    parser.add_option("--pkey", dest="primary_key", metavar="KEY", default=None, type="string")
    parser.add_option("--delimiter", dest="delimiter", metavar="CHARACTER", default=None, type="string")
    parser.add_option("--no-header", dest="no_header", action="store_true", default=False)
    parser.add_option("--custom-header", dest="custom_header", metavar="FIELD,FIELD...", default=None, type="string")
    parser.add_option("-h", "--help", dest="help", default=False, action="store_true")
    (options, args) = parser.parse_args()

    # Check validity of arguments
    if len(args) != 0:
        raise RuntimeError("Error: No positional arguments supported. Unrecognized option '%s'" % args[0])

    if options.help:
        print_import_help()
        exit(0)

    res = {}

    # Verify valid host:port --connect option
    (res["host"], res["port"]) = parse_connect_option(options.host)

    if options.clients < 1:
        raise RuntimeError("Error: --client option too low, must have at least one client connection")

    res["auth_key"] = options.auth_key
    res["clients"] = options.clients
    res["durability"] = "hard" if options.hard else "soft"
    res["force"] = options.force
    res["debug"] = options.debug
    res["create_sindexes"] = options.create_sindexes

    res["create_args"] = {k: getattr(options, k) for k in ['shards', 'replicas'] if getattr(options, k) != 0}
<commit_msg>fixing python2.6 compatibility in import script<commit_after>#!/usr/bin/env python
from __future__ import print_function

import signal

import sys, os, datetime, time, json, traceback, csv
import multiprocessing, multiprocessing.queues, subprocess, re, ctypes, codecs
from optparse import OptionParser
from ._backup import *
import rethinkdb as r

# Used because of API differences in the csv module, taken from
# http://python3porting.com/problems.html
PY3 = sys.version > '3'

#json parameters
json_read_chunk_size = 32 * 1024
json_max_buffer_size = 128 * 1024 * 1024
max_nesting_depth = 100
try:
    import cPickle as pickle
except ImportError:
    import pickle
try:
    from itertools import imap
except ImportError:
    imap = map
try:
    xrange
except NameError:
    xrange = range
try:
    from multiprocessing import SimpleQueue
except ImportError:
    from multiprocessing.queues import SimpleQueue

info = "'rethinkdb import` loads data into a RethinkDB cluster"
usage = "\
  rethinkdb import -d DIR [-c HOST:PORT] [-a AUTH_KEY] [--force]\n\
      [-i (DB | DB.TABLE)] [--clients NUM]\n\
      [--shards NUM_SHARDS] [--replicas NUM_REPLICAS]\n\
  rethinkdb import -f FILE --table DB.TABLE [-c HOST:PORT] [-a AUTH_KEY]\n\
      [--force] [--clients NUM] [--format (csv | json)] [--pkey PRIMARY_KEY]\n\
      [--shards NUM_SHARDS] [--replicas NUM_REPLICAS]\n\
      [--delimiter CHARACTER] [--custom-header FIELD,FIELD... [--no-header]]"

def print_import_help():
    print(info)
    print(usage)
    print("")
    print("  -h [ --help ]                    print this help")
    print("  -c [ --connect ] HOST:PORT       host and client port of a rethinkdb node to connect")
    print("                                   to (defaults to localhost:28015)")
    print("  -a [ --auth ] AUTH_KEY           authorization key for rethinkdb clients")
    print("  --clients NUM_CLIENTS            the number of client connections to use (defaults")
    print("                                   to 8)")
    print("  --hard-durability                use hard durability writes (slower, but less memory")
    print("                                   consumption on the server)")
    print("  --force                          import data even if a table already exists, and")
    print("                                   overwrite duplicate primary keys")
    print("  --fields                         limit which fields to use when importing one table")
    print("")
    print("Import directory:")
    print("  -d [ --directory ] DIR           the directory to import data from")
    print("  -i [ --import ] (DB | DB.TABLE)  limit restore to the given database or table (may")
    print("                                   be specified multiple times)")
    print("  --no-secondary-indexes           do not create secondary indexes for the imported tables")
    print("")
    print("Import file:")
    print("  -f [ --file ] FILE               the file to import data from")
    print("  --table DB.TABLE                 the table to import the data into")
    print("  --format (csv | json)            the format of the file (defaults to json and accepts")
    print("                                   newline delimited json)")
    print("  --pkey PRIMARY_KEY               the field to use as the primary key in the table")
    print("")
    print("Import CSV format:")
    print("  --delimiter CHARACTER            character separating fields, or '\\t' for tab")
    print("  --no-header                      do not read in a header of field names")
    print("  --custom-header FIELD,FIELD...   header to use (overriding file header), must be")
    print("                                   specified if --no-header")
    print("")
    print("Import JSON format:")
    print("  --max-document-size              the maximum size in bytes that a single JSON document")
    print("                                   can have (defaults to 134217728).")
    print("  --max-nesting-depth              the maximum nesting depth of the JSON documents")
    print("")
    print("EXAMPLES:")
    print("")
    print("rethinkdb import -d rdb_export -c mnemosyne:39500 --clients 128")
    print("  Import data into a cluster running on host 'mnemosyne' with a client port at 39500,")
    print("  using 128 client connections and the named export directory.")
    print("")
    print("rethinkdb import -f site_history.csv --format csv --table test.history --pkey count")
    print("  Import data into a local cluster and the table 'history' in the 'test' database,")
    print("  using the named CSV file, and using the 'count' field as the primary key.")
    print("")
    print("rethinkdb import -d rdb_export -c hades -a hunter2 -i test")
    print("  Import data into a cluster running on host 'hades' which requires authorization,")
    print("  using only the database 'test' from the named export directory.")
    print("")
    print("rethinkdb import -f subscriber_info.json --fields id,name,hashtag --force")
    print("  Import data into a local cluster using the named JSON file, and only the fields")
    print("  'id', 'name', and 'hashtag', overwriting any existing rows with the same primary key.")
    print("")
    print("rethinkdb import -f user_data.csv --delimiter ';' --no-header --custom-header id,name,number")
    print("  Import data into a local cluster using the named CSV file with no header and instead")
    print("  use the fields 'id', 'name', and 'number', the delimiter is a semicolon (rather than")
    print("  a comma).")

def parse_options():
    parser = OptionParser(add_help_option=False, usage=usage)
    parser.add_option("-c", "--connect", dest="host", metavar="HOST:PORT", default="localhost:28015", type="string")
    parser.add_option("-a", "--auth", dest="auth_key", metavar="AUTHKEY", default="", type="string")
    parser.add_option("--fields", dest="fields", metavar="FIELD,FIELD...", default=None, type="string")
    parser.add_option("--clients", dest="clients", metavar="NUM_CLIENTS", default=8, type="int")
    parser.add_option("--hard-durability", dest="hard", action="store_true", default=False)
    parser.add_option("--force", dest="force", action="store_true", default=False)
    parser.add_option("--debug", dest="debug", action="store_true", default=False)
    parser.add_option("--max-document-size", dest="max_document_size",  default=0,type="int")
    parser.add_option("--max-nesting-depth", dest="max_nesting_depth", default=0, type="int")

    # Replication settings
    parser.add_option("--shards", dest="shards", metavar="NUM_SHARDS", default=0, type="int")
    parser.add_option("--replicas", dest="replicas", metavar="NUM_REPLICAS", default=0, type="int")

    # Directory import options
    parser.add_option("-d", "--directory", dest="directory", metavar="DIRECTORY", default=None, type="string")
    parser.add_option("-i", "--import", dest="tables", metavar="DB | DB.TABLE", default=[], action="append", type="string")
    parser.add_option("--no-secondary-indexes", dest="create_sindexes", action="store_false", default=True)

    # File import options
    parser.add_option("-f", "--file", dest="import_file", metavar="FILE", default=None, type="string")
    parser.add_option("--format", dest="import_format", metavar="json | csv", default=None, type="string")
    parser.add_option("--table", dest="import_table", metavar="DB.TABLE", default=None, type="string")
    parser.add_option("--pkey", dest="primary_key", metavar="KEY", default=None, type="string")
    parser.add_option("--delimiter", dest="delimiter", metavar="CHARACTER", default=None, type="string")
    parser.add_option("--no-header", dest="no_header", action="store_true", default=False)
    parser.add_option("--custom-header", dest="custom_header", metavar="FIELD,FIELD...", default=None, type="string")
    parser.add_option("-h", "--help", dest="help", default=False, action="store_true")
    (options, args) = parser.parse_args()

    # Check validity of arguments
    if len(args) != 0:
        raise RuntimeError("Error: No positional arguments supported. Unrecognized option '%s'" % args[0])

    if options.help:
        print_import_help()
        exit(0)

    res = {}

    # Verify valid host:port --connect option
    (res["host"], res["port"]) = parse_connect_option(options.host)

    if options.clients < 1:
        raise RuntimeError("Error: --client option too low, must have at least one client connection")

    res["auth_key"] = options.auth_key
    res["clients"] = options.clients
    res["durability"] = "hard" if options.hard else "soft"
    res["force"] = options.force
    res["debug"] = options.debug
    res["create_sindexes"] = options.create_sindexes

    res["create_args"] = { }
    for k in ['shards', 'replicas']:
        if getattr(options, k) != 0:
            res["create_args"][k] = getattr(options, k)

{'fraction_non_alphanumeric': 0.13883709103237837, 'fraction_numerical': 0.008023958863084138, 'mean_word_length': 2.9699416778824586, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 18, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2252405', 'n_tokens_mistral': 5352, 'n_tokens_neox': 5137, 'n_words': 1627}