_import.py

#!/usr/bin/env python

# Copyright 2018 RethinkDB
#
# Licensed under the Apache License, Version 2.0 (the 'License');
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an 'AS IS' BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# This file incorporates work covered by the following copyright:
# Copyright 2010-2016 RethinkDB, all rights reserved.

"""`rethinkdb import` loads data into a RethinkDB cluster"""

from __future__ import print_function

import codecs
import collections
import csv
import ctypes
import json
import multiprocessing
import optparse
import os
import signal
import sys
import time
import traceback
from multiprocessing.queues import Queue, SimpleQueue

import six

from rethinkdb import ast, errors, query, utils_common
from rethinkdb.logger import default_logger

try:
    unicode
except NameError:
    unicode = str

try:
    from Queue import Empty, Full
except ImportError:
    from queue import Empty, Full


# json parameters
JSON_READ_CHUNK_SIZE = 128 * 1024
JSON_MAX_BUFFER_SIZE = 128 * 1024 * 1024
MAX_NESTING_DEPTH = 100

Error = collections.namedtuple("Error", ["message", "traceback", "file"])


class SourceFile(object):
    format = None  # set by subclasses

    name = None

    db = None
    table = None
    primary_key = None
    indexes = None
    write_hook = None
    source_options = None

    start_time = None
    end_time = None

    query_runner = None

    _source = None  # open filehandle for the source

    # - internal synchronization variables

    _bytes_size = None
    _bytes_read = None  # -1 until started

    _total_rows = None  # -1 until known
    _rows_read = None
    _rows_written = None

    def __init__(
        self,
        source,
        db,
        table,
        query_runner,
        primary_key=None,
        indexes=None,
        write_hook=None,
        source_options=None,
    ):

        if self.format is None:
            raise AssertionError(
                "{class_name} must have a format".format(
                    class_name=self.__class__.__name__
                )
            )

        if self.db == "rethinkdb":
            raise AssertionError("Can not import tables into the system database")

        # query_runner
        if not isinstance(query_runner, utils_common.RetryQuery):
            raise AssertionError("Query runner is not instance of RetryQuery")

        self.query_runner = query_runner

        # reporting information
        self._bytes_size = multiprocessing.Value(ctypes.c_longlong, -1)
        self._bytes_read = multiprocessing.Value(ctypes.c_longlong, -1)

        self._total_rows = multiprocessing.Value(ctypes.c_longlong, -1)
        self._rows_read = multiprocessing.Value(ctypes.c_longlong, 0)
        self._rows_written = multiprocessing.Value(ctypes.c_longlong, 0)

        # source
        if hasattr(source, "read"):
            if unicode != str or "b" in source.mode:
                # Python2.x or binary file, assume utf-8 encoding
                self._source = codecs.getreader("utf-8")(source)
            else:
                # assume that it has the right encoding on it
                self._source = source
        else:
            try:
                self._source = codecs.open(source, mode="r", encoding="utf-8")
            except IOError as exc:
                default_logger.exception(exc)
                raise ValueError(
                    'Unable to open source file "%s": %s' % (str(source), str(exc))
                )

        if (
            hasattr(self._source, "name")
            and self._source.name
            and os.path.isfile(self._source.name)
        ):
            self._bytes_size.value = os.path.getsize(source)
            if self._bytes_size.value == 0:
                raise ValueError("Source is zero-length: %s" % source)

        # table info
        self.db = db
        self.table = table
        self.primary_key = primary_key
        self.indexes = indexes or []
        self.write_hook = write_hook or []

        # options
        self.source_options = source_options or {
            "create_args": {"primary_key": self.primary_key}
        }

        # name
        if hasattr(self._source, "name") and self._source.name:
            self.name = os.path.basename(self._source.name)
        else:
            self.name = "%s.%s" % (self.db, self.table)

    def __hash__(self):
        return hash((self.db, self.table))

    def get_line(self):
        """Returns a single line from the file"""
        raise NotImplementedError(
            "This needs to be implemented on the %s subclass" % self.format
        )

    # - bytes
    @property
    def bytes_size(self):
        return self._bytes_size.value

    @bytes_size.setter
    def bytes_size(self, value):
        self._bytes_size.value = value

    @property
    def bytes_read(self):
        return self._bytes_read.value

    @bytes_read.setter
    def bytes_read(self, value):
        self._bytes_read.value = value

    # - rows
    @property
    def total_rows(self):
        return self._total_rows.value

    @total_rows.setter
    def total_rows(self, value):
        self._total_rows.value = value

    @property
    def rows_read(self):
        return self._rows_read.value

    @rows_read.setter
    def rows_read(self, value):
        self._rows_read.value = value

    @property
    def rows_written(self):
        return self._rows_written.value

    def add_rows_written(self, increment):  # we have multiple writers to coordinate
        with self._rows_written.get_lock():
            self._rows_written.value += increment

    # - percent done
    @property
    def percent_done(self):
        """return a float between 0 and 1 for a reasonable guess of percentage complete"""
        # assume that reading takes 50% of the time and writing the other 50%
        completed = 0.0  # of 2.0

        # - add read percentage
        if (
            self._bytes_size.value <= 0
            or self._bytes_size.value <= self._bytes_read.value
        ):
            completed += 1.0
        elif self._bytes_read.value < 0 and self._total_rows.value >= 0:
            # done by rows read
            if self._rows_read > 0:
                completed += float(self._rows_read) / float(self._total_rows.value)
        else:
            # done by bytes read
            if self._bytes_read.value > 0:
                completed += float(self._bytes_read.value) / float(
                    self._bytes_size.value
                )

        # - add written percentage
        if self._rows_read.value or self._rows_written.value:
            total_rows = float(self._total_rows.value)
            if total_rows == 0:
                completed += 1.0
            elif total_rows < 0:
                # a guesstimate
                per_row_size = float(self._bytes_read.value) / float(
                    self._rows_read.value
                )
                total_rows = float(self._rows_read.value) + (
                    float(self._bytes_size.value - self._bytes_read.value)
                    / per_row_size
                )
                completed += float(self._rows_written.value) / total_rows
            else:
                # accurate count
                completed += float(self._rows_written.value) / total_rows

        # - return the value
        return completed * 0.5

    def setup_table(self):
        """Ensure that the db, table, and indexes exist and are correct"""

        # - ensure the table exists and is ready
        self.query_runner(
            "create table: %s.%s" % (self.db, self.table),
            ast.expr([self.table])
            .set_difference(query.db(self.db).table_list())
            .for_each(
                query.db(self.db).table_create(
                    query.row,
                    **self.source_options["create_args"]
                    if "create_args" in self.source_options
                    else {}
                )
            ),
        )

        self.query_runner(
            "wait for %s.%s" % (self.db, self.table),
            query.db(self.db).table(self.table).wait(timeout=30),
        )

        # - ensure that the primary key on the table is correct
        primary_key = self.query_runner(
            "primary key %s.%s" % (self.db, self.table),
            query.db(self.db).table(self.table).info()["primary_key"],
        )
        if self.primary_key is None:
            self.primary_key = primary_key
        elif primary_key != self.primary_key:
            raise RuntimeError(
                "Error: table %s.%s primary key was `%s` rather than the expected: %s"
                % (self.db, self.table, primary_key, self.primary_key)
            )

    def restore_indexes(self, warning_queue):
        # recreate secondary indexes - dropping existing on the assumption they are wrong
        if self.indexes:
            existing_indexes = self.query_runner(
                "indexes from: %s.%s" % (self.db, self.table),
                query.db(self.db).table(self.table).index_list(),
            )
            try:
                created_indexes = []
                for index in self.indexes:
                    if index["index"] in existing_indexes:  # drop existing versions
                        self.query_runner(
                            "drop index: %s.%s:%s"
                            % (self.db, self.table, index["index"]),
                            query.db(self.db)
                            .table(self.table)
                            .index_drop(index["index"]),
                        )
                    self.query_runner(
                        "create index: %s.%s:%s"
                        % (self.db, self.table, index["index"]),
                        query.db(self.db)
                        .table(self.table)
                        .index_create(index["index"], index["function"]),
                    )
                    created_indexes.append(index["index"])

                # wait for all of the created indexes to build
                self.query_runner(
                    "waiting for indexes on %s.%s" % (self.db, self.table),
                    query.db(self.db)
                    .table(self.table)
                    .index_wait(query.args(created_indexes)),
                )
            except RuntimeError:
                exception_type, exception_class, trcback = sys.exc_info()
                warning_queue.put(
                    (
                        exception_type,
                        exception_class,
                        traceback.extract_tb(trcback),
                        self._source.name,
                    )
                )

        if self.write_hook:
            self.query_runner(
                "Write hook from: %s.%s" % (self.db, self.table),
                query.db(self.db).table(self.table).get_write_hook(),
            )
            try:
                self.query_runner(
                    "drop hook: %s.%s" % (self.db, self.table),
                    query.db(self.db).table(self.table).set_write_hook(None),
                )
                self.query_runner(
                    "create hook: %s.%s:%s" % (self.db, self.table, self.write_hook),
                    query.db(self.db)
                    .table(self.table)
                    .set_write_hook(self.write_hook["function"]),
                )
            except RuntimeError:
                exception_type, exception_class, trcback = sys.exc_info()
                warning_queue.put(
                    (
                        exception_type,
                        exception_class,
                        traceback.extract_tb(trcback),
                        self._source.name,
                    )
                )

    def batches(self, batch_size=None, warning_queue=None):

        # setup table
        self.setup_table()

        # default batch_size
        if batch_size is None:
            batch_size = utils_common.default_batch_size
        else:
            batch_size = int(batch_size)

        if batch_size <= 0:
            raise AssertionError("Batch size can not be less than one")

        # setup
        self.setup_file(warning_queue=warning_queue)

        # - yield batches

        batch = []
        try:
            need_more_data = False
            while True:
                if need_more_data:
                    self.fill_buffer()
                    need_more_data = False

                while len(batch) < batch_size:
                    try:
                        row = self.get_line()
                        # ToDo: validate the line
                        batch.append(row)
                    except NeedMoreData:
                        need_more_data = True
                        break
                else:
                    yield batch
                    batch = []

        except StopIteration as e:
            # yield any final batch
            if batch:
                yield batch

            # - check the end of the file

            self.teardown()

            # - rebuild indexes
            if self.indexes:
                self.restore_indexes(warning_queue)

    def setup_file(self, warning_queue=None):
        raise NotImplementedError("Subclasses need to implement this")

    def teardown(self):
        pass

    def read_to_queue(
        self,
        work_queue,
        exit_event,
        error_queue,
        warning_queue,
        timing_queue,
        fields=None,
        ignore_signals=True,
        batch_size=None,
    ):
        if (
            ignore_signals
        ):  # ToDo: work out when we are in a worker process automatically
            signal.signal(signal.SIGINT, signal.SIG_IGN)  # workers should ignore these

        if batch_size is None:
            batch_size = utils_common.default_batch_size

        self.start_time = time.time()
        try:
            timePoint = time.time()
            for batch in self.batches(warning_queue=warning_queue):
                timing_queue.put(("reader_work", time.time() - timePoint))
                timePoint = time.time()

                # apply the fields filter
                if fields:
                    for row in batch:
                        for key in [x for x in row.keys() if x not in fields]:
                            del row[key]

                while not exit_event.is_set():
                    try:
                        work_queue.put((self.db, self.table, batch), timeout=0.1)
                        self._rows_read.value += len(batch)
                        break
                    except Full:
                        pass
                else:
                    break
                timing_queue.put(("reader_wait", time.time() - timePoint))
                timePoint = time.time()

        # - report relevant errors
        except Exception as exc:
            default_logger.exception(exc)
            error_queue.put(Error(str(exc), traceback.format_exc(), self.name))
            exit_event.set()
            raise
        finally:
            self.end_time = time.time()


