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

  %s [flags]

    --debug             Show debugging information. (Optional)

    --cookie=<string>   Whole Cookie header to supply to the server, including

                        the parameter name (e.g., "ACSID=..."). (Optional)

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

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

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

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

                        (Required)

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

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

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

"""

import StringIO

import httplib

import logging

import csv

import getopt

import socket

import sys

import urllib

import urlparse

from google.appengine.ext.bulkload import constants

class Error(Exception):

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

class PostError(Error):

  """An error has occured while trying to post data to the server."""

class BadServerStatusError(PostError):

  """The server has returned an error while importing data."""

def ContentGenerator(csv_file,

                     batch_size,

                     create_csv_reader=csv.reader,

                     create_csv_writer=csv.writer):

  """Retrieves CSV data up to a batch size at a time.

  Args:

    csv_file: A file-like object for reading CSV data.

    batch_size: Maximum number of CSV rows to yield on each iteration.

    create_csv_reader, create_csv_writer: Used for dependency injection.

  Yields:

    Tuple (entity_count, csv_content) where:

      entity_count: Number of entities contained in the csv_content. Will be

        less than or equal to the batch_size and greater than 0.

      csv_content: String containing the CSV content containing the next

        entity_count entities.

  """

  try:

    csv.field_size_limit(800000)

  except AttributeError:

    pass

  reader = create_csv_reader(csv_file, skipinitialspace=True)

  exhausted = False

  while not exhausted:

    rows_written = 0

    content = StringIO.StringIO()

    writer = create_csv_writer(content)

    try:

      for i in xrange(batch_size):

        row = reader.next()

        writer.writerow(row)

        rows_written += 1

    except StopIteration:

      exhausted = True

    if rows_written > 0:

      yield rows_written, content.getvalue()

def PostEntities(host_port, uri, cookie, kind, content):

  """Posts Entity records to a remote endpoint over HTTP.

  Args:

   host_port: String containing the "host:port" pair; the port is optional.

   uri: Relative URI to access on the remote host (e.g., '/bulkload').

   cookie: String containing the Cookie header to use, if any.

   kind: Kind of the Entity records being posted.

   content: String containing the CSV data for the entities.

  Raises:

    BadServerStatusError if the server was contactable but returns an error.

    PostError If an error occurred while connecting to the server or reading

    or writing data.

  """

  logging.debug('Connecting to %s', host_port)

  try:

    body = urllib.urlencode({

      constants.KIND_PARAM: kind,

      constants.CSV_PARAM: content,

    })

    headers = {

      'Content-Type': 'application/x-www-form-urlencoded',

      'Content-Length': len(body),

      'Cookie': cookie,

    }

    logging.debug('Posting %d bytes to http://%s%s', len(body), host_port, uri)

    connection = httplib.HTTPConnection(host_port)

    try:

      connection.request('POST', uri, body, headers)

      response = connection.getresponse()

      status = response.status

      reason = response.reason

      content = response.read()

      logging.debug('Received response code %d: %s', status, reason)

      if status != httplib.OK:

        raise BadServerStatusError('Received code %d: %s\n%s' % (

                                   status, reason, content))

    finally:

      connection.close()

  except (IOError, httplib.HTTPException, socket.error), e:

    logging.debug('Encountered exception accessing HTTP server: %s', e)

    raise PostError(e)

def SplitURL(url):

  """Splits an HTTP URL into pieces.

  Args:

    url: String containing a full URL string (e.g.,

      'http://blah.com:8080/stuff?param=1#foo')

  Returns:

    Tuple (netloc, uri) where:

      netloc: String containing the host/port combination from the URL. The

        port is optional. (e.g., 'blah.com:8080').

      uri: String containing the relative URI of the URL. (e.g., '/stuff').

  """

  scheme, netloc, path, query, fragment = urlparse.urlsplit(url)

  return netloc, path

def ImportCSV(filename,

              post_url,

              cookie,

              batch_size,

              kind,

              split_url=SplitURL,

              openfile=file,

              create_content_generator=ContentGenerator,

              post_entities=PostEntities):

  """Imports CSV data using a series of HTTP posts.

  Args:

    filename: File on disk containing CSV data.

    post_url: URL to post the Entity data to.

    cookie: Full cookie header to use while connecting.

    batch_size: Maximum number of Entity objects to post with each request.

    kind: Entity kind of the objects being posted.

    split_url, openfile, create_content_generator, post_entities: Used for

      dependency injection.

  Returns:

    True if all entities were imported successfully; False otherwise.

  """

  host_port, uri = split_url(post_url)

  csv_file = openfile(filename, 'r')

  try:

    content_gen = create_content_generator(csv_file, batch_size)

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

    for num_entities, content in content_gen:

      logging.info('Importing %d entities in %d bytes',

                   num_entities, len(content))

      try:

        content = post_entities(host_port, uri, cookie, kind, content)

      except PostError, e:

        logging.error('An error occurred while importing: %s', e)

        return False

  finally:

    csv_file.close()

  return True

def PrintUsageExit(code):

  """Prints usage information and exits with a status code.

  Args:

    code: Status code to pass to sys.exit() after displaying usage information.

  """

  print sys.modules['__main__'].__doc__ % sys.argv[0]

  sys.stdout.flush()

  sys.stderr.flush()

  sys.exit(code)

def ParseArguments(argv):

  """Parses command-line arguments.

  Prints out a help message if -h or --help is supplied.

  Args:

    argv: List of command-line arguments.

  Returns:

    Tuple (url, filename, cookie, batch_size, kind) containing the values from

    each corresponding command-line flag.

  """

  opts, args = getopt.getopt(

    argv[1:],

    'h',

    ['debug',

     'help',

     'url=',

     'filename=',

     'cookie=',

     'batch_size=',

     'kind='])

  url = None

  filename = None

  cookie = ''

  batch_size = 10

  kind = None

  encoding = None

  for option, value in opts:

    if option == '--debug':

      logging.getLogger().setLevel(logging.DEBUG)

    if option in ('-h', '--help'):

      PrintUsageExit(0)

    if option == '--url':

      url = value

    if option == '--filename':

      filename = value

    if option == '--cookie':

      cookie = value

    if option == '--batch_size':

      batch_size = int(value)

      if batch_size <= 0:

        print >>sys.stderr, 'batch_size must be 1 or larger'

        PrintUsageExit(1)

    if option == '--kind':

      kind = value

  return (url, filename, cookie, batch_size, kind)

def main(argv):

  """Runs the importer."""

  logging.basicConfig(

    level=logging.INFO,

    format='%(levelname)-8s %(asctime)s %(filename)s] %(message)s')

  args = ParseArguments(argv)

  if [arg for arg in args if arg is None]:

    print >>sys.stderr, 'Invalid arguments'

    PrintUsageExit(1)

  url, filename, cookie, batch_size, kind = args

  if ImportCSV(filename, url, cookie, batch_size, kind):

    logging.info('Import succcessful')

    return 0

  logging.error('Import failed')

  return 1

if __name__ == '__main__':

  sys.exit(main(sys.argv))