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