class NeedMoreData(Exception):
    pass


class JsonSourceFile(SourceFile):
    format = "json"

    decoder = json.JSONDecoder()
    json_array = None
    found_first = False

    _buffer_size = JSON_READ_CHUNK_SIZE
    _buffer_str = None
    _buffer_pos = None
    _buffer_end = None

    def fill_buffer(self):
        if self._buffer_str is None:
            self._buffer_str = ""
            self._buffer_pos = 0
            self._buffer_end = 0
        elif self._buffer_pos == 0:
            # double the buffer under the assumption that the documents are too large to fit
            if self._buffer_size == JSON_MAX_BUFFER_SIZE:
                raise Exception(
                    "Error: JSON max buffer size exceeded on file %s (from position %d). Use '--max-document-size' to "
                    "extend your buffer." % (self.name, self.bytes_processed)
                )
            self._buffer_size = min(self._buffer_size * 2, JSON_MAX_BUFFER_SIZE)

        # add more data
        read_target = self._buffer_size - self._buffer_end + self._buffer_pos

        if read_target < 1:
            raise AssertionError("Can not set the read target and full the buffer")

        new_chunk = self._source.read(read_target)

        if len(new_chunk) == 0:
            raise StopIteration()  # file ended

        self._buffer_str = self._buffer_str[self._buffer_pos :] + new_chunk
        self._bytes_read.value += len(new_chunk)

        # reset markers
        self._buffer_pos = 0
        self._buffer_end = len(self._buffer_str) - 1

    def get_line(self):
        """Return a line from the current _buffer_str, or raise NeedMoreData trying"""

        # advance over any whitespace
        self._buffer_pos = json.decoder.WHITESPACE.match(
            self._buffer_str, self._buffer_pos
        ).end()
        if self._buffer_pos >= self._buffer_end:
            raise NeedMoreData()

        # read over a comma if we are not the first item in a json_array
        if (
            self.json_array
            and self.found_first
            and self._buffer_str[self._buffer_pos] == ","
        ):
            self._buffer_pos += 1
            if self._buffer_pos >= self._buffer_end:
                raise NeedMoreData()

        # advance over any post-comma whitespace
        self._buffer_pos = json.decoder.WHITESPACE.match(
            self._buffer_str, self._buffer_pos
        ).end()
        if self._buffer_pos >= self._buffer_end:
            raise NeedMoreData()

        # parse and return an object
        try:
            row, self._buffer_pos = self.decoder.raw_decode(
                self._buffer_str, idx=self._buffer_pos
            )
            self.found_first = True
            return row
        except (ValueError, IndexError):
            raise NeedMoreData()

    def setup_file(self, warning_queue=None):
        # - move to the first record

        # advance through any leading whitespace
        while True:
            self.fill_buffer()
            self._buffer_pos = json.decoder.WHITESPACE.match(self._buffer_str, 0).end()
            if self._buffer_pos == 0:
                break

        # check the first character
        try:
            if self._buffer_str[0] == "[":
                self.json_array = True
                self._buffer_pos = 1
            elif self._buffer_str[0] == "{":
                self.json_array = False
            else:
                raise ValueError(
                    "Error: JSON format not recognized - file does not begin with an object or array"
                )
        except IndexError:
            raise ValueError("Error: JSON file was empty of content")

    def teardown(self):

        # - check the end of the file
        # note: fill_buffer should have guaranteed that we have only the data in the end

        # advance through any leading whitespace
        self._buffer_pos = json.decoder.WHITESPACE.match(
            self._buffer_str, self._buffer_pos
        ).end()

        # check the end of the array if we have it
        if self.json_array:
            if self._buffer_str[self._buffer_pos] != "]":
                snippit = self._buffer_str[self._buffer_pos :]
                extra = (
                    ""
                    if len(snippit) <= 100
                    else " and %d more characters" % (len(snippit) - 100)
                )
                raise ValueError(
                    "Error: JSON array did not end cleanly, rather with: <<%s>>%s"
                    % (snippit[:100], extra)
                )
            self._buffer_pos += 1

        # advance through any trailing whitespace
        self._buffer_pos = json.decoder.WHITESPACE.match(
            self._buffer_str, self._buffer_pos
        ).end()
        snippit = self._buffer_str[self._buffer_pos :]
        if len(snippit) > 0:
            extra = (
                ""
                if len(snippit) <= 100
                else " and %d more characters" % (len(snippit) - 100)
            )
            raise ValueError(
                "Error: extra data after JSON data: <<%s>>%s" % (snippit[:100], extra)
            )


