Experimenting with OS/cubes

cubes_test.py

from cubes.providers import ModelProvider
from cubes.model import Cube, Measure, MeasureAggregate, create_dimension
from cubes.backends.sql.store import SQLStore
from cubes.errors import NoSuchCubeError, NoSuchDimensionError
from cubes import Workspace

from openspending.core import create_app, db
from openspending.model import Dataset
from openspending.model.dimension import AttributeDimension
from openspending.model.dimension import DateDimension
from openspending.model.dimension import CompoundDimension

app = create_app()


class OpenSpendingModelProvider(ModelProvider):
    __extension_name__ = 'openspending'

    def __init__(self, *args, **kwargs):
        super(OpenSpendingModelProvider, self).__init__(*args, **kwargs)

    def default_metadata(self, metadata=None):
        return {}

    def requires_store(self):
        return True

    def public_dimensions(self):
        return []

    def cube(self, name, locale=None):
        dataset = Dataset.by_name(name)
        if name is None:
            raise NoSuchCubeError("Unknown dataset %s" % name, name)

        mappings = {}
        joins = []
        fact_table = dataset.model.table.name

        aggregates = [MeasureAggregate('num_entries',
                                       label='Numer of entries',
                                       function='count')]
        measures = []
        for measure in dataset.model.measures:
            cubes_measure = Measure(measure.name, label=measure.label)
            measures.append(cubes_measure)
            aggregate = MeasureAggregate(measure.name,  # + '_sum',
                                         label=measure.label,
                                         measure=measure.name,
                                         function='sum')
            aggregates.append(aggregate)
            #mappings[measure.name] = '%s.%s' % (dataset.model.table.name, measure.name)

        dimensions = []
        for dim in dataset.model.dimensions:
            meta = {'name': dim.name, 'label': dim.label}
            if isinstance(dim, AttributeDimension):
                meta['key'] = dim.name
                meta['levels'] = [{
                    'name': dim.name,
                    'label': dim.label,
                    'attributes': [dim.name]
                }]
                col_name = '%s.%s' % (fact_table, dim.name)
                mappings['%s.%s' % (dim.name, dim.name)] = col_name
            elif isinstance(dim, DateDimension):
                continue
            elif isinstance(dim, CompoundDimension):
                meta['key'] = 'name'
                meta['levels'] = [{
                    'name': dim.name,
                    'label': dim.label,
                    'attributes': [a.name for a in dim.attributes]
                }]
                joins.append({
                    'master': '%s.%s' % (fact_table, dim.name + '_id'),
                    'detail': '%s.id' % dim.table.name
                })
                for attr in dim.attributes:
                    col_name = '%s.%s' % (dim.table.name, a.name)
                    mappings['%s.%s' % (dim.name, a.name)] = col_name
            dimensions.append(create_dimension(meta))

        print mappings
        cube = Cube(name=dataset.name,
                    fact=fact_table,
                    aggregates=aggregates,
                    measures=measures,
                    label=dataset.label,
                    description=dataset.description,
                    info=dataset.as_dict(),
                    dimensions=dimensions,
                    store=self.store,
                    mappings=mappings,
                    joins=joins)
        return cube

    def dimension(self, name, locale=None, templates=[]):
        raise NoSuchDimensionError('No global dimensions in OS', name)

    def list_cubes(self):
        cubes = []
        for dataset in Dataset.all_by_account(None):
            # TODO: current_user, check if it has a model.
            cube = {
                'name': dataset.name,
                'label': dataset.label
            }
            cubes.append(cube)
        return cubes


class OpenSpendingStore(SQLStore):
    related_model_provider = "openspending"

    def model_provider_name(self):
        return self.related_model_provider

    def __init__(self, **options):
        super(OpenSpendingStore, self).__init__(**options)


workspace = Workspace()
#provider = OpenSpendingModelProvider()


from cubes.extensions import extensions
extensions.store.extensions['openspending'] = OpenSpendingStore


with app.app_context():
    workspace.register_default_store('openspending', engine=db.engine)    

    print workspace.list_cubes()

    cube = workspace.cube('openspending-example')

    browser = workspace.browser('openspending-example')
    print browser, dir(browser)

    res = browser.aggregate(drilldown=['transactionid'])
    print res, res.summary
    print list(res)