#!/usr/bin/env python

#

# Copyright 2007 Google Inc.

#

# 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.

#

"""Imports data over HTTP.

Usage:

  %(arg0)s [flags]

    --debug                 Show debugging information. (Optional)

    --app_id=<string>       Application ID of endpoint (Optional for

                            *.appspot.com)

    --auth_domain=<domain>  The auth domain to use for logging in and for

                            UserProperties. (Default: gmail.com)

    --bandwidth_limit=<int> The maximum number of bytes per second for the

                            aggregate transfer of data to the server. Bursts

                            may exceed this, but overall transfer rate is

                            restricted to this rate. (Default 250000)

    --batch_size=<int>      Number of Entity objects to include in each post to

                            the URL endpoint. The more data per row/Entity, the

                            smaller the batch size should be. (Default 10)

    --config_file=<path>    File containing Model and Loader definitions.

                            (Required unless --dump or --restore are used)

    --db_filename=<path>    Specific progress database to write to, or to

                            resume from. If not supplied, then a new database

                            will be started, named:

                            bulkloader-progress-TIMESTAMP.

                            The special filename "skip" may be used to simply

                            skip reading/writing any progress information.

    --download              Export entities to a file.

    --dry_run               Do not execute any remote_api calls.

    --dump                  Use zero-configuration dump format.

    --email=<string>        The username to use. Will prompt if omitted.

    --exporter_opts=<string>

                            A string to pass to the Exporter.initialize method.

    --filename=<path>       Path to the file to import. (Required)

    --has_header            Skip the first row of the input.

    --http_limit=<int>      The maximum numer of HTTP requests per second to

                            send to the server. (Default: 8)

    --kind=<string>         Name of the Entity object kind to put in the

                            datastore. (Required)

    --loader_opts=<string>  A string to pass to the Loader.initialize method.

    --log_file=<path>       File to write bulkloader logs.  If not supplied

                            then a new log file will be created, named:

                            bulkloader-log-TIMESTAMP.

    --map                   Map an action across datastore entities.

    --mapper_opts=<string>  A string to pass to the Mapper.Initialize method.

    --num_threads=<int>     Number of threads to use for uploading entities

                            (Default 10)

    --passin                Read the login password from stdin.

    --restore               Restore from zero-configuration dump format.

    --result_db_filename=<path>

                            Result database to write to for downloads.

    --rps_limit=<int>       The maximum number of records per second to

                            transfer to the server. (Default: 20)

    --url=<string>          URL endpoint to post to for importing data.

                            (Required)

The exit status will be 0 on success, non-zero on import failure.

Works with the remote_api mix-in library for google.appengine.ext.remote_api.

Please look there for documentation about how to setup the server side.

Example:

%(arg0)s --url=http://app.appspot.com/remote_api --kind=Model \

 --filename=data.csv --config_file=loader_config.py

"""

import csv

import errno

import getopt

import getpass

import imp

import logging

import os

import Queue

import re

import shutil

import signal

import StringIO

import sys

import threading

import time

import traceback

import urllib2

import urlparse

from google.appengine.datastore import entity_pb

from google.appengine.api import apiproxy_stub_map

from google.appengine.api import datastore

from google.appengine.api import datastore_errors

from google.appengine.datastore import datastore_pb

from google.appengine.ext import db

from google.appengine.ext import key_range as key_range_module

from google.appengine.ext.db import polymodel

from google.appengine.ext.remote_api import remote_api_stub

from google.appengine.ext.remote_api import throttle as remote_api_throttle

from google.appengine.runtime import apiproxy_errors

from google.appengine.tools import adaptive_thread_pool

from google.appengine.tools import appengine_rpc

from google.appengine.tools.requeue import ReQueue

try:

  import sqlite3

except ImportError:

  pass

logger = logging.getLogger('google.appengine.tools.bulkloader')

KeyRange = key_range_module.KeyRange

DEFAULT_THREAD_COUNT = 10

DEFAULT_BATCH_SIZE = 10

DEFAULT_DOWNLOAD_BATCH_SIZE = 100

DEFAULT_QUEUE_SIZE = DEFAULT_THREAD_COUNT * 10

_THREAD_SHOULD_EXIT = '_THREAD_SHOULD_EXIT'

STATE_READ = 0

STATE_SENDING = 1

STATE_SENT = 2

STATE_NOT_SENT = 3

STATE_GETTING = 1

STATE_GOT = 2

STATE_ERROR = 3

DATA_CONSUMED_TO_HERE = 'DATA_CONSUMED_TO_HERE'

INITIAL_BACKOFF = 1.0

BACKOFF_FACTOR = 2.0

DEFAULT_BANDWIDTH_LIMIT = 250000

DEFAULT_RPS_LIMIT = 20