class CsvSourceFile(SourceFile):
    format = "csv"

    no_header_row = False
    custom_header = None

    _reader = None  # instance of csv.reader
    _columns = None  # name of the columns

    def __init__(self, *args, **kwargs):
        if "source_options" in kwargs and isinstance(kwargs["source_options"], dict):
            if "no_header_row" in kwargs["source_options"]:
                self.no_header_row = kwargs["source_options"]["no_header_row"]
            if "custom_header" in kwargs["source_options"]:
                self.custom_header = kwargs["source_options"]["custom_header"]

        super(CsvSourceFile, self).__init__(*args, **kwargs)

    def byte_counter(self):
        """Generator for getting a byte count on a file being used"""

        for line in self._source:
            self._bytes_read.value += len(line)
            if unicode != str:
                yield line.encode(
                    "utf-8"
                )  # Python2.x csv module does not really handle unicode
            else:
                yield line

    def setup_file(self, warning_queue=None):
        # - setup csv.reader with a byte counter wrapper

        self._reader = csv.reader(self.byte_counter())

        # - get the header information for column names

        if not self.no_header_row:
            self._columns = next(self._reader)

        # field names may override fields from the header
        if self.custom_header is not None:
            if not self.no_header_row:
                warning_queue.put(
                    "Ignoring header row on %s: %s" % (self.name, str(self._columns))
                )
            self._columns = self.custom_header
        elif self.no_header_row:
            raise ValueError("Error: No field name information available")

    def get_line(self):
        raw_row = next(self._reader)
        if len(self._columns) != len(raw_row):
            raise Exception(
                "Error: '%s' line %d has an inconsistent number of columns: %s"
                % (self.name, self._reader.line_num, str(raw_row))
            )

        row = {}
        for key, value in zip(
            self._columns, raw_row
        ):  # note: we import all csv fields as strings
            # treat empty fields as no entry rather than empty string
            if value == "":
                continue
            row[key] = value if str == unicode else unicode(value, encoding="utf-8")

        return row


