DatabaseTools.py

import grpc
from typing import Any, Dict, Tuple
from common.orm.Database import Database
from common.orm.HighLevel import get_or_create_object, update_or_create_object
from common.orm.backend.Tools import key_to_str
from common.rpc_method_wrapper.ServiceExceptions import InvalidArgumentException
from context.client.ContextClient import ContextClient
from device.proto.context_pb2 import Device, DeviceId
from device.service.driver_api.FilterFields import FilterFieldEnum
from .ConfigModel import delete_all_config_rules, grpc_config_rules_to_raw, update_config
from .ContextModel import ContextModel
from .DeviceModel import DeviceModel, DriverModel, grpc_to_enum__device_operational_status, set_drivers
from .EndPointModel import EndPointModel
from .TopologyModel import TopologyModel

def update_device_in_local_database(database : Database, device : Device) -> Tuple[DeviceModel, bool]:
    device_uuid = device.device_id.device_uuid.uuid

    for i,endpoint in enumerate(device.device_endpoints):
        endpoint_device_uuid = endpoint.endpoint_id.device_id.device_uuid.uuid
        if len(endpoint_device_uuid) == 0: endpoint_device_uuid = device_uuid
        if device_uuid != endpoint_device_uuid:
            raise InvalidArgumentException(
                'request.device_endpoints[{:d}].device_id.device_uuid.uuid'.format(i), endpoint_device_uuid,
                ['should be == {:s}({:s})'.format('request.device_id.device_uuid.uuid', device_uuid)])

    initial_config_result = update_config(database, device_uuid, 'initial', [])

    config_rules = grpc_config_rules_to_raw(device.device_config.config_rules)
    delete_all_config_rules(database, device_uuid, 'running')
    running_config_result = update_config(database, device_uuid, 'running', config_rules)

    result : Tuple[DeviceModel, bool] = update_or_create_object(database, DeviceModel, device_uuid, {
        'device_uuid'              : device_uuid,
        'device_type'              : device.device_type,
        'device_operational_status': grpc_to_enum__device_operational_status(device.device_operational_status),
        'device_initial_config_fk' : initial_config_result[0][0],
        'device_running_config_fk' : running_config_result[0][0],
    })
    db_device, updated = result
    set_drivers(database, db_device, device.device_drivers)

    for i,endpoint in enumerate(device.device_endpoints):
        endpoint_uuid = endpoint.endpoint_id.endpoint_uuid.uuid
        endpoint_device_uuid = endpoint.endpoint_id.device_id.device_uuid.uuid
        if len(endpoint_device_uuid) == 0: endpoint_device_uuid = device_uuid

        str_endpoint_key = key_to_str([device_uuid, endpoint_uuid])
        endpoint_attributes = {
            'device_fk'    : db_device,
            'endpoint_uuid': endpoint_uuid,
            'endpoint_type': endpoint.endpoint_type,
        }

        endpoint_topology_context_uuid = endpoint.endpoint_id.topology_id.context_id.context_uuid.uuid
        endpoint_topology_uuid = endpoint.endpoint_id.topology_id.topology_uuid.uuid
        if len(endpoint_topology_context_uuid) > 0 and len(endpoint_topology_uuid) > 0:
            result : Tuple[ContextModel, bool] = get_or_create_object(
                database, ContextModel, endpoint_topology_context_uuid, defaults={
                    'context_uuid': endpoint_topology_context_uuid,
                })
            db_context, _ = result

            str_topology_key = key_to_str([endpoint_topology_context_uuid, endpoint_topology_uuid])
            result : Tuple[TopologyModel, bool] = get_or_create_object(
                database, TopologyModel, str_topology_key, defaults={
                    'context_fk': db_context,
                    'topology_uuid': endpoint_topology_uuid,
                })
            db_topology, _ = result

            str_endpoint_key = key_to_str([str_endpoint_key, str_topology_key], separator=':')
            endpoint_attributes['topology_fk'] = db_topology

        result : Tuple[EndPointModel, bool] = update_or_create_object(
            database, EndPointModel, str_endpoint_key, endpoint_attributes)
        _, db_endpoint_updated = result
        updated = updated or db_endpoint_updated

    return db_device, updated

def sync_device_from_context(
    device_uuid : str, context_client : ContextClient, database : Database
    ) -> Tuple[DeviceModel, bool]:

    try:
        device : Device = context_client.GetDevice(DeviceId(device_uuid={'uuid': device_uuid}))
    except grpc.RpcError as e:
        if e.code() != grpc.StatusCode.NOT_FOUND: raise # pylint: disable=no-member
        return None
    return update_device_in_local_database(database, device)

def sync_device_to_context(db_device : DeviceModel, context_client : ContextClient) -> None:
    if db_device is None: return
    context_client.SetDevice(Device(**db_device.dump(
        include_config_rules=True, include_drivers=True, include_endpoints=True)))

def delete_device_from_context(db_device : DeviceModel, context_client : ContextClient) -> None:
    if db_device is None: return
    context_client.RemoveDevice(DeviceId(**db_device.dump_id()))

def get_device_driver_filter_fields(db_device : DeviceModel) -> Dict[FilterFieldEnum, Any]:
    if db_device is None: return {}
    database = db_device.database
    db_driver_pks = db_device.references(DriverModel)
    db_driver_names = [DriverModel(database, pk).driver.value for pk,_ in db_driver_pks]
    return {
        FilterFieldEnum.DEVICE_TYPE: db_device.device_type,
        FilterFieldEnum.DRIVER     : db_driver_names,
    }