DEFAULT_REQUEST_LIMIT = 8

MAXIMUM_INCREASE_DURATION = 5.0

MAXIMUM_HOLD_DURATION = 12.0

def ImportStateMessage(state):

  """Converts a numeric state identifier to a status message."""

  return ({

      STATE_READ: 'Batch read from file.',

      STATE_SENDING: 'Sending batch to server.',

      STATE_SENT: 'Batch successfully sent.',

      STATE_NOT_SENT: 'Error while sending batch.'

  }[state])

def ExportStateMessage(state):

  """Converts a numeric state identifier to a status message."""

  return ({

      STATE_READ: 'Batch read from file.',

      STATE_GETTING: 'Fetching batch from server',

      STATE_GOT: 'Batch successfully fetched.',

      STATE_ERROR: 'Error while fetching batch'

  }[state])

def MapStateMessage(state):

  """Converts a numeric state identifier to a status message."""

  return ({

      STATE_READ: 'Batch read from file.',

      STATE_GETTING: 'Querying for batch from server',

      STATE_GOT: 'Batch successfully fetched.',

      STATE_ERROR: 'Error while fetching or mapping.'

  }[state])

def ExportStateName(state):

  """Converts a numeric state identifier to a string."""

  return ({

      STATE_READ: 'READ',

      STATE_GETTING: 'GETTING',

      STATE_GOT: 'GOT',

      STATE_ERROR: 'NOT_GOT'

  }[state])

def ImportStateName(state):

  """Converts a numeric state identifier to a string."""

  return ({

      STATE_READ: 'READ',

      STATE_GETTING: 'SENDING',

      STATE_GOT: 'SENT',

      STATE_NOT_SENT: 'NOT_SENT'

  }[state])

class Error(Exception):

  """Base-class for exceptions in this module."""

class MissingPropertyError(Error):

  """An expected field is missing from an entity, and no default was given."""

class FatalServerError(Error):

  """An unrecoverable error occurred while posting data to the server."""

class ResumeError(Error):

  """Error while trying to resume a partial upload."""

class ConfigurationError(Error):

  """Error in configuration options."""

class AuthenticationError(Error):

  """Error while trying to authenticate with the server."""

class FileNotFoundError(Error):

  """A filename passed in by the user refers to a non-existent input file."""

class FileNotReadableError(Error):

  """A filename passed in by the user refers to a non-readable input file."""

class FileExistsError(Error):

  """A filename passed in by the user refers to an existing output file."""

class FileNotWritableError(Error):

  """A filename passed in by the user refers to a non-writable output file."""

class BadStateError(Error):

  """A work item in an unexpected state was encountered."""

class KeyRangeError(Error):

  """An error during construction of a KeyRangeItem."""

class FieldSizeLimitError(Error):

  """The csv module tried to read a field larger than the size limit."""

  def __init__(self, limit):

    self.message = """

A field in your CSV input file has exceeded the current limit of %d.

You can raise this limit by adding the following lines to your config file:

import csv

csv.field_size_limit(new_limit)

where new_limit is number larger than the size in bytes of the largest

field in your CSV.

""" % limit

    Error.__init__(self, self.message)

class NameClashError(Error):

  """A name clash occurred while trying to alias old method names."""

  def __init__(self, old_name, new_name, klass):

    Error.__init__(self, old_name, new_name, klass)

    self.old_name = old_name

    self.new_name = new_name

    self.klass = klass

def GetCSVGeneratorFactory(kind, csv_filename, batch_size, csv_has_header,

                           openfile=open, create_csv_reader=csv.reader):

  """Return a factory that creates a CSV-based UploadWorkItem generator.

  Args:

    kind: The kind of the entities being uploaded.

    csv_filename: File on disk containing CSV data.

    batch_size: Maximum number of CSV rows to stash into an UploadWorkItem.

    csv_has_header: Whether to skip the first row of the CSV.

    openfile: Used for dependency injection.

    create_csv_reader: Used for dependency injection.

  Returns:

    A callable (accepting the Progress Queue and Progress Generators

    as input) which creates the UploadWorkItem generator.

  """

  loader = Loader.RegisteredLoader(kind)

  loader._Loader__openfile = openfile

  loader._Loader__create_csv_reader = create_csv_reader

  record_generator = loader.generate_records(csv_filename)

  def CreateGenerator(request_manager, progress_queue, progress_generator):

    """Initialize a UploadWorkItem generator.

    Args:

      request_manager: A RequestManager instance.

      progress_queue: A ProgressQueue instance to send progress information.

      progress_generator: A generator of progress information or None.

    Returns:

      An UploadWorkItemGenerator instance.

    """

    return UploadWorkItemGenerator(request_manager,

                                   progress_queue,

                                   progress_generator,

                                   record_generator,

                                   csv_has_header,

                                   batch_size)

  return CreateGenerator