# ==


usage = """rethinkdb import -d DIR [-c HOST:PORT] [--tls-cert FILENAME] [-p] [--password-file FILENAME]
      [--force] [-i (DB | DB.TABLE)] [--clients NUM]
      [--shards NUM_SHARDS] [--replicas NUM_REPLICAS]
  rethinkdb import -f FILE --table DB.TABLE [-c HOST:PORT] [--tls-cert FILENAME] [-p] [--password-file FILENAME]
      [--force] [--clients NUM] [--format (csv | json)] [--pkey PRIMARY_KEY]
      [--shards NUM_SHARDS] [--replicas NUM_REPLICAS]
      [--delimiter CHARACTER] [--custom-header FIELD,FIELD... [--no-header]]"""

help_epilog = """
EXAMPLES:

rethinkdb import -d rdb_export -c mnemosyne:39500 --clients 128
  Import data into a cluster running on host 'mnemosyne' with a client port at 39500,
  using 128 client connections and the named export directory.

rethinkdb import -f site_history.csv --format csv --table test.history --pkey count
  Import data into a local cluster and the table 'history' in the 'test' database,
  using the named CSV file, and using the 'count' field as the primary key.

rethinkdb import -d rdb_export -c hades -p -i test
  Import data into a cluster running on host 'hades' which requires a password,
  using only the database 'test' from the named export directory.

rethinkdb import -f subscriber_info.json --fields id,name,hashtag --force
  Import data into a local cluster using the named JSON file, and only the fields
  'id', 'name', and 'hashtag', overwriting any existing rows with the same primary key.

rethinkdb import -f user_data.csv --delimiter ';' --no-header --custom-header id,name,number
  Import data into a local cluster using the named CSV file with no header and instead
  use the fields 'id', 'name', and 'number', the delimiter is a semicolon (rather than
  a comma).
"""


