data-access/nexustiles/backends/nexusproto/dao/CassandraProxy.py

# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import logging
import uuid
from configparser import NoOptionError

import nexusproto.DataTile_pb2 as nexusproto
import numpy as np
from cassandra.auth import PlainTextAuthProvider
from cassandra.cqlengine import columns, connection, CQLEngineException
from cassandra.cluster import NoHostAvailable
from cassandra.cqlengine.models import Model
from cassandra.policies import TokenAwarePolicy, DCAwareRoundRobinPolicy, WhiteListRoundRobinPolicy
from multiprocessing.synchronize import Lock
from nexusproto.serialization import from_shaped_array

INIT_LOCK = Lock(ctx=None)

logger = logging.getLogger(__name__)

class NexusTileData(Model):
    __table_name__ = 'sea_surface_temp'
    tile_id = columns.UUID(primary_key=True)
    tile_blob = columns.Blob()

    __nexus_tile = None

    def _get_nexus_tile(self):
        if self.__nexus_tile is None:
            self.__nexus_tile = nexusproto.TileData.FromString(self.tile_blob)

        return self.__nexus_tile

    def get_raw_data_array(self):

        nexus_tile = self._get_nexus_tile()
        the_tile_type = nexus_tile.tile.WhichOneof("tile_type")

        the_tile_data = getattr(nexus_tile.tile, the_tile_type)

        return from_shaped_array(the_tile_data.variable_data)

    def get_lat_lon_time_data_meta(self):
        """
        Retrieve data from data store and metadata from metadata store
        for this tile. For gridded tiles, the tile shape of the data
        will match the input shape. For example, if the input was a
        30x30 tile, all variables will also be 30x30. However, if the
        tile is a swath tile, the data will be transformed along the
        diagonal of the data matrix. For example, a 30x30 tile would
        become 900x900 where the 900 points are along the diagonal.

        Multi-variable tile will also include an extra dimension in the
        data array. For example, a 30 x 30 x 30 array would be
        transformed to N x 30 x 30 x 30 where N is the number of
         variables in this tile.

        latitude_data, longitude_data, np.array([grid_tile.time]), grid_tile_data, meta_data, is_multi_var

        :return: latitude data
        :return: longitude data
        :return: time data
        :return: data
        :return: meta data dictionary
        :return: boolean flag, True if this tile has more than one variable
        """
        is_multi_var = False

        if self._get_nexus_tile().HasField('grid_tile'):
            grid_tile = self._get_nexus_tile().grid_tile

            grid_tile_data = np.ma.masked_invalid(from_shaped_array(grid_tile.variable_data))
            latitude_data = np.ma.masked_invalid(from_shaped_array(grid_tile.latitude))
            longitude_data = np.ma.masked_invalid(from_shaped_array(grid_tile.longitude))

            if len(grid_tile_data.shape) == 2:
                grid_tile_data = grid_tile_data[np.newaxis, :]

            # Extract the meta data
            meta_data = {}
            for meta_data_obj in grid_tile.meta_data:
                name = meta_data_obj.name
                meta_array = np.ma.masked_invalid(from_shaped_array(meta_data_obj.meta_data))
                if len(meta_array.shape) == 2:
                    meta_array = meta_array[np.newaxis, :]
                meta_data[name] = meta_array

            return latitude_data, longitude_data, np.array([grid_tile.time]), grid_tile_data, meta_data, is_multi_var
        elif self._get_nexus_tile().HasField('swath_tile'):
            swath_tile = self._get_nexus_tile().swath_tile

            latitude_data = np.ma.masked_invalid(from_shaped_array(swath_tile.latitude)).reshape(-1)
            longitude_data = np.ma.masked_invalid(from_shaped_array(swath_tile.longitude)).reshape(-1)
            time_data = np.ma.masked_invalid(from_shaped_array(swath_tile.time)).reshape(-1)

            # Simplify the tile if the time dimension is the same value repeated
            if np.all(time_data == np.min(time_data)):
                time_data = np.array([np.min(time_data)])

            swath_tile_data = np.ma.masked_invalid(from_shaped_array(swath_tile.variable_data))

            tile_data = self._to_standard_index(swath_tile_data,
                                                (len(time_data), len(latitude_data), len(longitude_data)))

            # Extract the meta data
            meta_data = {}
            for meta_data_obj in swath_tile.meta_data:
                name = meta_data_obj.name
                actual_meta_array = np.ma.masked_invalid(from_shaped_array(meta_data_obj.meta_data))
                reshaped_meta_array = self._to_standard_index(actual_meta_array, tile_data.shape)
                meta_data[name] = reshaped_meta_array

            return latitude_data, longitude_data, time_data, tile_data, meta_data, is_multi_var
        elif self._get_nexus_tile().HasField('time_series_tile'):
            time_series_tile = self._get_nexus_tile().time_series_tile

            time_series_tile_data = np.ma.masked_invalid(from_shaped_array(time_series_tile.variable_data))
            time_data = np.ma.masked_invalid(from_shaped_array(time_series_tile.time)).reshape(-1)
            latitude_data = np.ma.masked_invalid(from_shaped_array(time_series_tile.latitude))
            longitude_data = np.ma.masked_invalid(from_shaped_array(time_series_tile.longitude))

            reshaped_array = np.ma.masked_all((len(time_data), len(latitude_data), len(longitude_data)))
            idx = np.arange(len(latitude_data))
            reshaped_array[:, idx, idx] = time_series_tile_data
            tile_data = reshaped_array
            # Extract the meta data
            meta_data = {}
            for meta_data_obj in time_series_tile.meta_data:
                name = meta_data_obj.name
                meta_array = np.ma.masked_invalid(from_shaped_array(meta_data_obj.meta_data))

                reshaped_meta_array = np.ma.masked_all((len(time_data), len(latitude_data), len(longitude_data)))
                idx = np.arange(len(latitude_data))
                reshaped_meta_array[:, idx, idx] = meta_array

                meta_data[name] = reshaped_meta_array

            return latitude_data, longitude_data, time_data, tile_data, meta_data, is_multi_var
        elif self._get_nexus_tile().HasField('swath_multi_variable_tile'):
            swath_tile = self._get_nexus_tile().swath_multi_variable_tile
            is_multi_var = True

            latitude_data = np.ma.masked_invalid(from_shaped_array(swath_tile.latitude)).reshape(-1)
            longitude_data = np.ma.masked_invalid(from_shaped_array(swath_tile.longitude)).reshape(-1)
            time_data = np.ma.masked_invalid(from_shaped_array(swath_tile.time)).reshape(-1)

            # Simplify the tile if the time dimension is the same value repeated
            if np.all(time_data == np.min(time_data)):
                time_data = np.array([np.min(time_data)])

            swath_tile_data = np.ma.masked_invalid(from_shaped_array(swath_tile.variable_data))

            desired_shape = (
                len(time_data),
                len(latitude_data),
                len(longitude_data),
            )
            tile_data = self._to_standard_index(swath_tile_data, desired_shape, is_multi_var=True)

            # Extract the meta data
            meta_data = {}
            for meta_data_obj in swath_tile.meta_data:
                name = meta_data_obj.name
                actual_meta_array = np.ma.masked_invalid(from_shaped_array(meta_data_obj.meta_data))
                reshaped_meta_array = self._to_standard_index(actual_meta_array, tile_data.shape)
                meta_data[name] = reshaped_meta_array

            return latitude_data, longitude_data, time_data, tile_data, meta_data, is_multi_var
        elif self._get_nexus_tile().HasField('grid_multi_variable_tile'):
            grid_multi_variable_tile = self._get_nexus_tile().grid_multi_variable_tile
            is_multi_var = True

            grid_tile_data = np.ma.masked_invalid(from_shaped_array(grid_multi_variable_tile.variable_data))
            latitude_data = np.ma.masked_invalid(from_shaped_array(grid_multi_variable_tile.latitude))
            longitude_data = np.ma.masked_invalid(from_shaped_array(grid_multi_variable_tile.longitude))

            # If there are 3 dimensions, that means the time dimension
            # was squeezed. Add back in
            if len(grid_tile_data.shape) == 3:
                grid_tile_data = np.expand_dims(grid_tile_data, axis=1)
            # If there are 4 dimensions, that means the time dimension
            # is present. Move the multivar dimension.
            if len(grid_tile_data.shape) == 4:
                grid_tile_data = np.moveaxis(grid_tile_data, -1, 0)

            # Extract the meta data
            meta_data = {}
            for meta_data_obj in grid_multi_variable_tile.meta_data:
                name = meta_data_obj.name
                meta_array = np.ma.masked_invalid(from_shaped_array(meta_data_obj.meta_data))
                if len(meta_array.shape) == 2:
                    meta_array = meta_array[np.newaxis, :]
                meta_data[name] = meta_array

            return latitude_data, longitude_data, np.array([grid_multi_variable_tile.time]), grid_tile_data, meta_data, is_multi_var
        else:
            raise NotImplementedError("Only supports grid_tile, swath_tile, swath_multi_variable_tile, and time_series_tile")

    def get_elevation_array(self):
        tile_type = self._get_nexus_tile().WhichOneof("tile_type")
        tile_data = getattr(self._get_nexus_tile(), tile_type)

        if not tile_data.HasField('elevation'):
            return None

        elevation_data = np.ma.masked_invalid(from_shaped_array(tile_data.elevation))

        if tile_type in ['swath_tile', 'swath_multi_variable_tile']:
            latitude_data = np.ma.masked_invalid(from_shaped_array(tile_data.latitude)).reshape(-1)
            longitude_data = np.ma.masked_invalid(from_shaped_array(tile_data.longitude)).reshape(-1)
            time_data = np.ma.masked_invalid(from_shaped_array(tile_data.time)).reshape(-1)

            if np.all(time_data == np.min(time_data)):
                time_data = np.array([np.min(time_data)])

            desired_shape = (
                len(time_data),
                len(latitude_data),
                len(longitude_data),
            )

            elevation_data = self._to_standard_index(
                elevation_data,
                desired_shape,
                tile_type == 'swath_multi_variable_tile'
            )

        return elevation_data


    @staticmethod
    def _to_standard_index(data_array, desired_shape, is_multi_var=False):
        """
        Transform swath data to a standard format where data runs along
        diagonal of ND matrix and the non-diagonal data points are
        masked

        :param data_array: The data array to be transformed
        :param desired_shape: The desired shape of the resulting array
        :param is_multi_var: True if this is a multi-variable tile
        :type data_array: np.array
        :type desired_shape: tuple
        :type is_multi_var: bool
        :return: Reshaped array
        :rtype: np.array
        """

        reshaped_array = []
        if is_multi_var:
            reshaped_data_array = np.moveaxis(data_array, -1, 0)
        else:
            reshaped_data_array = [data_array]

        for variable_data_array in reshaped_data_array:
            if desired_shape[0] == 1:
                variable_reshaped_array = np.ma.masked_all((desired_shape[1], desired_shape[2]))
            else:
                variable_reshaped_array = np.ma.masked_all(desired_shape)

            row, col = np.indices(variable_data_array.shape)

            variable_reshaped_array[
                np.diag_indices(desired_shape[1], len(variable_reshaped_array.shape))] = \
                variable_data_array[
                    row.flat, col.flat]
            variable_reshaped_array.mask[
                np.diag_indices(desired_shape[1], len(variable_reshaped_array.shape))] = \
                variable_data_array.mask[
                    row.flat, col.flat]

            if desired_shape[0] == 1:
                reshaped_array.append(variable_reshaped_array[np.newaxis, :])
            else:
                reshaped_array.append(variable_reshaped_array)

        if not is_multi_var:
            # If single var, squeeze extra dim out of array
            reshaped_array = reshaped_array[0]

        return reshaped_array


