def add_id_column(df, dataset_id):
    return add_constant_column(df, dataset_id, DATASET_ID) if not\
        DATASET_ID in df.columns else df DATASET_ID = BAMBOO_RESERVED_KEY_PREFIX + 'dataset_id' INDEX = BAMBOO_RESERVED_KEY_PREFIX + 'index' def now():
    return datetime.now() def parse_timestamp_query(query, schema):
    """Interpret date column queries as JSON."""
    if query:
        for date_column in schema.datetimes(query.keys()):
            query[date_column] = {
                key: datetime.fromtimestamp(int(value)) for (key, value) in
                query[date_column].items()
            }

    return query MONGO_ID = '_id' MONGO_ID_ENCODED = MONGO_RESERVED_KEY_PREFIX + '_id' def parse_columns(dataset, formula, name, dframe=None, no_index=False):
    """Parse a formula and return columns resulting from its functions.

    Parse a formula into a list of functions then apply those functions to
    the Data Frame and return the resulting columns.

    :param formula: The formula to parse.
    :param name: Name of the formula.
    :param dframe: A DataFrame to apply functions to.
    :param no_index: Drop the index on result columns.
    """
    functions = Parser.parse_functions(formula)
    dependent_columns = Parser.dependent_columns(formula, dataset)

    # make select from dependent_columns
    if dframe is None:
        select = {col: 1 for col in dependent_columns or [MONGO_ID]}

        dframe = dataset.dframe(
            query_args=QueryArgs(select=select),
            keep_mongo_keys=True).set_index(MONGO_ID_ENCODED)

        if not dependent_columns:
            # constant column, use dummy
            dframe['dummy'] = 0

    return __build_columns(dataset, dframe, functions, name, no_index) class QueryArgs(object):
    def __init__(self, query=None, select=None, distinct=None, limit=0,
                 order_by=None, dataset=None):
        """A holder for query arguments.

        :param query: An optional query.
        :param select: An optional select to limit the fields in the dframe.
        :param distinct: Return distinct entries for this field.
        :param limit: Limit on the number of rows in the returned dframe.
        :param order_by: Sort resulting rows according to a column value and
            sign indicating ascending or descending.

        Example of `order_by`:

          - ``order_by='mycolumn'``
          - ``order_by='-mycolumn'``
        """
        self.query = parse_dates_from_query(query, dataset)
        self.select = select
        self.distinct = distinct
        self.limit = limit
        self.order_by = parse_order_by(order_by)

    def encode(self, encoding, query):
        """Encode query, order_by, and select given an encoding.

        The query will be combined with the existing query.

        :param encoding: A dict to encode the QueryArgs fields with.
        :param query: An additional dict to combine with the existing query.
        """
        self.query = replace_keys(combine_dicts(self.query, query), encoding)
        self.order_by = self.order_by and replace_keys(dict(self.order_by),
                                                       encoding).items()
        self.select = self.select and replace_keys(self.select, encoding)

    def __nonzero__(self):
        return bool(self.query or self.select or self.distinct or self.limit
                    or self.order_by) def combine_dicts(*dicts):
    """Combine dicts with keys in later dicts taking precedence."""
    return dict(chain(*[_dict.iteritems() for _dict in dicts])) def invert_dict(dict_):
    return {v: k for (k, v) in dict_.items()} if dict_ else {} def replace_keys(original, mapping):
    """Recursively replace any keys in original with their values in mappnig.

    :param original: The dictionary to replace keys in.
    :param mapping: A dict mapping keys to new keys.

    :returns: Original with keys replaced via mapping.
    """
    return original if not type(original) in (dict, list) else {
        mapping.get(k, k): {
            dict: lambda: replace_keys(v, mapping),
            list: lambda: [replace_keys(vi, mapping) for vi in v]
        }.get(type(v), lambda: v)() for k, v in original.iteritems()} class AbstractModel(object):
    """An abstact class for all MongoDB models.

    Attributes:

    - __collection__: The MongoDB collection to communicate with.
    - STATE: A key with which to store state.
    - STATE_PENDING: A value for the pending state.
    - STATE_READY: A value for the ready state.

    """

    __collection__ = None
    __collectionname__ = None

    DB_READ_BATCH_SIZE = 1000
    DB_SAVE_BATCH_SIZE = 2000
    ERROR_MESSAGE = 'error_message'
    GROUP_DELIMITER = ','  # delimiter when passing multiple groups as a string
    MIN_BATCH_SIZE = 50
    STATE = 'state'
    STATE_FAILED = 'failed'
    STATE_PENDING = 'pending'
    STATE_READY = 'ready'

    @property
    def state(self):
        return self.record[self.STATE]

    @property
    def error_message(self):
        return self.record.get(self.ERROR_MESSAGE)

    @property
    def is_pending(self):
        return self.state == self.STATE_PENDING

    @property
    def is_ready(self):
        return self.state == self.STATE_READY

    @property
    def record_ready(self):
        return self.record is not None and self.state == self.STATE_READY

    @property
    def clean_record(self):
        """Remove reserved keys from records."""
        _dict = {
            key: value for (key, value) in self.record.items() if not key in
            BAMBOO_RESERVED_KEYS
        }
        return remove_mongo_reserved_keys(_dict)

    @classmethod
    def set_collection(cls, collection_name):
        """Return a MongoDB collection for the passed name.

        :param collection_name: The name of collection to return.

        :returns: A MongoDB collection from the current database.
        """
        return Database.db()[collection_name]

    @classproperty
    @classmethod
    def collection(cls):
        """Set the internal collection to the class' collection name."""
        if not cls.__collection__:
            cls.__collection__ = AbstractModel.set_collection(
                cls.__collectionname__)

        return cls.__collection__

    @classmethod
    def create(cls, *args):
        model = cls()
        return model.save(*args)

    @classmethod
    def find(cls, query_args, as_dict=False, as_cursor=False):
        """An interface to MongoDB's find functionality.

        :param query_args: An optional QueryArgs to hold the query arguments.
        :param as_cursor: If True, return the cursor.
        :param as_dict: If True, return dicts and not model instances.

        :returns: A list of dicts or model instances for each row returned.
        """
        cursor = cls.collection.find(query_args.query,
                                     query_args.select,
                                     sort=query_args.order_by,
                                     limit=query_args.limit)

        if as_cursor:
            return cursor
        else:
            return [record for record in cursor] if as_dict else [
                cls(record) for record in cursor
            ]

    @classmethod
    def find_one(cls, query, select=None, as_dict=False):
        """Return the first row matching `query` and `select` from MongoDB.

        :param query: A query to pass to MongoDB.
        :param select: An optional select to pass to MongoDB.
        :param as_dict: If true, return dicts and not model instances.

        :returns: A model instance of the row returned for this query and
            select.
        """
        record = cls.collection.find_one(query, select)

        return record if as_dict else cls(record)

    @classmethod
    def unset(cls, query, unset_query):
        """Call unset with the spec `query` the unset document `unset_query`.

        :param query: The spec restrict updates to.
        :param unset_query: The query to pass to unset.
        """
        cls.collection.update(query, {"$unset": unset_query}, multi=True)

    def __init__(self, record=None):
        """Instantiate with data in `record`."""
        self.record = record

    def __nonzero__(self):
        return self.record is not None

    def failed(self, message=None):
        """Perist the state of the current instance to `STATE_FAILED`.

        :params message: A string store as the error message, default None.
        """
        doc = {self.STATE: self.STATE_FAILED}

        if message:
            doc.update({self.ERROR_MESSAGE: message})

        self.update(doc)

    def pending(self):
        """Perist the state of the current instance to `STATE_PENDING`"""
        self.update({self.STATE: self.STATE_PENDING})

    def ready(self):
        """Perist the state of the current instance to `STATE_READY`"""
        self.update({self.STATE: self.STATE_READY})

    def delete(self, query):
        """Delete rows matching query.

        :param query: The query for rows to delete.
        """
        self.collection.remove(query)

    def save(self, record):
        """Save `record` in this model's collection.

        Save the record in the model instance's collection and set the internal
        record of this instance to the passed in record.

        :param record: The dict to save in the model's collection.

        :returns: The record passed in.
        """
        self.collection.insert(record)
        self.record = record

        return self

    def update(self, record):
        """Update the current instance with `record`.

        Update the current model instance based on its `_id`, set it to the
        passed in `record`.

        :param record: The record to replace the instance's data with.
        """
        record = dict_for_mongo(record)
        id_dict = {'_id': self.record['_id']}
        self.collection.update(id_dict, {'$set': record})

        # Set record to the latest record from the database
        self.record = self.__class__.collection.find_one(id_dict)

    def split_groups(self, groups):
        """Split a string based on the group delimiter"""
        return groups.split(self.GROUP_DELIMITER) if groups else []

    def join_groups(self, groups):
        return self.GROUP_DELIMITER.join(groups)