def parse_options(argv, prog=None):
    parser = utils_common.CommonOptionsParser(
        usage=usage, epilog=help_epilog, prog=prog
    )

    parser.add_option(
        "--clients",
        dest="clients",
        metavar="CLIENTS",
        default=8,
        help="client connections to use (default: 8)",
        type="pos_int",
    )
    parser.add_option(
        "--hard-durability",
        dest="durability",
        action="store_const",
        default="soft",
        help="use hard durability writes (slower, uses less memory)",
        const="hard",
    )
    parser.add_option(
        "--force",
        dest="force",
        action="store_true",
        default=False,
        help="import even if a table already exists, overwriting duplicate primary keys",
    )

    parser.add_option(
        "--batch-size",
        dest="batch_size",
        default=utils_common.default_batch_size,
        help=optparse.SUPPRESS_HELP,
        type="pos_int",
    )

    # Replication settings
    replication_options_group = optparse.OptionGroup(parser, "Replication Options")
    replication_options_group.add_option(
        "--shards",
        dest="create_args",
        metavar="SHARDS",
        help="shards to setup on created tables (default: 1)",
        type="pos_int",
        action="add_key",
    )
    replication_options_group.add_option(
        "--replicas",
        dest="create_args",
        metavar="REPLICAS",
        help="replicas to setup on created tables (default: 1)",
        type="pos_int",
        action="add_key",
    )
    parser.add_option_group(replication_options_group)

    # Directory import options
    dir_import_group = optparse.OptionGroup(parser, "Directory Import Options")
    dir_import_group.add_option(
        "-d",
        "--directory",
        dest="directory",
        metavar="DIRECTORY",
        default=None,
        help="directory to import data from",
    )
    dir_import_group.add_option(
        "-i",
        "--import",
        dest="db_tables",
        metavar="DB|DB.TABLE",
        default=[],
        help="restore only the given database or table (may be specified multiple times)",
        action="append",
        type="db_table",
    )
    dir_import_group.add_option(
        "--no-secondary-indexes",
        dest="indexes",
        action="store_false",
        default=None,
        help="do not create secondary indexes",
    )
    parser.add_option_group(dir_import_group)

    # File import options
    file_import_group = optparse.OptionGroup(parser, "File Import Options")
    file_import_group.add_option(
        "-f",
        "--file",
        dest="file",
        metavar="FILE",
        default=None,
        help="file to import data from",
        type="file",
    )
    file_import_group.add_option(
        "--table",
        dest="import_table",
        metavar="DB.TABLE",
        default=None,
        help="table to import the data into",
    )
    file_import_group.add_option(
        "--fields",
        dest="fields",
        metavar="FIELD,...",
        default=None,
        help="limit which fields to use when importing one table",
    )
    file_import_group.add_option(
        "--format",
        dest="format",
        metavar="json|csv",
        default=None,
        help="format of the file (default: json, accepts newline delimited json)",
        type="choice",
        choices=["json", "csv"],
    )
    file_import_group.add_option(
        "--pkey",
        dest="create_args",
        metavar="PRIMARY_KEY",
        default=None,
        help="field to use as the primary key in the table",
        action="add_key",
    )
    parser.add_option_group(file_import_group)

    # CSV import options
    csv_import_group = optparse.OptionGroup(parser, "CSV Options")
    csv_import_group.add_option(
        "--delimiter",
        dest="delimiter",
        metavar="CHARACTER",
        default=None,
        help="character separating fields, or '\\t' for tab",
    )
    csv_import_group.add_option(
        "--no-header",
        dest="no_header",
        action="store_true",
        default=None,
        help="do not read in a header of field names",
    )
    csv_import_group.add_option(
        "--custom-header",
        dest="custom_header",
        metavar="FIELD,...",
        default=None,
        help="header to use (overriding file header), must be specified if --no-header",
    )
    parser.add_option_group(csv_import_group)

    # JSON import options
    json_options_group = optparse.OptionGroup(parser, "JSON Options")
    json_options_group.add_option(
        "--max-document-size",
        dest="max_document_size",
        metavar="MAX_SIZE",
        default=0,
        help="maximum allowed size (bytes) for a single JSON document (default: 128MiB)",
        type="pos_int",
    )
    json_options_group.add_option(
        "--max-nesting-depth",
        dest="max_nesting_depth",
        metavar="MAX_DEPTH",
        default=0,
        help="maximum depth of the JSON documents (default: 100)",
        type="pos_int",
    )
    parser.add_option_group(json_options_group)

    options, args = parser.parse_args(argv)

    # Check validity of arguments

    if len(args) != 0:
        raise parser.error(
            "No positional arguments supported. Unrecognized option(s): %s" % args
        )

    # - create_args
    if options.create_args is None:
        options.create_args = {}

    # - options based on file/directory import

    if options.directory and options.file:
        parser.error("-f/--file and -d/--directory can not be used together")

    elif options.directory:
        if not os.path.exists(options.directory):
            parser.error("-d/--directory does not exist: %s" % options.directory)
        if not os.path.isdir(options.directory):
            parser.error("-d/--directory is not a directory: %s" % options.directory)
        options.directory = os.path.realpath(options.directory)

        # disallow invalid options
        if options.import_table:
            parser.error("--table option is not valid when importing a directory")
        if options.fields:
            parser.error("--fields option is not valid when importing a directory")
        if options.format:
            parser.error("--format option is not valid when importing a directory")
        if options.create_args:
            parser.error("--pkey option is not valid when importing a directory")

        if options.delimiter:
            parser.error("--delimiter option is not valid when importing a directory")
        if options.no_header:
            parser.error("--no-header option is not valid when importing a directory")
        if options.custom_header:
            parser.error(
                "table create options are not valid when importing a directory: %s"
                % ", ".join(
                    [x.lower().replace("_", " ") for x in options.custom_header.keys()]
                )
            )

        # check valid options
        if not os.path.isdir(options.directory):
            parser.error("Directory to import does not exist: %s" % options.directory)

        if options.fields and (
            len(options.db_tables) > 1 or options.db_tables[0].table is None
        ):
            parser.error(
                "--fields option can only be used when importing a single table"
            )

    elif options.file:
        if not os.path.exists(options.file):
            parser.error("-f/--file does not exist: %s" % options.file)
        if not os.path.isfile(options.file):
            parser.error("-f/--file is not a file: %s" % options.file)
        options.file = os.path.realpath(options.file)

        # format
        if options.format is None:
            options.format = os.path.splitext(options.file)[1].lstrip(".")

        # import_table
        if options.import_table:
            res = utils_common._tableNameRegex.match(options.import_table)
            if res and res.group("table"):
                options.import_table = utils_common.DbTable(
                    res.group("db"), res.group("table")
                )
            else:
                parser.error("Invalid --table option: %s" % options.import_table)
        else:
            parser.error("A value is required for --table when importing from a file")

        # fields
        options.fields = options.fields.split(",") if options.fields else None

        # disallow invalid options
        if options.db_tables:
            parser.error("-i/--import can only be used when importing a directory")
        if options.indexes:
            parser.error(
                "--no-secondary-indexes can only be used when importing a directory"
            )

        if options.format == "csv":
            # disallow invalid options
            if options.max_document_size:
                parser.error(
                    "--max_document_size only affects importing JSON documents"
                )

            # delimiter
            if options.delimiter is None:
                options.delimiter = ","
            elif options.delimiter == "\\t":
                options.delimiter = "\t"
            elif len(options.delimiter) != 1:
                parser.error(
                    "Specify exactly one character for the --delimiter option: %s"
                    % options.delimiter
                )

            # no_header
            if options.no_header is None:
                options.no_header = False
            elif options.custom_header is None:
                parser.error("--custom-header is required if --no-header is specified")

            # custom_header
            if options.custom_header:
                options.custom_header = options.custom_header.split(",")

        elif options.format == "json":
            # disallow invalid options
            if options.delimiter is not None:
                parser.error("--delimiter option is not valid for json files")
            if options.no_header:
                parser.error("--no-header option is not valid for json files")
            if options.custom_header is not None:
                parser.error("--custom-header option is not valid for json files")

            # default options
            options.format = "json"

            if options.max_document_size > 0:
                global JSON_MAX_BUFFER_SIZE
                JSON_MAX_BUFFER_SIZE = options.max_document_size

            options.file = os.path.abspath(options.file)

        else:
            parser.error("Unrecognized file format: %s" % options.format)

    else:
        parser.error("Either -f/--file or -d/--directory is required")

    # --

    # max_nesting_depth
    if options.max_nesting_depth > 0:
        global MAX_NESTING_DEPTH
        MAX_NESTING_DEPTH = options.max_nesting_depth

    # --

    return options


# This is run for each client requested, and accepts tasks from the reader processes


