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,
@@ -88,47 +79,60 @@ class PostgresSessionStore(SessionStore):
             );
         """)
 
-    @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)
-    
-    @ensure_cursor
+        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)
+    
+    @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

@@ -35,14 +35,14 @@ except ImportError:
 
 SESSION_TIMEOUT = 60 * 60 * 24 * 7
 
-def ensure_server(func):
+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.info("SessionStore connection failed! (%s/5)" % attempts)
+                _logger.warn("SessionStore connection failed! (%s/5)" % attempts)
                 if attempts >= 5:
                     raise error
     return wrapper
@@ -58,7 +58,6 @@ class RedisSessionStore(SessionStore):
             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
@@ -66,21 +65,17 @@ class RedisSessionStore(SessionStore):
     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
+    @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)
     
-    @ensure_server
+    @retry_redis
     def delete(self, session):
         self.server.delete(self._get_session_key(session.sid))
     
-    @ensure_server
+    @retry_redis
     def get(self, sid):
         if not self.is_valid_key(sid):
             return self.new()