gisaf-backend/src/gisaf/models/dashboard.py

from io import BytesIO
from pickle import loads
from pathlib import Path
from datetime import datetime
import logging

from sqlmodel import Field, Relationship, String
from pydantic import BaseModel
import pandas as pd

from gisaf.config import conf
from gisaf.models.metadata import gisaf
from gisaf.models.models_base import Model
from gisaf.models.misc import NotADataframeError

logger = logging.getLogger(__name__)

try:
    import matplotlib.pyplot as plt
except ImportError:
    plt = None # type: ignore


class DashboardPageSource(Model, table=True):
    __tablename__ = 'dashboard_page_source' # type: ignore
    __table_args__ = gisaf.table_args

    id: str = Field(primary_key=True)
    name: str


class DashboardPageCommon:
    """
    Base class for DashboardPage and DashboardPageSection, where some methods
    are common, eg. attachments
    """

    def get_attachment_url(self):
        ## Serve through web front-end (nginx static file)
        if not self.attachment:
            return
        base_storage_url = conf.dashboard.base_storage_url
        if not base_storage_url:
            base_storage_url = '/dashboard-attachment/'
        return f'{base_storage_url}{self.group}/{self.attachment}'

    def save_plot(self, plot):
        """
        Render the matplotlib plot (or figure) and save it in the filesystem
        :param plot: matplotlib plot or figure...
        :return:
        """
        self.ensure_dir_exists()

        ## Different types of figures supported
        fig = None
        if plt:
            if isinstance(plot, plt.Axes):
                fig = plot.figure
            elif isinstance(plot, plt.Figure):
                fig = plot
            if fig:
                fig.savefig(self.get_plot_file_path(), bbox_inches='tight')
                plt.close(fig)

        if plot and not fig:
            logger.warning('Cannot save dashboard attachment (unknown attachment type)')
            return

        #logger.info(f'Saved attachment of dashboard page {self.group}/{self.name} '
                    #f'in {self.get_attachment_file_name()}')
        return self.get_plot_file_name()

    def save_attachment(self, attached, name=None):
        """
        Save the attachment in the filesystem
        :param attached: matplotlib plot or figure...
        :return:
        """
        if not self.attachment:
            ## Not set yet (creation)
            self.attachment = f'{self.name}.png'

        self.ensure_dir_exists()

        ## Different types of figures supported
        fig = None
        if plt:
            if isinstance(attached, plt.Axes):
                fig = attached.figure
            elif isinstance(attached, plt.Figure):
                fig = attached
            if fig:
                fig.savefig(self.get_attachment_file_name(), bbox_inches='tight')
                plt.close(fig)

        if attached and not fig:
            logger.warning('Cannot save dashboard attachment (unknown attachment type)')
            return

        #logger.info(f'Saved attachment of dashboard page {self.group}/{self.name} '
                    #f'in {self.get_attachment_file_name()}')
        return self.attachment

    def get_page_df(self):
        """
        Get the dataframe of the page
        """
        if not self.df:
            return
        try:
            return pd.read_pickle(BytesIO(self.df), compression=None)
        except KeyError:
            raise NotADataframeError()

    def get_plot(self):
        return loads(self.plot)


class DashboardPage(Model, DashboardPageCommon, table=True):
    __tablename__ = 'dashboard_page' # type: ignore
    __table_args__ = gisaf.table_args

    class Admin:
        menu = 'Dashboard'

    id: int = Field(primary_key=True)
    name: str
    notebook: str
    group: str
    description: str
    attachment: str
    html: str
    viewable_role: str
    df: bytes
    plot: bytes
    source_id: int = Field(foreign_key=gisaf.table('dashboard_page_source.id'))
    time: datetime
    expanded_panes: str
    source: DashboardPageSource = Relationship()

    def __str__(self):
        return f'{self.group:s}/{self.name:s}'

    def __repr__(self):
        return f'<models.DashboardPage {self.group:s}/{self.name:s}>'

    @classmethod
    def selectinload(cls):
        return [cls.source]

    def ensure_dir_exists(self):
        """
        Make sure the directory exists, before saving the file
        """
        dir_name = Path(conf.dashboard.base_storage_dir) / self.group
        dir_name.mkdir(exist_ok=True)
        return dir_name

    def get_attachment_file_name(self):
        """
        Get the file name of the attachment
        :return:
        """
        if self.attachment:
            base_dir = Path(conf.dashboard.base_storage_dir)
            if base_dir:
                return base_dir/self.group/self.attachment
            else:
                raise UserWarning('Cannot save attachment: no notebook/base_storage_dir in gisaf config')

    def get_notebook_url(self):
        if self.notebook:
            base_url = conf.dashboard.base_source_url
            if base_url:
                return f'{base_url}{self.notebook}'
            else:
                logger.debug('Notebook: no base_url in gisaf config')