def table_writer(
    tables, options, work_queue, error_queue, warning_queue, exit_event, timing_queue
):
    signal.signal(signal.SIGINT, signal.SIG_IGN)  # workers should ignore these
    db = table = batch = None

    try:
        conflict_action = "replace" if options.force else "error"
        timePoint = time.time()
        while not exit_event.is_set():
            # get a batch
            try:
                db, table, batch = work_queue.get(timeout=0.1)
            except Empty:
                continue
            timing_queue.put(("writer_wait", time.time() - timePoint))
            timePoint = time.time()

            # shut down when appropriate
            if isinstance(batch, StopIteration):
                return

            # find the table we are working on
            table_info = tables[(db, table)]
            tbl = query.db(db).table(table)

            # write the batch to the database
            try:
                res = options.retryQuery(
                    "write batch to %s.%s" % (db, table),
                    tbl.insert(
                        ast.expr(batch, nesting_depth=MAX_NESTING_DEPTH),
                        durability=options.durability,
                        conflict=conflict_action,
                    ),
                )

                if res["errors"] > 0:
                    raise RuntimeError(
                        "Error when importing into table '%s.%s': %s"
                        % (db, table, res["first_error"])
                    )
                modified = res["inserted"] + res["replaced"] + res["unchanged"]
                if modified != len(batch):
                    raise RuntimeError(
                        "The inserted/replaced/unchanged number did not match when importing into table '%s.%s': %s"
                        % (db, table, res["first_error"])
                    )

                table_info.add_rows_written(modified)

            except errors.ReqlError:
                # the error might have been caused by a comm or temporary error causing a partial batch write

                for row in batch:
                    if table_info.primary_key not in row:
                        raise RuntimeError(
                            "Connection error while importing.  Current row does not have the specified primary key "
                            "(%s), so cannot guarantee absence of duplicates"
                            % table_info.primary_key
                        )
                    res = None
                    if conflict_action == "replace":
                        res = options.retryQuery(
                            "write row to %s.%s" % (db, table),
                            tbl.insert(
                                ast.expr(row, nesting_depth=MAX_NESTING_DEPTH),
                                durability=options.durability,
                                conflict=conflict_action,
                                ignore_write_hook=True,
                            ),
                        )
                    else:
                        existingRow = options.retryQuery(
                            "read row from %s.%s" % (db, table),
                            tbl.get(row[table_info.primary_key]),
                        )
                        if not existingRow:
                            res = options.retryQuery(
                                "write row to %s.%s" % (db, table),
                                tbl.insert(
                                    ast.expr(row, nesting_depth=MAX_NESTING_DEPTH),
                                    durability=options.durability,
                                    conflict=conflict_action,
                                    ignore_write_hook=True,
                                ),
                            )
                        elif existingRow != row:
                            raise RuntimeError(
                                "Duplicate primary key `%s`:\n%s\n%s"
                                % (table_info.primary_key, str(row), str(existingRow))
                            )

                    if res["errors"] > 0:
                        raise RuntimeError(
                            "Error when importing into table '%s.%s': %s"
                            % (db, table, res["first_error"])
                        )
                    if res["inserted"] + res["replaced"] + res["unchanged"] != 1:
                        raise RuntimeError(
                            "The inserted/replaced/unchanged number was not 1 when inserting on '%s.%s': %s"
                            % (db, table, res)
                        )
                    table_info.add_rows_written(1)
            timing_queue.put(("writer_work", time.time() - timePoint))
            timePoint = time.time()

    except Exception as e:
        error_queue.put(Error(str(e), traceback.format_exc(), "%s.%s" % (db, table)))
        exit_event.set()


def update_progress(tables, debug, exit_event, sleep=0.2):
    signal.signal(signal.SIGINT, signal.SIG_IGN)  # workers should not get these

    # give weights to each of the tables based on file size
    totalSize = sum([x.bytes_size for x in tables])
    for table in tables:
        table.weight = float(table.bytes_size) / totalSize

    lastComplete = None
    startTime = time.time()
    readWrites = collections.deque(maxlen=5)  # (time, read, write)
    readWrites.append((startTime, 0, 0))
    readRate = None
    writeRate = None
    while True:
        try:
            if exit_event.is_set():
                break
            complete = read = write = 0
            currentTime = time.time()
            for table in tables:
                complete += table.percent_done * table.weight
                if debug:
                    read += table.rows_read
                    write += table.rows_written
            readWrites.append((currentTime, read, write))
            if complete != lastComplete:
                timeDelta = readWrites[-1][0] - readWrites[0][0]
                if debug and len(readWrites) > 1 and timeDelta > 0:
                    readRate = max(
                        (readWrites[-1][1] - readWrites[0][1]) / timeDelta, 0
                    )
                    writeRate = max(
                        (readWrites[-1][2] - readWrites[0][2]) / timeDelta, 0
                    )
                utils_common.print_progress(
                    complete, indent=2, read=readRate, write=writeRate
                )
                lastComplete = complete
            time.sleep(sleep)
        except KeyboardInterrupt:
            break
        except Exception as e:
            if debug:
                print(e)
                traceback.print_exc()


