Source code for omegaml.documents

  1import datetime
  2
  3from mongoengine.base.fields import ObjectIdField
  4from mongoengine.document import Document
  5from mongoengine.fields import DateTimeField, DictField, FileField, StringField
  6from mongoengine.pymongo_support import LEGACY_JSON_OPTIONS
  7from pymongo.errors import OperationFailure
  8
  9from omegaml.util import settings
 10
 11
 12# default kinds of objects
 13class MDREGISTRY:
 14    PANDAS_DFROWS = 'pandas.dfrows'  # dataframe
 15    PANDAS_SEROWS = 'pandas.serows'  # series
 16    PANDAS_HDF = 'pandas.hdf'
 17    PYTHON_DATA = 'python.data'
 18    PANDAS_DFGROUP = 'pandas.dfgroup'
 19    SKLEARN_JOBLIB = 'sklearn.joblib'
 20    OMEGAML_JOBS = 'script.ipynb'
 21    SPARK_MLLIB = 'spark.mllib'
 22    OMEGAML_RUNNING_JOBS = 'job.run'
 23    MINIBATCH_STREAM = 'stream.minibatch'
 24
 25    #: the list of accepted data types. extend using OmegaStore.register_backend
 26    KINDS = []
 27
 28

[docs]
 29class Metadata:
 30    """
 31    Metadata stores information about objects in OmegaStore
 32    """
 33
 34    # NOTE THIS IS ONLY HERE FOR DOCUMENTATION PURPOSE.
 35    #
 36    # If you use this class to save a document, it will raise a NameError
 37    #
 38    # The actual Metadata class is created in make_Metadata() below.
 39    # Rationale: If we let mongoengine create Metadata here the class
 40    # is bound to a specific MongoClient instance. Using make_Metadata
 41    # binds the class to the specific instance that exists at the time
 42    # of creation. Open to better ways.
 43
 44    # fields
 45    #: this is the name of the data
 46    name = StringField(unique_with=['bucket', 'prefix'])
 47    #: bucket
 48    bucket = StringField()
 49    #: prefix
 50    prefix = StringField()
 51    #: kind of data
 52    kind = StringField(choices=MDREGISTRY.KINDS)
 53    #: for PANDAS_HDF and SKLEARN_JOBLIB this is the gridfile
 54    gridfile = FileField()
 55    #: for PANDAS_DFROWS this is the collection
 56    collection = StringField()
 57    #: for PYTHON_DATA this is the actual document
 58    objid = ObjectIdField()
 59    #: omegaml technical attributes, e.g. column indicies
 60    kind_meta = DictField()
 61    #: customer-defined other meta attributes
 62    attributes = DictField()
 63    #: s3file attributes
 64    s3file = DictField()
 65    #: location URI
 66    uri = StringField()
 67    #: created datetime
 68    created = DateTimeField(default=datetime.datetime.now)
 69    #: created datetime
 70    modified = DateTimeField(default=datetime.datetime.now)

 71
 72
 73def make_Metadata(db_alias='omega', collection=None):
 74    # this is to create context specific Metadata class that takes the
 75    # database from the given alias at the time of use
 76    from omegaml.documents import Metadata as Metadata_base
 77
 78    collection = collection or settings().OMEGA_MONGO_COLLECTION
 79
 80    class Metadata(Metadata_base, Document):
 81        # override db_alias in gridfile
 82        gridfile = FileField(db_alias=db_alias, collection_name=collection)
 83        # the actual db is defined at runtime
 84        meta = {
 85            'db_alias': db_alias,
 86            'strict': False,
 87            'indexes': [
 88                # unique entry
 89                {'fields': ['bucket', 'prefix', 'name']},
 90                'created',  # most recent is last, i.e. [-1]
 91            ],
 92        }
 93
 94        def __new__(cls, *args, **kwargs):
 95            # undo the Metadata.__new__ protection
 96            newcls = super().__real_new__(cls)
 97            return newcls
 98
 99        def __eq__(self, other):
100            return self.objid == other.objid
101
102        def __unicode__(self):
103            fields = ('name', 'bucket', 'prefix', 'created', 'kind')
104            kwargs = ('%s=%s' % (k, getattr(self, k)) for k in self._fields.keys() if k in fields)
105            return "Metadata(%s)" % ','.join(kwargs)
106
107        def save(self, *args, **kwargs):
108            assert self.name is not None, "a dataset name is needed before saving"
109            self.modified = datetime.datetime.now()
110            return super(Metadata_base, self).save(*args, **kwargs)
111
112        def to_json(self, **kwargs):
113            kwargs['json_options'] = kwargs.get('json_options', LEGACY_JSON_OPTIONS)
114            return super().to_json(**kwargs)
115
116        def to_dict(self):
117            return self.to_mongo().to_dict()
118
119        @classmethod
120        def ensure_indexes(cls):
121            # work around to https://github.com/MongoEngine/mongoengine/issues/2502
122            # read-only users cannot create indexes
123            try:
124                super().ensure_indexes()
125            except OperationFailure:
126                pass
127
128    return Metadata
129
130
131def make_QueryCache(db_alias='omega'):
132    class QueryCache(Document):
133        collection = StringField()
134        key = StringField()
135        value = DictField()
136        meta = {'db_alias': db_alias, 'indexes': ['key']}
137
138    return QueryCache
139
140
141def raise_on_use(exc):
142    def inner(*args, **kwargs):
143        raise exc
144
145    return inner
146
147
148Metadata.__real_new__ = Metadata.__new__
149Metadata.__new__ = raise_on_use(NameError("You must use make_Metadata()() to instantiate a working object"))