#!/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)

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

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

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

                            (Default 10)

    --passin                Read the login password from stdin.

    --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 cPickle

import csv

import errno

import getopt

import getpass

import imp

import logging

import os

import Queue

import re

import signal

import StringIO

import sys

import threading

import time

import urllib2

import urlparse

from google.appengine.api import datastore_errors

from google.appengine.ext import db

from google.appengine.ext.db import polymodel

from google.appengine.ext.remote_api import remote_api_stub

from google.appengine.runtime import apiproxy_errors

from google.appengine.tools import appengine_rpc

try:

  import sqlite3

except ImportError:

  pass

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

DEFAULT_THREAD_COUNT = 10

DEFAULT_BATCH_SIZE = 10

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_NOT_GOT = 3

MINIMUM_THROTTLE_SLEEP_DURATION = 0.001

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

BANDWIDTH_UP = 'http-bandwidth-up'

BANDWIDTH_DOWN = 'http-bandwidth-down'

REQUESTS = 'http-requests'

HTTPS_BANDWIDTH_UP = 'https-bandwidth-up'

HTTPS_BANDWIDTH_DOWN = 'https-bandwidth-down'

HTTPS_REQUESTS = 'https-requests'

RECORDS = 'records'

MAXIMUM_INCREASE_DURATION = 8.0

MAXIMUM_HOLD_DURATION = 10.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_NOT_GOT: 'Error while fetching batch'

  }[state])

def ExportStateName(state):

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

  return ({

      STATE_READ: 'READ',

      STATE_GETTING: 'GETTING',

      STATE_GOT: 'GOT',

      STATE_NOT_GOT: '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_GOT: '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 KeyRangeError(Error):

  """Error while trying to generate a KeyRange."""

class BadStateError(Error):

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

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 WorkItem 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 a WorkItem.

    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 WorkItem 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(progress_queue, progress_generator):

    """Initialize a WorkItem generator linked to a progress generator and queue.

    Args:

      progress_queue: A ProgressQueue instance to send progress information.

      progress_generator: A generator of progress information or None.

    Returns:

      A WorkItemGenerator instance.

    """

    return WorkItemGenerator(progress_queue,

                             progress_generator,

                             record_generator,

                             csv_has_header,

                             batch_size)

  return CreateGenerator

class WorkItemGenerator(object):

  """Reads rows from a row generator and generates WorkItems of batches."""

  def __init__(self,

               progress_queue,

               progress_generator,

               record_generator,

               skip_first,

               batch_size):

    """Initialize a WorkItemGenerator.

    Args:

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

      else:

        if self.column_count != len(row):

          raise ResumeError('Column count mismatch, %d: %s' %

                            (self.column_count, str(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 WorkItem containing the given rows, with the given keys.

    Args:

      key_start: The start key for the WorkItem.

      key_end: The end key for the WorkItem.

      rows: A list of the rows for the WorkItem.

      progress_key: The progress key for the WorkItem

    Returns:

      A WorkItem instance for the given batch.

    """

    assert rows

    item = WorkItem(self.progress_queue, rows,

                    key_start, key_end,

                    progress_key=progress_key)

    return item

  def Batches(self):

    """Reads from the record_generator and generates WorkItems.

    Yields:

      Instances of class WorkItem

    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

        number of rows.

    """

    csv_file = self.openfile(self.csv_filename, 'rb')

    reader = self.create_csv_reader(csv_file, skipinitialspace=True)

    return reader

class KeyRangeGenerator(object):

  """Generates ranges of keys to download.

  Reads progress information from the progress database and creates

  KeyRange objects corresponding to incompletely downloaded parts of an

  export.

  """

  def __init__(self, kind, progress_queue, progress_generator):

    """Initialize the KeyRangeGenerator.

    Args:

      kind: The kind of entities being transferred.

      progress_queue: A queue used for tracking progress information.

      progress_generator: A generator of prior progress information, or None

        if there is no prior status.

    """

    self.kind = kind

    self.row_count = 0

    self.xfer_count = 0

    self.progress_queue = progress_queue

    self.progress_generator = progress_generator

  def Batches(self):

    """Iterate through saved progress information.

    Yields:

      KeyRange instances corresponding to undownloaded key ranges.

    """

    if self.progress_generator is not None:

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

        if state is not None and state != STATE_GOT and key_start is not None:

          key_start = ParseKey(key_start)

          key_end = ParseKey(key_end)

          result = KeyRange(self.progress_queue,