def import_tables(options, sources, files_ignored=None):
    # Make sure this isn't a pre-`reql_admin` cluster - which could result in data loss
    # if the user has a database named 'rethinkdb'
    utils_common.check_minimum_version(options, "1.6")

    start_time = time.time()

    tables = dict(((x.db, x.table), x) for x in sources)  # (db, table) => table

    if six.PY3:
        ctx = multiprocessing.get_context(multiprocessing.get_start_method())
        error_queue = SimpleQueue(ctx=ctx)
        warning_queue = SimpleQueue(ctx=ctx)
        timing_queue = SimpleQueue(ctx=ctx)
    else:
        error_queue = SimpleQueue()
        warning_queue = SimpleQueue()
        timing_queue = SimpleQueue()

    max_queue_size = options.clients * 3
    work_queue = multiprocessing.Manager().Queue(max_queue_size)

    exit_event = multiprocessing.Event()
    interrupt_event = multiprocessing.Event()

    errors = []
    warnings = []
    timing_sums = {}

    pools = []
    progress_bar = None
    progress_bar_sleep = 0.2

    # - setup KeyboardInterupt handler
    signal.signal(signal.SIGINT, lambda a, b: utils_common.abort(pools, exit_event))

    # - queue draining
    def drain_queues():
        # error_queue
        while not error_queue.empty():
            errors.append(error_queue.get())

        # warning_queue
        while not warning_queue.empty():
            warnings.append(warning_queue.get())

        # timing_queue
        while not timing_queue.empty():
            key, value = timing_queue.get()
            if key not in timing_sums:
                timing_sums[key] = value
            else:
                timing_sums[key] += value

    # - setup dbs and tables

    # create missing dbs
    needed_dbs = set([x.db for x in sources])
    if "rethinkdb" in needed_dbs:
        raise RuntimeError(
            "Error: Cannot import tables into the system database: 'rethinkdb'"
        )
    options.retryQuery(
        "ensure dbs: %s" % ", ".join(needed_dbs),
        ast.expr(needed_dbs)
        .set_difference(query.db_list())
        .for_each(query.db_create(query.row)),
    )

    # check for existing tables, or if --force is enabled ones with mis-matched primary keys
    existing_tables = dict(
        [
            ((x["db"], x["name"]), x["primary_key"])
            for x in options.retryQuery(
                "list tables",
                query.db("rethinkdb")
                .table("table_config")
                .pluck(["db", "name", "primary_key"]),
            )
        ]
    )
    already_exist = []
    for source in sources:
        if (source.db, source.table) in existing_tables:
            if not options.force:
                already_exist.append("%s.%s" % (source.db, source.table))
            elif source.primary_key is None:
                source.primary_key = existing_tables[(source.db, source.table)]
            elif source.primary_key != existing_tables[(source.db, source.table)]:
                raise RuntimeError(
                    "Error: Table '%s.%s' already exists with a different primary key: %s (expected: %s)"
                    % (
                        source.db,
                        source.table,
                        existing_tables[(source.db, source.table)],
                        source.primary_key,
                    )
                )

    if len(already_exist) == 1:
        raise RuntimeError(
            "Error: Table '%s' already exists, run with --force to import into the existing table"
            % already_exist[0]
        )
    elif len(already_exist) > 1:
        already_exist.sort()
        raise RuntimeError(
            "Error: The following tables already exist, run with --force to import into the existing tables:\n  %s"
            % "\n  ".join(already_exist)
        )

    # - start the import

    try:
        # - start the progress bar
        if not options.quiet:
            progress_bar = multiprocessing.Process(
                target=update_progress,
                name="progress bar",
                args=(sources, options.debug, exit_event, progress_bar_sleep),
            )
            progress_bar.start()
            pools.append([progress_bar])

        # - start the writers
        writers = []
        pools.append(writers)
        for i in range(options.clients):
            writer = multiprocessing.Process(
                target=table_writer,
                name="table writer %d" % i,
                kwargs={
                    "tables": tables,
                    "options": options,
                    "work_queue": work_queue,
                    "error_queue": error_queue,
                    "warning_queue": warning_queue,
                    "timing_queue": timing_queue,
                    "exit_event": exit_event,
                },
            )
            writers.append(writer)
            writer.start()

        # - read the tables options.clients at a time
        readers = []
        pools.append(readers)
        file_iter = iter(sources)
        try:
            while not exit_event.is_set():
                # add a workers to fill up the readers pool
                while len(readers) < options.clients:
                    table = next(file_iter)
                    reader = multiprocessing.Process(
                        target=table.read_to_queue,
                        name="table reader %s.%s" % (table.db, table.table),
                        kwargs={
                            "fields": options.fields,
                            "batch_size": options.batch_size,
                            "work_queue": work_queue,
                            "error_queue": error_queue,
                            "warning_queue": warning_queue,
                            "timing_queue": timing_queue,
                            "exit_event": exit_event,
                        },
                    )
                    readers.append(reader)
                    reader.start()

                # drain the queues
                drain_queues()

                # reap completed tasks
                for reader in readers[:]:
                    if not reader.is_alive():
                        readers.remove(reader)
                    if len(readers) == options.clients:
                        time.sleep(0.05)
        except StopIteration:
            pass  # ran out of new tables

        # - wait for the last batch of readers to complete
        while readers:
            # drain the queues
            drain_queues()

            # drain the work queue to prevent readers from stalling on exit
            if exit_event.is_set():
                try:
                    while True:
                        work_queue.get(timeout=0.1)
                except Empty:
                    pass

            # watch the readers
            for reader in readers[:]:
                try:
                    reader.join(0.1)
                except Exception as exc:
                    default_logger.exception(exc)
                if not reader.is_alive():
                    readers.remove(reader)

        # - append enough StopIterations to signal all writers
        for _ in writers:
            while True:
                if exit_event.is_set():
                    break
                try:
                    work_queue.put((None, None, StopIteration()), timeout=0.1)
                    break
                except Full:
                    pass

        # - wait for all of the writers
        for writer in writers[:]:
            while writer.is_alive():
                writer.join(0.1)
            writers.remove(writer)

        # - stop the progress bar
        if progress_bar:
            progress_bar.join(progress_bar_sleep * 2)
            if not interrupt_event.is_set():
                utils_common.print_progress(1, indent=2)
            if progress_bar.is_alive():
                progress_bar.terminate()

        # - drain queues
        drain_queues()

        # - final reporting
        if not options.quiet:
            # if successful, make sure 100% progress is reported
            if len(errors) == 0 and not interrupt_event.is_set():
                utils_common.print_progress(1.0, indent=2)

            # advance past the progress bar
            print("")

            # report statistics
            def plural(num, text):
                return "%d %s%s" % (num, text, "" if num == 1 else "s")

            print(
                "  %s imported to %s in %.2f secs"
                % (
                    plural(sum(x.rows_written for x in sources), "row"),
                    plural(len(sources), "table"),
                    time.time() - start_time,
                )
            )

            # report debug statistics
            if options.debug:
                print("Debug timing:")
                for key, value in sorted(timing_sums.items(), key=lambda x: x[0]):
                    print("  %s: %.2f" % (key, value))
    finally:
        signal.signal(signal.SIGINT, signal.SIG_DFL)

    drain_queues()

    for error in errors:
        print("%s" % error.message, file=sys.stderr)
        if options.debug and error.traceback:
            print("  Traceback:\n%s" % error.traceback, file=sys.stderr)
        if len(error.file) == 4:
            print("  In file: %s" % error.file, file=sys.stderr)

    for warning in warnings:
        print("%s" % warning[1], file=sys.stderr)
        if options.debug:
            print(
                "%s traceback: %s" % (warning[0].__name__, warning[2]), file=sys.stderr
            )
        if len(warning) == 4:
            print("In file: %s" % warning[3], file=sys.stderr)

    if interrupt_event.is_set():
        raise RuntimeError("Interrupted")
    if errors:
        raise RuntimeError("Errors occurred during import")
    if warnings:
        raise RuntimeError("Warnings occurred during import")