class UploadWorkItemGenerator(object):

  """Reads rows from a row generator and generates UploadWorkItems."""

  def __init__(self,

               request_manager,

               progress_queue,

               progress_generator,

               record_generator,

               skip_first,

               batch_size):

    """Initialize a WorkItemGenerator.

    Args:

      request_manager: A RequestManager instance with which to associate

        WorkItems.

      progress_queue: A progress queue with which to associate WorkItems.

      progress_generator: A generator of progress information.

      record_generator: A generator of data records.

      skip_first: Whether to skip the first data record.

      batch_size: The number of data records per WorkItem.

    """

    self.request_manager = request_manager

    self.progress_queue = progress_queue

    self.progress_generator = progress_generator

    self.reader = record_generator

    self.skip_first = skip_first

    self.batch_size = batch_size

    self.line_number = 1

    self.column_count = None

    self.read_rows = []

    self.row_count = 0

    self.xfer_count = 0

  def _AdvanceTo(self, line):

    """Advance the reader to the given line.

    Args:

      line: A line number to advance to.

    """

    while self.line_number < line:

      self.reader.next()

      self.line_number += 1

      self.row_count += 1

      self.xfer_count += 1

  def _ReadRows(self, key_start, key_end):

    """Attempts to read and encode rows [key_start, key_end].

    The encoded rows are stored in self.read_rows.

    Args:

      key_start: The starting line number.

      key_end: The ending line number.

    Raises:

      StopIteration: if the reader runs out of rows

      ResumeError: if there are an inconsistent number of columns.

    """

    assert self.line_number == key_start

    self.read_rows = []

    while self.line_number <= key_end:

      row = self.reader.next()

      self.row_count += 1

      if self.column_count is None:

        self.column_count = len(row)

      self.read_rows.append((self.line_number, row))

      self.line_number += 1

  def _MakeItem(self, key_start, key_end, rows, progress_key=None):

    """Makes a UploadWorkItem containing the given rows, with the given keys.

    Args:

      key_start: The start key for the UploadWorkItem.

      key_end: The end key for the UploadWorkItem.

      rows: A list of the rows for the UploadWorkItem.

      progress_key: The progress key for the UploadWorkItem

    Returns:

      An UploadWorkItem instance for the given batch.

    """

    assert rows

    item = UploadWorkItem(self.request_manager, self.progress_queue, rows,

                          key_start, key_end, progress_key=progress_key)

    return item

  def Batches(self):

    """Reads from the record_generator and generates UploadWorkItems.

    Yields:

      Instances of class UploadWorkItem

    Raises:

      ResumeError: If the progress database and data file indicate a different

        number of rows.

    """

    if self.skip_first:

      logger.info('Skipping header line.')

      try:

        self.reader.next()

      except StopIteration:

        return

    exhausted = False

    self.line_number = 1

    self.column_count = None

    logger.info('Starting import; maximum %d entities per post',

                self.batch_size)

    state = None

    if self.progress_generator:

      for progress_key, state, key_start, key_end in self.progress_generator:

        if key_start:

          try:

            self._AdvanceTo(key_start)

            self._ReadRows(key_start, key_end)

            yield self._MakeItem(key_start,

                                 key_end,

                                 self.read_rows,

                                 progress_key=progress_key)

          except StopIteration:

            logger.error('Mismatch between data file and progress database')

            raise ResumeError(

                'Mismatch between data file and progress database')

        elif state == DATA_CONSUMED_TO_HERE:

          try:

            self._AdvanceTo(key_end + 1)

          except StopIteration:

            state = None

    if self.progress_generator is None or state == DATA_CONSUMED_TO_HERE:

      while not exhausted:

        key_start = self.line_number

        key_end = self.line_number + self.batch_size - 1

        try:

          self._ReadRows(key_start, key_end)

        except StopIteration:

          exhausted = True

          key_end = self.line_number - 1

        if key_start <= key_end:

          yield self._MakeItem(key_start, key_end, self.read_rows)

class CSVGenerator(object):

  """Reads a CSV file and generates data records."""

  def __init__(self,

               csv_filename,

               openfile=open,

               create_csv_reader=csv.reader):

    """Initializes a CSV generator.

    Args:

      csv_filename: File on disk containing CSV data.

      openfile: Used for dependency injection of 'open'.

      create_csv_reader: Used for dependency injection of 'csv.reader'.

    """

    self.csv_filename = csv_filename

    self.openfile = openfile

    self.create_csv_reader = create_csv_reader

  def Records(self):

    """Reads the CSV data file and generates row records.

    Yields:

      Lists of strings

    Raises:

      ResumeError: If the progress database and data file indicate a different