Dash with layout restored from a store

dash-layout-from-store.py

import flask
import psycopg2
from flask import request, Flask

import dash


class DashLayoutShim(dash.Dash):
    @staticmethod
    def _update_layout_value(layout, store: dict):
        assert isinstance(store, dict)
        for k, v in store.items():
            try:
                component_id, component_prop = k.split('.')
            except ValueError as e:
                raise ValueError(k, e)
            try:
                my_object = layout[component_id]
            except KeyError:
                print(component_id, 'not found')  # For when component ids are not yet in the layout
                continue
            setattr(my_object, component_prop, v)
        return layout

    def serve_layout(self):
        layout_flask_response = super().serve_layout()
        external_store_id = request.cookies.get('externalStoreId', None)
        if external_store_id is None:  # We're good, just do the default layout
            return layout_flask_response
        else:
            def _conn_str_pg() -> str:
                return 'my postgresql connection string here'

            store_data = None

            # SQL injection checks should go here
            
            query_str = """
                SELECT
                    store_data
                FROM application_state
                WHERE app_id='{app_id}' AND store_id='{store_id}';
                """.format(app_id=APP_NAME, store_id=external_store_id)
            with psycopg2.connect(_conn_str_pg()) as conn:
                with conn.cursor() as curs:
                    try:
                        curs.execute(query_str)
                    except Exception as e:
                        print('ERROR ON', query_str)
                        conn.rollback()
                        raise e

                    res = curs.fetchone()
                    store_data = res[0]
                    assert isinstance(store_data, dict), 'store must be dictionary!'
                    print(store_data)

            if store_data is not None:
                layout = self._update_layout_value(
                    layout=self._layout_value(),
                    store=store_data)
            else:
                raise ValueError(
                    'store data not found for {} {}'.format(APP_NAME, external_store_id))

            # TODO - Set browser cache limit - pass hash into frontend
            return flask.Response(
                json.dumps(layout,
                           cls=plotly.utils.PlotlyJSONEncoder),
                mimetype='application/json')
                

# Define Flask server and Dash application here
flask_app = Flask(import_name=APP_NAME)

app = DashLayoutShim(server=flask_app, name=APP_NAME)
    

# Add a cookie when the user hits `?requestFromExternalStore` that makes `_dash-layout` serve the right layout
@app.server.after_request
def update_store_cookie(response):
    # If we are requesting an external store id
    if request.args.get('restoreFromExternalStore', None) is not None:
        response.set_cookie('externalStoreId', request.args.get('restoreFromExternalStore'))

    # If we have gotten the layout, let's wipe the unnecessary store id
    if request.path == '{}_dash-layout'.format(app.config['routes_pathname_prefix']):
        response.set_cookie('externalStoreId', '', expires=0)

    return response