class DashboardPageSection(Model, DashboardPageCommon, table=True):
    __tablename__ = 'dashboard_page_section' # type: ignore
    __table_args__ = gisaf.table_args

    class Admin:
        menu = 'Dashboard'

    id: str = Field(primary_key=True)
    name: str
    dashboard_page_id: int = Field(foreign_key=gisaf.table('dashboard_page.id'))
    dashboard_page: DashboardPage = Relationship()

    description: str
    attachment: str
    html: str
    df: bytes
    plot: str

    def __str__(self):
        return f'{self.name} for dashboard page #{self.dashboard_page_id}'

    def __repr__(self):
        return f'<models.DashboardPageSection #{self.id}>'

    @classmethod
    def selectinload(cls):
        return [cls.dashboard_page]

    def get_plot_url(self):
        ## Serve through web front-end (nginx static file)
        if not self.plot:
            return
        return conf.dashboard.base_storage_url \
            + self.dashboard_page.group + '/' \
            + self.dashboard_page.name + '/' \
            + self.name + '.png'

    def ensure_dir_exists(self):
        """
        Make sure the directory exists, before saving the file
        """
        dir_name = Path(conf.dashboard.base_storage_dir) / self.page.group / self.page.name
        dir_name.mkdir(exist_ok=True)
        return dir_name

    def get_attachment_file_name(self):
        """
        Get the file name of the attachment
        :return:
        """
        if not self.attachment:
            return
        base_dir = Path(conf.dashboard.base_storage_dir)
        if not base_dir:
            raise UserWarning('Cannot save attachment: no notebook/base_storage_dir in gisaf config')
        return base_dir/self.page.group/self.page.name/self.attachment

    def get_plot_file_name(self):
        return f'{self.name}.png'

    def get_plot_file_path(self):
        """
        Get the file name of the plot
        :return:
        """
        if not self.plot:
            return
        base_dir = Path(conf.dashboard.base_storage_dir)
        if not base_dir:
            raise UserWarning('Cannot save attachment: no notebook/base_storage_dir in gisaf config')
        return base_dir/self.page.group/self.page.name/self.get_plot_file_name()


class Widget(Model, table=True):
    __tablename__ = 'widget' # type: ignore
    __table_args__ = gisaf.table_args
    ## CREATE TABLE gisaf.widget (name char(50) not null PRIMARY KEY, title varchar, subtitle varchar, notebook varchar, content varchar, time timestamp);
    name: str = Field(primary_key=True, sa_type=String(50)) # type: ignore
    title: str
    subtitle: str
    content: str
    time: datetime
    notebook: str

    class Admin:
        menu = 'Dashboard'


class DashboadPageSectionType(BaseModel):
    name: str
    plot: str


class DashboardPage_(BaseModel):
    name: str
    group: str
    description: str
    time: datetime | None = Field(default_factory=datetime.now)
    html: str | None = None
    attachment: str | None = None
    dfData: str | None = None
    plotData: str | None = None
    notebook: str | None = None
    expandedPanes: list[str] | None = None
    sections: list[DashboadPageSectionType] | None = None


class DashboardGroup(BaseModel):
    name: str
    pages: list[DashboardPage_]


class DashboardHome(BaseModel):
    title: str
    content: str
    footer: str