class BasePyModelModel(MongoModel):
_phantom = False
_db = None
DB_RECONNECT_ATTEMPTS = 5
deleted = fields.BooleanField(default=False)
delete_stamp = fields.UUIDField(blank=True)
created_at = DefaultDateTimeField()
modified_at = fields.DateTimeField()
@classmethod
def reload_mappings(cls):
logging.info("Reloading collections.")
for clazz in cls._lazy_get_classes():
if hasattr(clazz, "_reload_mapping"):
clazz._reload_mapping()
@classmethod
def create_indexes(cls, database_name=None):
logging.info("Creating indexes for model objects..")
for clazz in cls._lazy_get_classes():
if hasattr(clazz, "_ensure_indexes"):
clazz._ensure_indexes()
if hasattr(clazz, "_ensure_index"):
clazz._ensure_index(clazz.F_DELETED, name=clazz.F_DELETED, sparse=True)
@classmethod
def _create_collection(cls):
"""
This method provides a default implementation. It is called when the model is initialized.
To provide a specialized implementation (such as for the activity model), this method is overridden in
the derived class.
"""
try:
cls._db.validate_collection(cls._COLLECTION_NAME)
except pymongo.errors.OperationFailure as e:
# Create only when the collection do not exist.
try:
cls._db.create_collection(cls._COLLECTION_NAME)
except CollectionInvalid as e:
# the collection already exists
logging.debug(
'Collection "%s" was not created. Reason: %s'
% (cls._COLLECTION_NAME, str(e))
)
@classmethod
def create_collections(cls):
logging.info("Creating collections.")
for clazz in cls._lazy_get_classes():
if hasattr(clazz, "_create_collection"):
clazz._create_collection()
@classmethod
def remove_all_data_from_collections(cls, database_name="keyflow_tests"):
"""
This method removes all documents from all collections.
Should normally not be used but can come
in handy in order to speed up tests for instance.
"""
for collection in cls._db.list_collection_names():
try:
if collection == "activities" or collection == "keyflow_rules_log":
cls._db[collection].drop()
else:
cls._db[collection].delete_many({})
except OperationFailure:
pass
@classmethod
def _lazy_get_classes(cls):
from keyflow.models.party import Party
from keyflow.models.guest_account import GuestAccount
from keyflow.models.party_requests import PartyRequest
from keyflow.models.party_chat_message import PartyChatMessage
all_classes = (Party, GuestAccount, PartyRequest, PartyChatMessage)
return all_classes
@classmethod
def initialize(
cls,
host=None,
database_name=None,
ssl=None,
ssl_cert=None,
use_ssl=True,
use_authentication=True,
user=None,
password=None,
):
from urllib.parse import quote_plus
logging.info("Initializing database")
hosts_string = host
connect_string = "mongodb://"
if not hosts_string:
# This could be like =["db-01:27017", "db-02:27017","db-03:27017"]
hosts_string = ",".join(options.mongodb_hosts)
if not database_name:
database_name = options.mongodb_database
if use_authentication:
# Could still be unset for development
username = user or options.mongodb_user
password = password or options.mongodb_password
if username and password:
connect_string = (
f"{connect_string}"
f"{quote_plus(username)}:"
f"{quote_plus(password)}@"
)
connect_string = f"{connect_string}{hosts_string}/{database_name}?"
if use_ssl:
if options.mongodb_ca_cert:
ssl_string = (
f"ssl={str(options.mongodb_ssl).lower()}&tlsCAFile="
f"{quote_plus(options.mongodb_ca_cert)}"
)
connect_string = f"{connect_string}{ssl_string}"
connect(connect_string)
cls._db = pymodm.connection._get_db()
for clazz in cls._lazy_get_classes():
if hasattr(clazz, "_store_db_functions"):
clazz._store_db_functions()
@classmethod
def _generate_id(cls):
for _ in range(cls.DB_RECONNECT_ATTEMPTS):
try:
result = cls._db.command(
"findandmodify",
"seq",
query={"_id": cls._COLLECTION_NAME},
update={"$inc": {"seq": int(1)}},
new=True,
upsert=True,
)
return result["value"]["seq"]
break
except AutoReconnect as e:
logging.warning("Reconnecting to DB: %s", str(e))
time.sleep(1)
def _prepare_phantom(self, data=None):
self.id = self._generate_id()
self.created_at = datetime.utcnow()
self.modified_at = self.created_at
# Adding the deleted flag to be false by default
self.deleted = False
def prepare_for_insert(self, prepare_phantom=True):
if prepare_phantom:
self._prepare_phantom()
return
def save(
self,
validate=True,
cascade=None,
full_clean=True,
force_insert=False,
*args,
**kwargs,
):
if self._phantom or not self.id:
self._prepare_phantom()
for _ in range(self.DB_RECONNECT_ATTEMPTS):
try:
super(MongoModel, self).save(
cascade=None, full_clean=True, force_insert=False
)
break
except errors.ValidationError as ex:
raise ex
break
except AutoReconnect as e:
logging.warning("Reconnecting to DB: %s", str(e))
time.sleep(1)
else:
return False
self._phantom = False
return True
class Conf(MongoModel):
"""Runtime configuration.
Attributes:
project_id: The project identifier. Notice how there is only one Conf
instance active at any given time and that the project_id of the
various services loaded by the Conf is always set to the conf
project_id value
bootstrap: the list of services to be loaded at bootstrap with their
configuration
storage: the configuration of the services
services: the services
"""
project_id = fields.UUIDField(primary_key=True)
bootstrap = fields.DictField(required=False, blank=True)
storage = fields.DictField(required=False, blank=True)
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
# active services
self.services = {}
# Save pointer to LVAPPManager
self.lvapp_manager = \
srv_or_die("empower.services.ranmanager.lvapp.lvappmanager")
# Save pointer to VBSPManager
self.vbsp_manager = \
srv_or_die("empower.services.ranmanager.vbsp.vbspmanager")
def register_service(self, service_id, name, params):
"""Register new service."""
if str(service_id) in self.bootstrap:
raise ValueError("Worker %s already registered" % service_id)
params['service_id'] = service_id
params['project_id'] = self.project_id
service = self._start_service(service_id, name, params)
out = json.dumps(service.params, sort_keys=True, indent=4,
cls=JSONSerializer)
self.bootstrap[str(service_id)] = {
"name": service.name,
"params": json.loads(out)
}
self.save()
return self.services[service_id]
def unregister_service(self, service_id):
"""Unregister new service."""
if str(service_id) not in self.bootstrap:
raise ValueError("Application %s not registered" % service_id)
self._stop_service(service_id)
del self.bootstrap[str(service_id)]
del self.storage[str(service_id)]
self.save()
def reconfigure_service(self, service_id, params):
"""Reconfigure service."""
if str(service_id) not in self.bootstrap:
raise ValueError("Application %s not registered" % service_id)
service = self.services[service_id]
for param in params:
if param not in self.bootstrap[str(service_id)]['params']:
raise KeyError("Param %s undefined" % param)
setattr(service, param, params[param])
self.bootstrap[str(service_id)] = {
"name": service.name,
"params": service.params
}
self.save()
return self.services[service_id]
def start_services(self):
"""Start registered services."""
for service_id in list(self.storage):
if service_id not in self.bootstrap:
del self.storage[service_id]
self.save()
for service_id in self.bootstrap:
name = self.bootstrap[service_id]['name']
params = self.bootstrap[service_id]['params']
self._start_service(UUID(service_id), name, params)
def stop_services(self):
"""Start registered services."""
for service_id in self.bootstrap:
self._stop_service(UUID(service_id))
def _start_service(self, service_id, name, params):
"""Start an service."""
if service_id in self.services:
raise ValueError("Service %s is already running" % service_id)
init_method = getattr(import_module(name), "launch")
service = init_method(**params)
self.services[service_id] = service
self.services[service_id].start()
return service
def _stop_service(self, service_id):
"""Stop an service."""
if service_id not in self.services:
raise ValueError("Service %s not running" % service_id)
self.services[service_id].stop()
del self.services[service_id]
def to_dict(self):
"""Return JSON-serializable representation of the object."""
output = {}
output['project_id'] = self.project_id
output['bootstrap'] = self.bootstrap
return output
def to_str(self):
"""Return an ASCII representation of the object."""
return str(self.project_id)
def __str__(self):
return self.to_str()
def __hash__(self):
return hash(self.project_id)
def __eq__(self, other):
if isinstance(other, Conf):
return self.project_id == other.project_id
return False
def __ne__(self, other):
return not self.__eq__(other)
def __repr__(self):
return self.__class__.__name__ + "('" + self.to_str() + "')"
class Env(MongoModel):
"""Env class
Attributes:
project_id: The project identifier. Notice how there is only one Env
instance active at any given time and that the project_id of the
workers loaded by Env is always set to the Env's project_id
bootstrap: the list of services to be loaded at bootstrap with their
configuration
storage: the configuration of the services
services: the services
"""
project_id = fields.UUIDField(primary_key=True)
bootstrap = fields.DictField(required=False, blank=True)
storage = fields.DictField(required=False, blank=True)
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
# List of services in this Env/Project
self.services = {}
# Save pointer to EnvManager
self.manager = srv_or_die("envmanager")
@property
def wtps(self):
"""Return the WTPs."""
return srv_or_die("lvappmanager").devices
@property
def vbses(self):
"""Return the VBSes."""
return srv_or_die("vbspmanager").devices
def write_points(self, points):
"""Write points to time-series manager."""
ts_manager = srv_or_die("tsmanager")
ts_manager.write_points(points)
def save_service_state(self, service_id):
"""Save service state."""
service = self.services[service_id]
self.bootstrap[str(service.service_id)] = {
"name": service.name,
"params": serialize.serialize(service.params)
}
self.storage[str(service.service_id)] = \
serialize.serialize(service.storage)
self.save()
def remove_service_state(self, service_id):
"""Remove service state."""
del self.bootstrap[str(service_id)]
del self.storage[str(service_id)]
self.save()
def register_service(self, name, params, service_id=None):
"""Register service."""
if not service_id:
# Load a service object using the specified parameters
requested = self.load_service(uuid.uuid4(), name, params)
# Check if a service with the same parameters already exists
services = self.services.values()
found = next((s for s in services if s == requested), None)
if found:
return found
# Generate a service_id if necessary
if not service_id:
service_id = uuid.uuid4()
# Start the service
service = self.start_service(service_id, name, params)
# Save service state
self.save_service_state(service.service_id)
return service
def unregister_service(self, service_id):
"""Unregister service."""
# If not found abort
if service_id not in self.services:
raise KeyError("Service %s not registered" % service_id)
# Stop the service
self.stop_service(service_id)
# Remove service state
self.remove_service_state(service_id)
def reconfigure_service(self, service_id, params):
"""Reconfigure service."""
if service_id not in self.services:
raise KeyError("Service %s not registered" % service_id)
service = self.services[service_id]
for param in params:
if param not in self.bootstrap[str(service_id)]['params']:
raise KeyError("Param %s undefined" % param)
setattr(service, param, params[param])
self.save_service_state(service_id)
return self.services[service_id]
def start_services(self):
"""Start registered services."""
for service_id in list(self.bootstrap):
try:
name = self.bootstrap[service_id]['name']
params = self.bootstrap[service_id]['params']
storage = self.storage[service_id]
service_id = uuid.UUID(service_id)
self.start_service(service_id, name, params, storage)
except TypeError as ex:
self.manager.log.error("Unable to start service %s: %s",
name, ex)
self.remove_service_state(service_id)
def stop_services(self):
"""Stop registered services."""
for service_id in self.bootstrap:
self.stop_service(uuid.UUID(service_id))
def load_service(self, service_id, name, params):
"""Load a service instance."""
init_method = getattr(import_module(name), "launch")
service = init_method(context=self, service_id=service_id, **params)
if not isinstance(service, EWorker):
raise ValueError("Service %s not EWorker type" % name)
return service
def start_service(self, service_id, name, params, storage=None):
"""Start a service."""
# wait we are trying to start a service that already exists, abort
if service_id in self.services:
raise ValueError("Service %s is already running" % service_id)
# this will look for the launch method and call it
self.manager.log.info("Loading service: %s (%s)", name, service_id)
self.manager.log.info(" - params: %s", params)
service = self.load_service(service_id, name, params)
# add to service list
self.services[service.service_id] = service
# set storage
service.set_storage(storage)
# register handlers
for handler in service.HANDLERS:
api_manager = srv_or_die("apimanager")
api_manager.register_handler(handler)
handler.service = service
# start service
self.manager.log.info("Starting service: %s (%s)", name, service_id)
service.start()
return service
def stop_service(self, service_id):
"""Stop a service."""
if service_id not in self.services:
raise KeyError("Service %s not running" % service_id)
self.services[service_id].stop()
del self.services[service_id]
def to_dict(self):
"""Return JSON-serializable representation of the object."""
output = {}
output['project_id'] = self.project_id
output['bootstrap'] = self.bootstrap
output['storage'] = self.storage
output['platform'] = {
"machine": platform.machine(),
"node": platform.node(),
"platform": platform.platform(),
"processor": platform.processor(),
"python_version": platform.python_version(),
"release": platform.release(),
"system": platform.system(),
"version": platform.version(),
}
return output
def to_str(self):
"""Return an ASCII representation of the object."""
return str(self.project_id)
def __str__(self):
return self.to_str()
def __hash__(self):
return hash(self.project_id)
def __eq__(self, other):
if isinstance(other, Env):
return self.project_id == other.project_id
return False
def __ne__(self, other):
return not self.__eq__(other)
def __repr__(self):
return self.__class__.__name__ + "('" + self.to_str() + "')"
class Request(MongoModel):
request_id = fields.UUIDField(primary_key=True)
user_id = fields.ReferenceField(User)