publish muk_session_store - 12.0

muk_session_store/__manifest__.py

@@ -20,7 +20,7 @@
 {
     "name": "MuK Session Store",
     "summary": """Session Store Options""",
-    "version": "12.0.1.0.5",
+    "version": "12.0.1.0.6",
     "category": "Extra Tools",
     "license": "AGPL-3",
     "website": "http://www.mukit.at",

muk_session_store/store/postgres.py

@@ -23,6 +23,7 @@ import psycopg2
 import functools
 
 from contextlib import closing
+from contextlib import contextmanager
 from datetime import datetime, date
 
 from werkzeug.contrib.sessions import SessionStore
@@ -32,17 +33,15 @@ from odoo.tools import config
 
 _logger = logging.getLogger(__name__)
 
-def ensure_cursor(func):
+def retry_database(func):
     @functools.wraps(func)
     def wrapper(self, *args, **kwargs):
         for attempts in range(1, 6):
             try:
                 return func(self, *args, **kwargs)
             except psycopg2.InterfaceError as error:
-                _logger.info("SessionStore connection failed! (%s/5)" % attempts)
-                if attempts < 5:
-                    self._open_connection()
-                else:
+                _logger.warn("SessionStore connection failed! (%s/5)" % attempts)
+                if attempts >= 5:
                     raise error
     return wrapper
 
@@ -51,35 +50,27 @@ class PostgresSessionStore(SessionStore):
     def __init__(self, *args, **kwargs):
         super(PostgresSessionStore, self).__init__(*args, **kwargs)
         self.dbname = config.get('session_store_dbname', 'session_store')
-        self._open_connection()
-        self._setup_db()
+        self._setup_database(raise_exception=False)
     
+    def _setup_database(self, raise_exception=True):
+        try:
+            with db_connect(self.dbname, allow_uri=True).cursor() as cursor:
+                cursor.autocommit(True)
+                self._create_table(cursor)
+        except:
+            self._create_database()
+            self._setup_database()
+
     def _create_database(self):
-        with closing(db_connect("postgres").cursor()) as cursor:
+        with db_connect("postgres").cursor() as cursor:
             cursor.autocommit(True)
             cursor.execute("""
                 CREATE DATABASE {dbname} 
                 ENCODING 'unicode'
                 TEMPLATE 'template0';
             """.format(dbname=self.dbname))
-        self._setup_db()
-        
-    def _open_connection(self, create_db=True):
-        try:
-            connection = db_connect(self.dbname, allow_uri=True)
-            self.cursor = connection.cursor()
-            self.cursor.autocommit(True)
-        except:
-            if not create_db:
-                raise
-            self._create_database()
-            return self._open_connection(create_db=False)
-    
-    def __del__(self):
-        self.cursor.close()
 
-    @ensure_cursor
-    def _setup_db(self):
+    def _create_table(self, cursor):
         self.cursor.execute("""
             CREATE TABLE IF NOT EXISTS sessions (
                 sid varchar PRIMARY KEY,
@@ -87,48 +78,61 @@ class PostgresSessionStore(SessionStore):
                 payload text NOT NULL
             );
         """)
-    
-    @ensure_cursor
+
+    @contextmanager
+    def open_cursor(self):
+        connection = db_connect(self.dbname, allow_uri=True)
+        cursor = connection.cursor()
+        cursor.autocommit(True)
+        yield cursor
+        cursor.close()
+   
+    @retry_database
     def save(self, session):
-        self.cursor.execute("""
-            INSERT INTO sessions (sid, write_date, payload)
-            VALUES (%(sid)s, now() at time zone 'UTC', %(payload)s)
-            ON CONFLICT (sid)
-            DO UPDATE SET payload = %(payload)s, write_date = now() at time zone 'UTC';
-        """, dict(sid=session.sid, payload=json.dumps(dict(session))))
+        with open_cursor() as cursor:
+            cursor.execute("""
+                INSERT INTO sessions (sid, write_date, payload)
+                VALUES (%(sid)s, now() at time zone 'UTC', %(payload)s)
+                ON CONFLICT (sid)
+                DO UPDATE SET payload = %(payload)s, write_date = now() at time zone 'UTC';
+            """, dict(sid=session.sid, payload=json.dumps(dict(session))))
         
-    @ensure_cursor
+    @retry_database
     def delete(self, session):
-        self.cursor.execute("DELETE FROM sessions WHERE sid=%s;", [session.sid])
+        with open_cursor() as cursor:
+            cursor.execute("DELETE FROM sessions WHERE sid=%s;", [session.sid])
 
-    @ensure_cursor
+    @retry_database
     def get(self, sid):
         if not self.is_valid_key(sid):
             return self.new()
-        self.cursor.execute("""
-            SELECT payload, write_date 
-            FROM sessions WHERE sid=%s;
-        """, [sid])
-        try:
-            payload, write_date = self.cursor.fetchone()
-            if write_date.date() != datetime.today().date():
-                self.cursor.execute("""
-                    UPDATE sessions 
-                    SET write_date = now() at time zone 'UTC' 
-                    WHERE sid=%s;
-                """, [sid])
-            return self.session_class(json.loads(payload), sid, False)
-        except Exception:
-            return self.session_class({}, sid, False)
+        with open_cursor() as cursor:
+            cursor.execute("""
+                SELECT payload, write_date 
+                FROM sessions WHERE sid=%s;
+            """, [sid])
+            try:
+                payload, write_date = cursor.fetchone()
+                if write_date.date() != datetime.today().date():
+                    cursor.execute("""
+                        UPDATE sessions 
+                        SET write_date = now() at time zone 'UTC' 
+                        WHERE sid=%s;
+                    """, [sid])
+                return self.session_class(json.loads(payload), sid, False)
+            except Exception:
+                return self.session_class({}, sid, False)
     
-    @ensure_cursor
+    @retry_database
     def list(self):
-        self.cursor.execute("SELECT sid FROM sessions;")
-        return [record[0] for record in self.cursor.fetchall()]
+        with open_cursor() as cursor:
+            cursor.execute("SELECT sid FROM sessions;")
+            return [record[0] for record in cursor.fetchall()]
     
-    @ensure_cursor
+    @retry_database
     def clean(self):
-        self.cursor.execute("""
-            DELETE FROM sessions 
-            WHERE now() at time zone 'UTC' - write_date > '7 days';
-        """)
+        with open_cursor() as cursor:
+            cursor.execute("""
+                DELETE FROM sessions 
+                WHERE now() at time zone 'UTC' - write_date > '7 days';
+            """)

muk_session_store/store/redis.py

@@ -1,93 +1,88 @@
-import json
-###################################################################################
-#
-#    Copyright (C) 2017 MuK IT GmbH
-#
-#    This program is free software: you can redistribute it and/or modify
-#    it under the terms of the GNU Affero General Public License as
-#    published by the Free Software Foundation, either version 3 of the
-#    License, or (at your option) any later version.
-#
-#    This program is distributed in the hope that it will be useful,
-#    but WITHOUT ANY WARRANTY; without even the implied warranty of
-#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
-#    GNU Affero General Public License for more details.
-#
-#    You should have received a copy of the GNU Affero General Public License
-#    along with this program.  If not, see <http://www.gnu.org/licenses/>.
-#
-###################################################################################
-
-import pickle
-import logging
-import functools
-
-from werkzeug.contrib.sessions import SessionStore
-
-from odoo.tools import config
-
-_logger = logging.getLogger(__name__)
-
-try:
-    import redis
-except ImportError:
-    pass
-
-SESSION_TIMEOUT = 60 * 60 * 24 * 7
-
-def ensure_server(func):
-    @functools.wraps(func)
-    def wrapper(self, *args, **kwargs):
-        for attempts in range(1, 6):
-            try:
-                return func(self, *args, **kwargs)
-            except redis.ConnectionError as error:
-                _logger.info("SessionStore connection failed! (%s/5)" % attempts)
-                if attempts >= 5:
-                    raise error
-    return wrapper
-
-class RedisSessionStore(SessionStore):
-    
-    def __init__(self, *args, **kwargs):
-        super(RedisSessionStore, self).__init__(*args, **kwargs)
-        self.prefix = config.get('session_store_prefix', '')
-        self.server = redis.Redis(
-            host=config.get('session_store_host', 'localhost'),
-            port=int(config.get('session_store_port', 6379)),
-            db=int(config.get('session_store_dbindex', 1)),
-            password=config.get('session_store_pass', None)
-        )
-        self._check_server()
-    
-    def _encode_session_key(self, kex):
-        return key.encode('utf-8') if isinstance(key, str) else key
-    
-    def _get_session_key(self, sid):
-        return self._encode_session_key(self.key_prefix + sid)
-    
-    @ensure_server
-    def _check_server(self):
-        self.server.ping()
-    
-    @ensure_server
-    def save(self, session):
-        key = self._get_session_key(session.sid)
-        payload = pickle.dumps(dict(session), pickle.HIGHEST_PROTOCOL)
-        self.server.setex(name=key, value=payload, time=SESSION_TIMEOUT)
-    
-    @ensure_server
-    def delete(self, session):
-        self.server.delete(self._get_session_key(session.sid))
-    
-    @ensure_server
-    def get(self, sid):
-        if not self.is_valid_key(sid):
-            return self.new()
-        key = self._get_session_key(sid)
-        payload = self.server.get(key)
-        if payload:
-            self.server.setex(name=key, value=payload, time=SESSION_TIMEOUT)
-            return self.session_class(pickle.loads(payload), sid, False)
-        else:
+import json
+###################################################################################
+#
+#    Copyright (C) 2017 MuK IT GmbH
+#
+#    This program is free software: you can redistribute it and/or modify
+#    it under the terms of the GNU Affero General Public License as
+#    published by the Free Software Foundation, either version 3 of the
+#    License, or (at your option) any later version.
+#
+#    This program is distributed in the hope that it will be useful,
+#    but WITHOUT ANY WARRANTY; without even the implied warranty of
+#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+#    GNU Affero General Public License for more details.
+#
+#    You should have received a copy of the GNU Affero General Public License
+#    along with this program.  If not, see <http://www.gnu.org/licenses/>.
+#
+###################################################################################
+
+import pickle
+import logging
+import functools
+
+from werkzeug.contrib.sessions import SessionStore
+
+from odoo.tools import config
+
+_logger = logging.getLogger(__name__)
+
+try:
+    import redis
+except ImportError:
+    pass
+
+SESSION_TIMEOUT = 60 * 60 * 24 * 7
+
+def retry_redis(func):
+    @functools.wraps(func)
+    def wrapper(self, *args, **kwargs):
+        for attempts in range(1, 6):
+            try:
+                return func(self, *args, **kwargs)
+            except redis.ConnectionError as error:
+                _logger.warn("SessionStore connection failed! (%s/5)" % attempts)
+                if attempts >= 5:
+                    raise error
+    return wrapper
+
+class RedisSessionStore(SessionStore):
+    
+    def __init__(self, *args, **kwargs):
+        super(RedisSessionStore, self).__init__(*args, **kwargs)
+        self.prefix = config.get('session_store_prefix', '')
+        self.server = redis.Redis(
+            host=config.get('session_store_host', 'localhost'),
+            port=int(config.get('session_store_port', 6379)),
+            db=int(config.get('session_store_dbindex', 1)),
+            password=config.get('session_store_pass', None)
+        )
+    
+    def _encode_session_key(self, kex):
+        return key.encode('utf-8') if isinstance(key, str) else key
+    
+    def _get_session_key(self, sid):
+        return self._encode_session_key(self.key_prefix + sid)
+    
+    @retry_redis
+    def save(self, session):
+        key = self._get_session_key(session.sid)
+        payload = pickle.dumps(dict(session), pickle.HIGHEST_PROTOCOL)
+        self.server.setex(name=key, value=payload, time=SESSION_TIMEOUT)
+    
+    @retry_redis
+    def delete(self, session):
+        self.server.delete(self._get_session_key(session.sid))
+    
+    @retry_redis
+    def get(self, sid):
+        if not self.is_valid_key(sid):
+            return self.new()
+        key = self._get_session_key(sid)
+        payload = self.server.get(key)
+        if payload:
+            self.server.setex(name=key, value=payload, time=SESSION_TIMEOUT)
+            return self.session_class(pickle.loads(payload), sid, False)
+        else:
             return self.session_class({}, sid, False)