class CassandraProxy(object):
    def __init__(self, config):
        self.config = config
        self.__cass_url = config.get("cassandra", "host")
        self.__cass_username = config.get("cassandra", "username")
        self.__cass_password = config.get("cassandra", "password")
        self.__cass_keyspace = config.get("cassandra", "keyspace")
        self.__cass_local_DC = config.get("cassandra", "local_datacenter")
        self.__cass_protocol_version = config.getint("cassandra", "protocol_version")
        self.__cass_dc_policy = config.get("cassandra", "dc_policy")

        try:
            self.__cass_port = config.getint("cassandra", "port")
        except NoOptionError:
            self.__cass_port = 9042

        with INIT_LOCK:
            try:
                connection.get_cluster()
            except CQLEngineException:
                self.__open()

    def __open(self):
        if self.__cass_dc_policy == 'DCAwareRoundRobinPolicy':
            dc_policy = DCAwareRoundRobinPolicy(self.__cass_local_DC)
            token_policy = TokenAwarePolicy(dc_policy)
        elif self.__cass_dc_policy == 'WhiteListRoundRobinPolicy':
            token_policy = WhiteListRoundRobinPolicy([self.__cass_url])

        if self.__cass_username and self.__cass_password:
            auth_provider = PlainTextAuthProvider(username=self.__cass_username, password=self.__cass_password)
        else:
            auth_provider = None
        try:
            connection.setup(
                [host for host in self.__cass_url.split(',')], self.__cass_keyspace,
                protocol_version=self.__cass_protocol_version, load_balancing_policy=token_policy,
                port=self.__cass_port,
                auth_provider=auth_provider
            )
        except NoHostAvailable as e:
            logger.error("Cassandra is not accessible, SDAP will not server local datasets", e)

    def fetch_nexus_tiles(self, *tile_ids):
        tile_ids = [uuid.UUID(str(tile_id)) for tile_id in tile_ids if
                    (isinstance(tile_id, str) or isinstance(tile_id, str))]

        res = []
        for tile_id in tile_ids:
            filterResults = NexusTileData.objects.filter(tile_id=tile_id)
            if len(filterResults) > 0:
                res.append(filterResults[0])

        return res