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

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

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

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

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

                            (Default 10)

                            may exceed this, but overall transfer rate is

                            restricted to this rate. (Default 250000)

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

import getpass

import logging

import new

import os

import Queue

import signal

import sys

import threading

import time

import traceback

import urllib2

import urlparse

from google.appengine.ext import db

from google.appengine.ext.remote_api import remote_api_stub

from google.appengine.tools import appengine_rpc

try:

  import sqlite3

except ImportError:

  pass

UPLOADER_VERSION = '1'

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

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'

def StateMessage(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])

class Error(Exception):

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

class FatalServerError(Error):

  """An unrecoverable error occurred while trying to post 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."""

def GetCSVGeneratorFactory(csv_filename, batch_size,

                           openfile=open, create_csv_reader=csv.reader):

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

  Args:

    csv_filename: File on disk containing CSV data.

    batch_size: Maximum number of CSV rows to stash into a WorkItem.

    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.

  """

  def CreateGenerator(progress_queue, progress_generator):

    """Initialize a CSV 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 CSVGenerator instance.

    """

    return CSVGenerator(progress_queue,

                        progress_generator,

                        csv_filename,

                        batch_size,

                        openfile,

                        create_csv_reader)

  return CreateGenerator

class CSVGenerator(object):

  """Reads a CSV file and generates WorkItems containing batches of records."""

  def __init__(self,

               progress_queue,

               progress_generator,

               csv_filename,

               batch_size,

               openfile,

               create_csv_reader):

    """Initializes a CSV generator.

    Args:

      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.

      csv_filename: File on disk containing CSV data.

      batch_size: Maximum number of CSV rows to stash into a WorkItem.

      openfile: Used for dependency injection of 'open'.

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

    """

    self.progress_queue = progress_queue

    self.progress_generator = progress_generator

    self.csv_filename = csv_filename

    self.batch_size = batch_size

    self.openfile = openfile

    self.create_csv_reader = create_csv_reader

    self.line_number = 1

    self.column_count = None

    self.read_rows = []

    self.reader = None

    self.row_count = 0

    self.sent_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.sent_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 the CSV data file and generates WorkItems.

    Yields:

      Instances of class WorkItem

    Raises:

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

        number of rows.

    """

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

    csv_content = csv_file.read()

    if csv_content:

      has_headers = csv.Sniffer().has_header(csv_content)

    else:

      has_headers = False

    csv_file.seek(0)

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

    if has_headers:

      logging.info('The CSV file appears to have a header line, skipping.')

      self.reader.next()

    exhausted = False

    self.line_number = 1

    self.column_count = None

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

                 self.batch_size)

    state = None

    if self.progress_generator is not None:

      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:

            logging.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 ReQueue(object):

  """A special thread-safe queue.

  A ReQueue allows unfinished work items to be returned with a call to

  reput().  When an item is reput, task_done() should *not* be called

  in addition, getting an item that has been reput does not increase

  the number of outstanding tasks.

  This class shares an interface with Queue.Queue and provides the

  additional Reput method.

  """

  def __init__(self,

               queue_capacity,

               requeue_capacity=None,

               queue_factory=Queue.Queue,

               get_time=time.time):

    """Initialize a ReQueue instance.

    Args:

      queue_capacity: The number of items that can be put in the ReQueue.

      requeue_capacity: The numer of items that can be reput in the ReQueue.

      queue_factory: Used for dependency injection.

      get_time: Used for dependency injection.

    """

    if requeue_capacity is None:

      requeue_capacity = queue_capacity

    self.get_time = get_time

    self.queue = queue_factory(queue_capacity)

    self.requeue = queue_factory(requeue_capacity)

    self.lock = threading.Lock()

    self.put_cond = threading.Condition(self.lock)

    self.get_cond = threading.Condition(self.lock)

  def _DoWithTimeout(self,

                     action,

                     exc,

                     wait_cond,

                     done_cond,

                     lock,

                     timeout=None,

                     block=True):

    """Performs the given action with a timeout.

    The action must be non-blocking, and raise an instance of exc on a

    recoverable failure.  If the action fails with an instance of exc,

    we wait on wait_cond before trying again.  Failure after the

    timeout is reached is propagated as an exception.  Success is

    signalled by notifying on done_cond and returning the result of

    the action.  If action raises any exception besides an instance of

    exc, it is immediately propagated.

    Args:

      action: A callable that performs a non-blocking action.

      exc: An exception type that is thrown by the action to indicate

        a recoverable error.

      wait_cond: A condition variable which should be waited on when

        action throws exc.

      done_cond: A condition variable to signal if the action returns.

      lock: The lock used by wait_cond and done_cond.

      timeout: A non-negative float indicating the maximum time to wait.

      block: Whether to block if the action cannot complete immediately.

    Returns:

      The result of the action, if it is successful.

    Raises:

      ValueError: If the timeout argument is negative.

    """

    if timeout is not None and timeout < 0.0:

      raise ValueError('\'timeout\' must not be a negative  number')

    if not block:

      timeout = 0.0

    result = None

    success = False

    start_time = self.get_time()

    lock.acquire()

    try:

      while not success:

        try:

          result = action()

          success = True

        except Exception, e:

          if not isinstance(e, exc):

            raise e

          if timeout is not None:

            elapsed_time = self.get_time() - start_time

            timeout -= elapsed_time

            if timeout <= 0.0:

              raise e

          wait_cond.wait(timeout)

    finally:

      if success:

        done_cond.notify()

      lock.release()

    return result

  def put(self, item, block=True, timeout=None):

    """Put an item into the requeue.

    Args:

      item: An item to add to the requeue.

      block: Whether to block if the requeue is full.

      timeout: Maximum on how long to wait until the queue is non-full.

    Raises:

      Queue.Full if the queue is full and the timeout expires.

    """

    def PutAction():

      self.queue.put(item, block=False)

    self._DoWithTimeout(PutAction,

                        Queue.Full,

                        self.get_cond,

                        self.put_cond,

                        self.lock,

                        timeout=timeout,

                        block=block)

  def reput(self, item, block=True, timeout=None):

    """Re-put an item back into the requeue.

    Re-putting an item does not increase the number of outstanding

    tasks, so the reput item should be uniquely associated with an

    item that was previously removed from the requeue and for which

    task_done has not been called.

    Args:

      item: An item to add to the requeue.

      block: Whether to block if the requeue is full.

      timeout: Maximum on how long to wait until the queue is non-full.

    Raises:

      Queue.Full is the queue is full and the timeout expires.

    """

    def ReputAction():

      self.requeue.put(item, block=False)

    self._DoWithTimeout(ReputAction,

                        Queue.Full,

                        self.get_cond,

                        self.put_cond,

                        self.lock,

                        timeout=timeout,

                        block=block)

  def get(self, block=True, timeout=None):

    """Get an item from the requeue.

    Args:

      block: Whether to block if the requeue is empty.

      timeout: Maximum on how long to wait until the requeue is non-empty.

    Returns:

      An item from the requeue.

    Raises:

      Queue.Empty if the queue is empty and the timeout expires.

    """

    def GetAction():

      try:

        result = self.requeue.get(block=False)

        self.requeue.task_done()

      except Queue.Empty:

        result = self.queue.get(block=False)

      return result

    return self._DoWithTimeout(GetAction,

                               Queue.Empty,

                               self.put_cond,

                               self.get_cond,

                               self.lock,

                               timeout=timeout,

                               block=block)

  def join(self):

    """Blocks until all of the items in the requeue have been processed."""

    self.queue.join()