#!/usr/bin/env python
# -*- coding: utf-8 -*-

##############################################################################
#  Copyright Kitware 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.
##############################################################################

from pymongo import MongoClient

from . import base


class MongoConnector(base.DatabaseConnector):
    name = 'mongo'

    def __init__(self, *args, **kwargs):
        super(MongoConnector, self).__init__(*args, **kwargs)
        if not self.validate(**kwargs):
            return

        self.database = kwargs.get('database')
        self.collection = kwargs.get('collection')
        self.databaseUrl = kwargs.get('url')

        self.initialized = True

    def connect(self):
        self.conn = MongoClient(self.databaseUrl)

    def disconnect(self):
        self.conn.close()
        self.conn = None

    def performSelect(self, fields, queryProps={}, filters=[], client=None):
        self.connect()

        result = super(MongoConnector, self).performSelect(fields, queryProps, filters)
        coll = self.conn[self.database][self.collection]

        opts = {}
        for k, v in queryProps.iteritems():
            target = None
            if k == 'fields':
                target = 'projection'
                if v == []:
                    v = None
            elif k == 'offset':
                target = 'skip'
            elif k in ['filter', 'limit', 'no_cursor_timeout', 'cursor_type', 'sort', 'allow_partial_results', 'oplog_replay', 'modifiers']:
                target = k

            if target is not None:
                opts[target] = v

        result['data'] = list(coll.find(**opts))

        self.disconnect()

        return result

    @staticmethod
    def validate(url=None, database=None, collection=None, **kwargs):
        return url and database and collection


base.registerConnectorClass(MongoConnector.name, MongoConnector)

@@ -834,16 +834,252 @@
  base%0A%0A%0A
+def inferFields(records):%0A    fields = set()%0A    for r in records:%0A        for field in r:%0A            fields.add(field)%0A    return list(fields)%0A%0A%0Adef convertFields(headers, row):%0A    return map(lambda field: row.get(field), headers)%0A%0A%0A
 class Mo
@@ -2408,16 +2408,9 @@
 sult
-%5B'data'%5D
+s
  = l
@@ -2431,16 +2431,156 @@
 **opts))
+%0A        headers = inferFields(results)%0A        results = %5BconvertFields(headers, row) for row in results%5D%0A%0A        result%5B'data'%5D = results
 %0A%0A      

Return a list-of-lists from performSelect()