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