#!/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 "".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)