From d813601d928ad2d12e60307a321ad6f9b9fe1a8f Mon Sep 17 00:00:00 2001
From: rr- <rr-@sakuya.pl>
Date: Sun, 8 May 2016 19:41:35 +0200
Subject: [PATCH] server/api: log queries in debug mode

---
 server/.pylintrc                           |  1 +
 server/requirements.txt                    |  1 +
 server/szurubooru/app.py                   |  9 ++++-
 server/szurubooru/db/__init__.py           |  5 ++-
 server/szurubooru/db/session.py            | 42 ++++++++++++++++------
 server/szurubooru/middleware/db_session.py |  8 ++++-
 6 files changed, 53 insertions(+), 13 deletions(-)

diff --git a/server/.pylintrc b/server/.pylintrc
index 0e68f4d..067e53d 100644
--- a/server/.pylintrc
+++ b/server/.pylintrc
@@ -1,6 +1,7 @@
 [basic]
 function-rgx=^_?[a-z_][a-z0-9_]{2,}$|^test_
 method-rgx=^[a-z_][a-z0-9_]{2,}$|^test_
+good-names=ex,_,logger
 
 [variables]
 dummy-variables-rgx=_|dummy
diff --git a/server/requirements.txt b/server/requirements.txt
index c521c28..b1079d7 100644
--- a/server/requirements.txt
+++ b/server/requirements.txt
@@ -6,3 +6,4 @@ SQLAlchemy>=1.0.12
 pytest>=2.9.1
 pytest-cov>=2.2.1
 freezegun>=0.3.6
+coloredlogs==5.0
diff --git a/server/szurubooru/app.py b/server/szurubooru/app.py
index f9d16bb..c767a0f 100644
--- a/server/szurubooru/app.py
+++ b/server/szurubooru/app.py
@@ -1,7 +1,9 @@
 ''' Exports create_app. '''
 
+import logging
+import coloredlogs
 import falcon
-from szurubooru import api, errors, middleware
+from szurubooru import api, config, errors, middleware
 
 def _on_auth_error(ex, _request, _response, _params):
     raise falcon.HTTPForbidden(
@@ -38,6 +40,11 @@ def create_app():
     ''' Create a WSGI compatible App object. '''
     falcon.responders.create_method_not_allowed = create_method_not_allowed
 
+    coloredlogs.install(fmt='[%(asctime)-15s] %(name)s %(message)s')
+    if config.config['debug']:
+        logging.getLogger('szurubooru').setLevel(logging.INFO)
+        logging.getLogger('sqlalchemy.engine').setLevel(logging.INFO)
+
     app = falcon.API(
         request_type=api.Request,
         middleware=[
diff --git a/server/szurubooru/db/__init__.py b/server/szurubooru/db/__init__.py
index 68d04b9..49ca156 100644
--- a/server/szurubooru/db/__init__.py
+++ b/server/szurubooru/db/__init__.py
@@ -18,4 +18,7 @@ from szurubooru.db.comment import (
     Comment,
     CommentScore)
 from szurubooru.db.snapshot import Snapshot
-from szurubooru.db.session import session
+from szurubooru.db.session import (
+    session,
+    reset_query_count,
+    get_query_count)
diff --git a/server/szurubooru/db/session.py b/server/szurubooru/db/session.py
index 393dbd5..36f25f1 100644
--- a/server/szurubooru/db/session.py
+++ b/server/szurubooru/db/session.py
@@ -1,14 +1,36 @@
 import sqlalchemy
 from szurubooru import config
 
+class QueryCounter(object):
+    _query_count = 0
+
+    @staticmethod
+    def bump():
+        QueryCounter._query_count += 1
+
+    @staticmethod
+    def reset():
+        QueryCounter._query_count = 0
+
+    @staticmethod
+    def get():
+        return QueryCounter._query_count
+
+def create_session():
+    _engine = sqlalchemy.create_engine(
+        '{schema}://{user}:{password}@{host}:{port}/{name}'.format(
+            schema=config.config['database']['schema'],
+            user=config.config['database']['user'],
+            password=config.config['database']['pass'],
+            host=config.config['database']['host'],
+            port=config.config['database']['port'],
+            name=config.config['database']['name']))
+    sqlalchemy.event.listen(
+        _engine, 'after_execute', lambda *args: QueryCounter.bump())
+    _session_maker = sqlalchemy.orm.sessionmaker(bind=_engine)
+    return sqlalchemy.orm.scoped_session(_session_maker)
+
 # pylint: disable=invalid-name
-_engine = sqlalchemy.create_engine(
-    '{schema}://{user}:{password}@{host}:{port}/{name}'.format(
-        schema=config.config['database']['schema'],
-        user=config.config['database']['user'],
-        password=config.config['database']['pass'],
-        host=config.config['database']['host'],
-        port=config.config['database']['port'],
-        name=config.config['database']['name']))
-_session_maker = sqlalchemy.orm.sessionmaker(bind=_engine)
-session = sqlalchemy.orm.scoped_session(_session_maker)
+session = create_session()
+reset_query_count = QueryCounter.reset
+get_query_count = QueryCounter.get
diff --git a/server/szurubooru/middleware/db_session.py b/server/szurubooru/middleware/db_session.py
index 8bf6848..3a412f9 100644
--- a/server/szurubooru/middleware/db_session.py
+++ b/server/szurubooru/middleware/db_session.py
@@ -1,10 +1,16 @@
-from szurubooru import db
+import logging
+from szurubooru import config, db
+
+logger = logging.getLogger(__name__)
 
 class DbSession(object):
     ''' Attaches database session to the context of every request. '''
 
     def process_request(self, request, _response):
         request.context.session = db.session()
+        db.reset_query_count()
 
     def process_response(self, _request, _response, _resource):
         db.session.remove()
+        if config.config['debug']:
+            logger.info('Executed %d queries', db.get_query_count())