def parse_sources(options, files_ignored=None):
    def parse_info_file(path):
        primary_key = None
        indexes = []
        write_hook = None
        with open(path, "r") as info_file:
            metadata = json.load(info_file)
            if "primary_key" in metadata:
                primary_key = metadata["primary_key"]
            if "indexes" in metadata and options.indexes is not False:
                indexes = metadata["indexes"]
            if "write_hook" in metadata:
                write_hook = metadata["write_hook"]
        return primary_key, indexes, write_hook

    has_write_hooks = utils_common.check_minimum_version(options, "2.3.7", False)

    sources = set()
    if files_ignored is None:
        files_ignored = []
    if options.directory and options.file:
        raise RuntimeError(
            "Error: Both --directory and --file cannot be specified together"
        )
    elif options.file:
        db, table = options.import_table
        path, ext = os.path.splitext(options.file)
        table_type_options = None
        if ext == ".json":
            table_type = JsonSourceFile
        elif ext == ".csv":
            table_type = CsvSourceFile
            table_type_options = {
                "no_header_row": options.no_header,
                "custom_header": options.custom_header,
            }
        else:
            raise Exception("The table type is not recognised: %s" % ext)

        # - parse the info file if it exists
        primary_key = (
            options.create_args.get("primary_key", None)
            if options.create_args
            else None
        )
        indexes = []
        write_hook = None
        info_path = path + ".info"
        if (primary_key is None or options.indexes is not False) and os.path.isfile(
            info_path
        ):
            info_primary_key, info_indexes, info_write_hook = parse_info_file(info_path)
            if primary_key is None:
                primary_key = info_primary_key
            if options.indexes is not False:
                indexes = info_indexes
            if write_hook is None:
                write_hook = info_write_hook
            if write_hook and not has_write_hooks:
                raise Exception("this RDB version doesn't support write-hooks")

        sources.add(
            table_type(
                source=options.file,
                db=db,
                table=table,
                query_runner=options.retryQuery,
                primary_key=primary_key,
                indexes=indexes,
                write_hook=write_hook,
                source_options=table_type_options,
            )
        )
    elif options.directory:
        # Scan for all files, make sure no duplicated tables with different formats
        dbs = False
        files_ignored = []
        for root, dirs, files in os.walk(options.directory):
            if not dbs:
                files_ignored.extend([os.path.join(root, f) for f in files])
                # The first iteration through should be the top-level directory, which contains the db folders
                dbs = True

                # don't recurse into folders not matching our filter
                db_filter = set([db_table[0] for db_table in options.db_tables or []])
                if db_filter:
                    for dir_name in dirs[:]:  # iterate on a copy
                        if dir_name not in db_filter:
                            dirs.remove(dir_name)
            else:
                if dirs:
                    files_ignored.extend([os.path.join(root, d) for d in dirs])
                    del dirs[:]

                db = os.path.basename(root)
                for filename in files:
                    path = os.path.join(root, filename)
                    table, ext = os.path.splitext(filename)
                    table = os.path.basename(table)

                    if ext not in [".json", ".csv", ".info"]:
                        files_ignored.append(os.path.join(root, filename))
                    elif ext == ".info":
                        pass  # Info files are included based on the data files
                    elif not os.path.exists(os.path.join(root, table + ".info")):
                        files_ignored.append(os.path.join(root, filename))
                    else:
                        # apply db/table filters
                        if options.db_tables:
                            for filter_db, filter_table in options.db_tables:
                                if db == filter_db and filter_table in (None, table):
                                    break  # either all tables in this db, or specific pair
                            else:
                                files_ignored.append(os.path.join(root, filename))
                                continue  # not a chosen db/table

                        # collect the info
                        primary_key = None
                        indexes = []
                        write_hook = None
                        info_path = os.path.join(root, table + ".info")
                        if not os.path.isfile(info_path):
                            files_ignored.append(os.path.join(root, filename))
                        else:
                            primary_key, indexes, write_hook = parse_info_file(
                                info_path
                            )
                        if write_hook and not has_write_hooks:
                            raise Exception(
                                "RDB versions below doesn't support write-hooks"
                            )

                        table_type = None
                        if ext == ".json":
                            table_type = JsonSourceFile
                        elif ext == ".csv":
                            table_type = CsvSourceFile
                        else:
                            raise Exception(
                                "The table type is not recognised: %s" % ext
                            )
                        source = table_type(
                            source=path,
                            query_runner=options.retryQuery,
                            db=db,
                            table=table,
                            primary_key=primary_key,
                            indexes=indexes,
                            write_hook=write_hook,
                        )

                        # ensure we don't have a duplicate
                        if table in sources:
                            raise RuntimeError(
                                "Error: Duplicate db.table found in directory tree: %s.%s"
                                % (source.db, source.table)
                            )

                        sources.add(source)

        # Warn the user about the files that were ignored
        if len(files_ignored) > 0:
            print(
                "Unexpected files found in the specified directory.  Importing a directory expects",
                file=sys.stderr,
            )
            print(
                " a directory from `rethinkdb export`.  If you want to import individual tables",
                file=sys.stderr,
            )
            print(
                " import them as single files.  The following files were ignored:",
                file=sys.stderr,
            )
            for ignored_file in files_ignored:
                print("%s" % str(ignored_file), file=sys.stderr)
    else:
        raise RuntimeError("Error: Neither --directory or --file specified")

    return sources


def main(argv=None, prog=None):
    start_time = time.time()

    if argv is None:
        argv = sys.argv[1:]
    options = parse_options(argv, prog=prog)

    try:
        sources = parse_sources(options)
        import_tables(options, sources)
    except RuntimeError as ex:
        print(ex, file=sys.stderr)
        if str(ex) == "Warnings occurred during import":
            return 2
        return 1
    if not options.quiet:
        print("  Done (%d seconds)" % (time.time() - start_time))
    return 0


if __name__ == "__main__":
    sys.exit(main())