def workflow(cls): """ @workflow decorator denotes models which have .state field referring to WF State. Methods contributed to class: * set_state - change .state field with calling State.on_state_enter * fire_event - Perform transition using event name * fire_transition - Perform transition :return: """ cls.fire_event = fire_event cls.fire_transition = fire_transition cls._has_workflow = True cls._has_expired = False if is_document(cls): # MongoEngine model from mongoengine import signals as mongo_signals cls.set_state = document_set_state mongo_signals.post_save.connect(_on_document_post_save, sender=cls) if ("last_seen" in cls._fields and "expired" in cls._fields and "first_discovered" in cls._fields): cls.touch = document_touch cls._has_expired = True else: # Django model from django.db.models import signals as django_signals cls.set_state = model_set_state django_signals.post_save.connect(_on_model_post_save, sender=cls) cls.fire_transition = fire_transition cls.fire_event = fire_event return cls
def handle(self, host=None, port=None, *args, **options): connect() db = get_db() collections = set(db.list_collection_names()) for model_id in iter_model_id(): model = get_model(model_id) if not model: self.die("Invalid model: %s" % model_id) if not is_document(model): continue # Rename collections when necessary legacy_collections = model._meta.get("legacy_collections", []) for old_name in legacy_collections: if old_name in collections: new_name = model._meta["collection"] self.print("[%s] Renaming %s to %s" % (model_id, old_name, new_name)) db[old_name].rename(new_name) break # Ensure only documents with auto_create_index == False if model._meta.get("auto_create_index", True): continue # Index model self.index_model(model_id, model) # Index datastreams self.index_datastreams() # Index GridVCS self.index_gridvcs() # Index mongo cache self.index_cache() # Index datasource cache self.index_datasource_cache() # @todo: Detect changes self.print("OK")
def get_transform(self, value): def inner_document(qs): if value and not is_objectid(value): raise HTTPException( status_code=HTTPStatus.UNPROCESSABLE_ENTITY, detail=f"'{value}' is not a valid ObjectId", ) vv = self.model.get_by_id(value) if value else None if value and not vv: raise HTTPException( status_code=HTTPStatus.NOT_FOUND, detail=f"NotFond {str(self.model)}: {value}") return qs.filter(**{self.name: vv}) def inner_model(qs): if value and not is_int(value): raise HTTPException( status_code=HTTPStatus.UNPROCESSABLE_ENTITY, detail=f"'{value}' is not a Integer", ) vv = self.model.get_by_id(int(value)) if value else None if value and not vv: raise HTTPException( status_code=HTTPStatus.NOT_FOUND, detail=f"NotFond {str(self.model)}: {value}") return qs.filter(**{self.name: vv}) if is_document(self.model): return inner_document return inner_model
def cleaned_query(self, q): nq = {} for p in q: if p.endswith("__exists"): v = BooleanParameter().clean(q[p]) nq[p.replace("__exists", "__isnull")] = not v continue if "__" in p: np, lt = p.split("__", 1) else: np, lt = p, None # Skip ignored params if np in self.ignored_params or p in ( self.limit_param, self.page_param, self.start_param, self.format_param, self.sort_param, self.query_param, self.only_param, ): continue v = q[p] if self.in_param in p: v = v.split(",") if v == "\x00": v = None # Pass through interface cleaners if lt == "referred": # Unroll __referred app, fn = v.split("__", 1) model = self.site.apps[app].model if not is_document(model): extra_where = '%s."%s" IN (SELECT "%s" FROM %s)' % ( self.model._meta.db_table, self.model._meta.pk.name, model._meta.get_field(fn).attname, model._meta.db_table, ) if None in nq: nq[None] += [extra_where] else: nq[None] = [extra_where] continue elif lt and hasattr(self, "lookup_%s" % lt): # Custom lookup getattr(self, "lookup_%s" % lt)(nq, np, v) continue elif np in self.fk_fields and lt: # dereference try: nq[np] = self.fk_fields[np].objects.get(**{lt: v}) except self.fk_fields[np].DoesNotExist: nq[np] = 0 # False search continue elif np in self.clean_fields: # @todo: Check for valid lookup types v = self.clean_fields[np].clean(v) # Write back nq[p] = v return nq
def get_model_references(): """ Build model reference map :return: [(model id, [(remote model, remote field), ..], ..] """ from noc.lib.nosql import PlainReferenceField, ForeignKeyField from noc.core.model.fields import DocumentReferenceField from django.db.models import ForeignKey def add_ref(model, ref_model, ref_field): model_id = get_model_id(model) refs[model_id] += [(ref_model, ref_field)] refs = defaultdict(list) # model -> [(ref model, ref field)] for model_id in iter_model_id(): model = get_model(model_id) if not model: continue if is_document(model): # mongoengine document for fn in model._fields: f = model._fields[fn] if isinstance(f, PlainReferenceField): add_ref(f.document_type, model_id, fn) elif isinstance(f, ForeignKeyField): add_ref(f.document_type, model_id, fn) else: # Django model for f in model._meta.fields: if isinstance(f, ForeignKey): add_ref(f.rel.to, model_id, f.name) elif isinstance(f, DocumentReferenceField): add_ref(f.document, model_id, f.name) return [(m, refs[m]) for m in refs]
def reset_model_labels(cls, model_id: str, labels: List[str]): """ Unset labels from effective_labels field on models :param model_id: :param labels: :return: """ from django.db import connection model = get_model(model_id) if is_document(model): coll = model._get_collection() coll.bulk_write([ UpdateMany[{ "effective_labels": { "$in": labels } }, { "$pull": { "effective_labels": { "$in": labels } } }, ] ]) else: sql = f""" UPDATE {model._meta.db_table} SET effective_labels=array( SELECT unnest(effective_labels) EXCEPT SELECT unnest(%s::varchar[]) ) WHERE effective_labels && %s::varchar[] """ cursor = connection.cursor() cursor.execute(sql, [labels, labels])
def fix(): for model_id in BI_SYNC_MODELS: model = get_model(model_id) print("[%s]" % model_id) if is_document(model): fix_document(model) else: fix_model(model)
def resourcegroup(cls): if is_document(cls): mongo_signals.pre_save.connect(_apply_document_effective_groups, sender=cls) else: django_signals.pre_save.connect(_apply_model_effective_groups, sender=cls) return cls
def test_model_meta(model_id): model = get_model(model_id) assert model if is_document(model): pytest.skip("Not a model") assert model._meta assert model._meta.app_label assert model._meta.db_table
def test_document_meta(model_id): model = get_model(model_id) assert model if not is_document(model): pytest.skip("Not a document") assert model._meta.get( "allow_inheritance" ) is None, "'allow_inheritance' is obsolete and must not be used" assert not model._meta.get( "strict", True), "Document must be declared as {'strict': False}" assert not model._meta.get( "auto_create_index", True ), "Index autocreation must not be used (Use auto_create_index: False)"
def iter_id(self, model): if not isinstance(model, tuple): model = (model, ) for m in model: if is_document(m): for d in m._get_collection().find( {}, { "_id": 1 }, no_cursor_timeout=True).sort("_id"): yield d["_id"] else: for id in m.objects.values_list("id", flat=True): yield id
def iter_references(): for model_id in iter_model_id(): model = get_model(model_id) if not model: continue if is_document(model): # MongoEngine document for fn in model._fields: f = model._fields[fn] if isinstance(f, PlainReferenceField): yield f.document_type, model_id, fn elif isinstance(f, ForeignKeyField): yield f.document_type, model_id, fn else: # Django model for f in model._meta.fields: if isinstance(f, ForeignKey): yield f.remote_field.model, model_id, f.name elif isinstance(f, DocumentReferenceField): f_doc = f.document if not is_document(f_doc): f_doc = get_model(f_doc) yield f_doc, model_id, f.name
def bi_sync(cls): """ Denote class to add bi_id defaults :param cls: :return: """ if is_document(cls): f = cls._fields.get(BI_ID_FIELD) assert f, "%s field must be defined" % BI_ID_FIELD else: f = [f for f in cls._meta.fields if f.name == BI_ID_FIELD] assert f, "%s field must be defined" % BI_ID_FIELD f = f[0] f.default = new_bi_id return cls
def inner(m_cls): # Install handlers if is_document(m_cls): from mongoengine import signals as mongo_signals mongo_signals.post_save.connect(on_post_save, sender=m_cls, weak=False) else: from django.db.models import signals as django_signals django_signals.post_save.connect(on_post_save, sender=m_cls, weak=False) return m_cls
def test_load_data(initial_data): global model_refs, m2m_refs data = initial_data assert "$model" in data model = get_model(data["$model"]) assert model # Get reference fields refs = model_refs.get(data["$model"]) # name -> model mrefs = m2m_refs.get(data["$model"]) # name -> model if refs is None: refs = {} mrefs = {} if is_document(model): pass else: # Django models for f in model._meta.fields: if isinstance(f, (models.ForeignKey, CachedForeignKey)): refs[f.name] = f.remote_field.model elif isinstance(f, DocumentReferenceField): refs[f.name] = f.document for f in model._meta.many_to_many: if isinstance(f, models.ManyToManyField): mrefs[f.name] = f.remote_field.model model_refs[data["$model"]] = refs m2m_refs[data["$model"]] = mrefs # kwargs = {} m2m = {} for k in data: if k.startswith("$"): continue if k in refs: kwargs[k] = _dereference(refs[k], data[k]) elif k in mrefs: m2m[k] = [_dereference(mrefs[k], x) for x in data[k]] else: kwargs[k] = data[k] d = model(**kwargs) d.save() assert d.pk # M2M fields for k in m2m: for r in m2m[k]: getattr(d, k).add(r)
def api_model_fields_lookup(self, request, model_id): try: model = get_model(model_id=model_id) except AssertionError: return self.render_json( { "status": False, "message": "Not found model by id: %s" % model_id }, status=self.NOT_FOUND, ) # Get links if is_document(model): fields = model._fields else: fields = [f.name for f in model._meta.fields] return [{ "id": name, "label": name } for name in fields if name not in IGNORED_FIELDS]
def iter_id(self, model): if not isinstance(model, tuple): model = (model, ) for m in model: if is_document(m): match = {} while True: print(match) cursor = (m._get_collection().find( match, { "_id": 1 }, no_cursor_timeout=True).sort("_id").limit(BATCH_SIZE)) for d in cursor: yield d["_id"] if match and match["_id"]["$gt"] == d["_id"]: break match = {"_id": {"$gt": d["_id"]}} else: for id in m.objects.values_list("id", flat=True): yield id
def unset_cient_group(self, model_id: str): from django.db import connection model = get_model(model_id) if is_document(model): coll = model._get_collection() coll.bulk_write([ UpdateMany[{ "effective_client_groups": { "$in": [self.id] } }, { "$pull": { "effective_client_groups": { "$in": [self.id] } } }, ] ]) else: sql = f"UPDATE {model._meta.db_table} SET effective_client_groups=array_remove(effective_client_groups, '{str(self.id)}') WHERE '{str(self.id)}'=ANY (effective_service_groups)" cursor = connection.cursor() cursor.execute(sql)
def get_documents(): for model_id in iter_model_id(): model = get_model(model_id) if model and is_document(model): yield model
def model(cls, m_cls): """ Decorator to denote models with labels. Contains field validation and `effective_labels` generation. Usage: ``` @Label.model class MyModel(...) ``` Adds pre-save hook to check and process `.labels` fields. Raises ValueError if any of the labels is not exists. Target model may have `iter_effective_labels` method with following signature: ``` def iter_effective_labels(self) -> Iterable[List[str]] ``` which may yield a list of effective labels from related objects to form `effective_labels` field. :param m_cls: Target model class :return: """ def default_iter_effective_labels(instance) -> Iterable[List[str]]: yield instance.labels or [] def on_pre_save(sender, instance=None, document=None, *args, **kwargs): instance = instance or document # Clean up labels labels = Label.merge_labels( default_iter_effective_labels(instance)) instance.labels = labels # Validate instance labels can_set_label = getattr(sender, "can_set_label", lambda x: False) for label in set(instance.labels): if not can_set_label(label): # Check can_set_label method raise ValueError(f"Invalid label: {label}") # Build and clean up effective labels. Filter can_set_labels labels_iter = getattr(sender, "iter_effective_labels", default_iter_effective_labels) instance.effective_labels = [ ll for ll in Label.merge_labels(labels_iter(instance)) if can_set_label(ll) or ll[-1] in MATCH_OPS ] # Install handlers if is_document(m_cls): from mongoengine import signals as mongo_signals mongo_signals.pre_save.connect(on_pre_save, sender=m_cls, weak=False) else: from django.db.models import signals as django_signals django_signals.pre_save.connect(on_pre_save, sender=m_cls, weak=False) return m_cls