distil/artifice/database.py

86 lines
3.6 KiB
Python

from sqlalchemy import func
from .models import Resource, UsageEntry, Tenant
import json
from transformers import TransformerValidationError
from datetime import datetime
class Database(object):
def __init__(self, session):
self.session = session
# engine = create_engine(os.environ["DATABASE_URL"])
# Base.metadata.create_all(engine)
def insert_tenant(self, tenant_id, tenant_name, metadata):
"""If a tenant exists does nothing,
and if it doesn't, creates and inserts it."""
# Have we seen this tenant before?
query = self.session.query(Tenant).\
filter(Tenant.id == tenant_id)
if query.count() == 0:
self.session.add(Tenant(id=tenant_id,
info=metadata,
name=tenant_name,
created=datetime.now()
))
self.session.flush()
def enter(self, usage, start, end):
"""Creates a new database entry for every usage strategy
in a resource, for all the resources given"""
for resource in usage:
resource_id = resource.resource_id
tenant_id = resource.tenant_id
try:
for service, volume in resource.usage().iteritems():
# Have we seen this resource before?
query = self.session.query(Resource).\
filter(Resource.id == resource_id,
Resource.tenant_id == tenant_id)
if query.count() == 0:
info = json.dumps(resource.info)
self.session.add(Resource(id=resource_id,
info=str(info),
tenant_id=tenant_id,
created=datetime.now()
))
entry = UsageEntry(service=service,
volume=volume,
resource_id=resource_id,
tenant_id=tenant_id,
start=start,
end=end,
created=datetime.now()
)
self.session.add(entry)
self.session.flush()
except TransformerValidationError:
# log something related to the resource usage failing
# transform.
pass
def usage(self, start, end, tenant_id):
"""Returns a query of usage entries for a given tenant,
in the given range.
start, end: define the range to query
tenant: a tenant entry (tenant_id for now)"""
if start > end:
raise AttributeError("End must be a later date than start.")
# build a query set in the format:
# tenant_id | resource_id | service | sum(volume)
query = self.session.query(UsageEntry.tenant_id,
UsageEntry.resource_id,
UsageEntry.service,
func.sum(UsageEntry.volume).label("volume")).\
filter(UsageEntry.start >= start, UsageEntry.end <= end).\
filter(UsageEntry.tenant_id == tenant_id).\
group_by(UsageEntry.tenant_id, UsageEntry.resource_id,
UsageEntry.service)
return query