MikroWizard Initial commit | MikroMan Welcome to the world :)

py/libs/db/db.py

@@ -0,0 +1,170 @@
+#!/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)
+