MikroWizard.mikroman/py/libs/db/db.py

#!/usr/bin/python
# -*- coding: utf-8 -*-

# db_user_tasks.py: Models and functions for accsessing db
# MikroWizard.com , Mikrotik router management solution
# Author: Tomi.Mickelsson@iki.fi modified by sepehr.ha@gmail.com

from peewee import *
from playhouse.shortcuts import model_to_dict
from psycopg2.extensions import ISOLATION_LEVEL_SERIALIZABLE

from flask import abort
import config

import logging
log = logging.getLogger("db")

if config.IS_SQLITE:
    # config.DATABASE_HOST is full path to sqlite file
    database = SqliteDatabase(config.DATABASE_HOST, pragmas={})
else:
    from playhouse.postgres_ext import PostgresqlExtDatabase, ArrayField, BinaryJSONField, BooleanField, JSONField
    # support for arrays of uuid
    import psycopg2.extras
    psycopg2.extras.register_uuid()

    database = PostgresqlExtDatabase(config.DATABASE_NAME,
        user=config.DATABASE_USER, password=config.DATABASE_PASSWORD,
        host=config.DATABASE_HOST, port=config.DATABASE_PORT , isolation_level=ISOLATION_LEVEL_SERIALIZABLE)


# --------------------------------------------------------------------------
# Base model and common methods

class BaseModel(Model):
    """Base class for all database models."""

    # exclude these fields from the serialized dict
    EXCLUDE_FIELDS = []

    def serialize(self):
        """Serialize the model into a dict."""
        d = model_to_dict(self, recurse=False, exclude=self.EXCLUDE_FIELDS)
        d["id"] = str(d["id"]) # unification: id is always a string
        return d

    class Meta:
        database = database


def get_object_or_404(model, **kwargs):
    """Retrieve a single object or abort with 404."""

    try:
        return model.get(**kwargs)
    except model.DoesNotExist:
        log.warning("NO OBJECT {} {}".format(model, kwargs))
        abort(200)

def get_object_or_none(model, **kwargs):
    """Retrieve a single object or return None."""

    try:
        return model.get(**kwargs)
    except model.DoesNotExist:
        return None


# --------------------------------------------------------------------------
# USER

class User(BaseModel):

    # Should user.id be an integer or uuid? Both have pros and cons.
    # Since user.id is sensitive data, I selected uuid here.
    if not config.IS_SQLITE:
        id = UUIDField(primary_key=True)
        id.auto_increment = True # is auto generated by server

    username = TextField()
    password = TextField()
    hash = TextField()
    first_name = TextField()
    last_name = TextField()
    role = TextField()
    email = TextField()
    adminperms = TextField()
    if not config.IS_SQLITE:
        tags = ArrayField(TextField)
    else:
        tags = TextField()

    created = DateTimeField()
    modified = DateTimeField()

    EXCLUDE_FIELDS = [password,hash] # never expose password


    def is_superuser(self):
        return self.role == "superuser"

    def full_name(self):
        return "{} {}".format(self.first_name, self.last_name or '')

    def serialize(self):
        """Serialize this object to dict/json."""

        d = super(User, self).serialize()

        # add extra data
        d["fullname"] = self.full_name()
        d["tags"] = self.tags or [] # never None
        return d

    def __str__(self):
        return "<User {}, {}, role={}>".format(self.id,
                self.username, self.role)

    class Meta:
        db_table = 'users'


def get_user(uid):
    """Return user object or throw."""
    return get_object_or_404(User, id=uid)


def get_user_by_username(username):
    """Return user object or None"""

    if not username:
        return None

    try:
        # return User.select().where(User.username == username).get()
        # case insensitive query
        if config.IS_SQLITE:
            sql = "SELECT * FROM users where username = ? LIMIT 1"
            args = username.lower()
        else:
            sql = "SELECT * FROM users where LOWER(username) = LOWER(%s) LIMIT 1"
            args = (username,)
        return list(User.raw(sql, args))[0]

    except IndexError:
        return None


def query_users(page=0, limit=1000, search=None):
    """Return list of users. Desc order"""

    page = int(page or 0)
    limit = int(limit or 1000)

    q = User.select()
    if search:
        search = "%"+search+"%"
        q = q.where(User.first_name ** search | User.last_name ** search |
                User.username ** search)
    q = q.paginate(page, limit).order_by(User.id.desc())
    return q


# --------------------------------------------------------------------------

if __name__ == '__main__':

    # quick adhoc tests
    logging.basicConfig(level=logging.DEBUG)