<reponame>treeverse/lakeview
from typing import List

from werkzeug.exceptions import HTTPException
from flask import Blueprint, render_template, request, jsonify
from flask_caching import Cache

import loader
from du import size_for, diff_for
from view import wants_json


def error_handler(e):
    error = e.original_exception
    if wants_json(request):
        return jsonify(error=error), 500
    return render_template('error.html', error=error)


class AthenaController(object):

    def __init__(self, db: loader.Inventory, cache: Cache):
        self.db = db
        self.cache = cache

    def as_blueprint(self) -> Blueprint:
        bp = Blueprint('athena_controllers', __name__)
        bp.route('/')(self.index_handler)
        bp.route('/du')(self.du_handler)
        bp.errorhandler(HTTPException)(error_handler)
        return bp

    def index_handler(self):
        """
        Render the index page (listing all available inventory dates)
        """
        return render_template('index.html', dates=reversed(self._enumerate_dates()))

    def du_handler(self):
        """
        Find the sizes of prefixes. Diff said sizes with other inventory dates
        """
        prefix = request.args.get('prefix') or ''
        delimiter = request.args.get('delimiter') or '/'
        date = request.args.get('date')

        compare_with_date = request.args.get('compare')
        if compare_with_date:
            cache_key = f'diff:{date}:{compare_with_date}:{delimiter}:{prefix}'
            res = self.cache.get(cache_key)
            if not res:
                res = diff_for(self.db, date, compare_with_date, prefix, delimiter)
                self.cache.set(cache_key, res)
        else:
            cache_key = f'du:{date}:{delimiter}:{prefix}'
            res = self.cache.get(cache_key)
            if not res:
                res = size_for(self.db, date, prefix, delimiter)
                self.cache.set(cache_key, res)

        if wants_json(request):
            json_kwargs = {
                'prefix': prefix,
                'delimiter': delimiter,
                'date': date,
                'compare': compare_with_date,
                'response': res,
            }
            return jsonify(**json_kwargs)

        return render_template(
            'du.html',
            prefix=prefix,
            delimiter=delimiter,
            date=date,
            compare_with_date=compare_with_date,
            response=res,
            available_dates=reversed(self._enumerate_dates())
        )

    def _enumerate_dates(self) -> List[str]:
        """
        list all available dates as present in the S3 inventory
        :return: a list of date strings as exists in the inventory listing
        """
        cached_dates = self.cache.get('inventory_dates')
        if cached_dates:
            return cached_dates
        res = self.db.query('SHOW PARTITIONS {table_name}')
        cached_dates = list([f.get('row').lstrip('dt=') for f in res])
        self.cache.set('inventory_dates', cached_dates, timeout=600)
        return cached_dates


from typing import List

from werkzeug.exceptions import HTTPException
from flask import Blueprint, render_template, request, jsonify
from flask_caching import Cache

import loader
from du import size_for, diff_for
from view import wants_json


def error_handler(e):
    error = e.original_exception
    if wants_json(request):
        return jsonify(error=error), 500
    return render_template('error.html', error=error)


class AthenaController(object):

    def __init__(self, db: loader.Inventory, cache: Cache):
        self.db = db
        self.cache = cache

    def as_blueprint(self) -> Blueprint:
        bp = Blueprint('athena_controllers', __name__)
        bp.route('/')(self.index_handler)
        bp.route('/du')(self.du_handler)
        bp.errorhandler(HTTPException)(error_handler)
        return bp

    def index_handler(self):
        """
        Render the index page (listing all available inventory dates)
        """
        return render_template('index.html', dates=reversed(self._enumerate_dates()))

    def du_handler(self):
        """
        Find the sizes of prefixes. Diff said sizes with other inventory dates
        """
        prefix = request.args.get('prefix') or ''
        delimiter = request.args.get('delimiter') or '/'
        date = request.args.get('date')

        compare_with_date = request.args.get('compare')
        if compare_with_date:
            cache_key = f'diff:{date}:{compare_with_date}:{delimiter}:{prefix}'
            res = self.cache.get(cache_key)
            if not res:
                res = diff_for(self.db, date, compare_with_date, prefix, delimiter)
                self.cache.set(cache_key, res)
        else:
            cache_key = f'du:{date}:{delimiter}:{prefix}'
            res = self.cache.get(cache_key)
            if not res:
                res = size_for(self.db, date, prefix, delimiter)
                self.cache.set(cache_key, res)

        if wants_json(request):
            json_kwargs = {
                'prefix': prefix,
                'delimiter': delimiter,
                'date': date,
                'compare': compare_with_date,
                'response': res,
            }
            return jsonify(**json_kwargs)

        return render_template(
            'du.html',
            prefix=prefix,
            delimiter=delimiter,
            date=date,
            compare_with_date=compare_with_date,
            response=res,
            available_dates=reversed(self._enumerate_dates())
        )

    def _enumerate_dates(self) -> List[str]:
        """
        list all available dates as present in the S3 inventory
        :return: a list of date strings as exists in the inventory listing
        """
        cached_dates = self.cache.get('inventory_dates')
        if cached_dates:
            return cached_dates
        res = self.db.query('SHOW PARTITIONS {table_name}')
        cached_dates = list([f.get('row').lstrip('dt=') for f in res])
        self.cache.set('inventory_dates', cached_dates, timeout=600)
        return cached_dates
Render the index page (listing all available inventory dates) Find the sizes of prefixes. Diff said sizes with other inventory dates list all available dates as present in the S3 inventory         :return: a list of date strings as exists in the inventory listing