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

#

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.

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

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

import Queue

import signal

import sys

import threading

import time

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

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'

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

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

                           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.

  """

  def CreateGenerator(progress_queue, progress_generator):

    Args:

      progress_queue: A ProgressQueue instance to send progress information.

      progress_generator: A generator of progress information or None.

    Returns:

    """

  return CreateGenerator

  def __init__(self,

               progress_queue,

               progress_generator,

    Args:

    """

    self.progress_queue = progress_queue

    self.progress_generator = progress_generator

    self.batch_size = batch_size

    self.line_number = 1

    self.column_count = None

    self.read_rows = []

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

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

    Yields:

      Instances of class WorkItem

    Raises:

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

        number of rows.

    """

    exhausted = False

    self.line_number = 1

    self.column_count = None

    state = 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:

            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)