from math import ceil
from pandas import concat, DataFrame
from pymongo.errors import AutoReconnect
from bamboo.core.frame import add_id_column, DATASET_ID, INDEX
from bamboo.lib.datetools import now, parse_timestamp_query
from bamboo.lib.mongo import MONGO_ID, MONGO_ID_ENCODED
from bamboo.lib.parsing import parse_columns
from bamboo.lib.query_args import QueryArgs
from bamboo.lib.utils import combine_dicts, invert_dict, replace_keys
from bamboo.models.abstract_model import AbstractModel

    :param dframe: The DataFrame to encode.
    :param dataset: The Dataset to use a mapping for.
    :param append_index: Add index to the DataFrame, default True.

    :returns: A modified `dframe`.
    """
    if append_index:
        dframe = add_index(dframe)

    dframe = add_id_column(dframe, dataset.dataset_id)
    encoded_columns_map = dataset.schema.rename_map_for_dframe(dframe)

    return dframe.rename(columns=encoded_columns_map)


def update_calculations(record, dataset):
    calculations = dataset.calculations(include_aggs=False)

    if len(calculations):
        dframe = DataFrame(data=record, index=[0])
        labels_to_slugs = dataset.schema.labels_to_slugs

        for c in calculations:
            columns = parse_columns(dataset, c.formula, c.name, dframe=dframe)
            record[labels_to_slugs[c.name]] = columns[0][0]

    return record


class Observation(AbstractModel):

    __collectionname__ = 'observations'

    DELETED_AT = '-1'  # use a short code for key
    ENCODING = 'enc'
    ENCODING_DATASET_ID = '%s_%s' % (DATASET_ID, ENCODING)

    @classmethod
    def delete(cls, dataset, index):
        """Delete observation at index for dataset.

        :param dataset: The dataset to delete the observation from.
        :param index: The index of the observation to delete.
        """
        query = {INDEX: index, DATASET_ID: dataset.dataset_id}
        query = cls.encode(query, dataset=dataset)

        cls.__soft_delete(query)

    @classmethod
    def delete_all(cls, dataset, query=None):
        """Delete the observations for `dataset`.

        :param dataset: The dataset to delete observations for.
        :param query: An optional query to restrict deletion.
        """
        query = query or {}
        query.update({DATASET_ID: dataset.dataset_id})
        query = cls.encode(query, dataset=dataset)

        super(cls, cls()).delete(query)

    @classmethod
    def delete_columns(cls, dataset, columns):
        """Delete a column from the dataset."""
        encoding = cls.encoding(dataset)

        cls.unset({cls.ENCODING_DATASET_ID: dataset.dataset_id},
                  {"%s.%s" % (cls.ENCODING, c): 1 for c in columns})

        cls.unset(
            cls.encode({DATASET_ID: dataset.dataset_id}, encoding=encoding),
            cls.encode({c: 1 for c in columns}, encoding=encoding))

    @classmethod
    def delete_encoding(cls, dataset):
        query = {cls.ENCODING_DATASET_ID: dataset.dataset_id}

        super(cls, cls()).delete(query)

    @classmethod
    def encoding(cls, dataset, encoded_dframe=None):
        record = super(cls, cls).find_one({
            cls.ENCODING_DATASET_ID: dataset.dataset_id}).record

        if record is None and encoded_dframe is not None:
            encoding = cls.__make_encoding(encoded_dframe)
            cls.__store_encoding(dataset, encoding)

            return cls.encoding(dataset)

        return record[cls.ENCODING] if record else None

    @classmethod
    def encode(cls, dict_, dataset=None, encoding=None):
        if dataset:
            encoding = cls.encoding(dataset)

        return replace_keys(dict_, encoding) if encoding else dict_

    @classmethod
    def decoding(cls, dataset):
        return invert_dict(cls.encoding(dataset))

    @classmethod
    def find(cls, dataset, query_args=None, as_cursor=False,
             include_deleted=False):
        """Return observation rows matching parameters.

        :param dataset: Dataset to return rows for.
        :param include_deleted: If True, return delete records, default False.
        :param query_args: An optional QueryArgs to hold the query arguments.

        :raises: `JSONError` if the query could not be parsed.

        :returns: A list of dictionaries matching the passed in `query` and
            other parameters.
        """
        encoding = cls.encoding(dataset) or {}

        query_args = query_args or QueryArgs()