You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

137 lines
4.9 KiB

  1. ###################################################################################
  2. #
  3. # Copyright (C) 2017 MuK IT GmbH
  4. #
  5. # This program is free software: you can redistribute it and/or modify
  6. # it under the terms of the GNU Affero General Public License as
  7. # published by the Free Software Foundation, either version 3 of the
  8. # License, or (at your option) any later version.
  9. #
  10. # This program is distributed in the hope that it will be useful,
  11. # but WITHOUT ANY WARRANTY; without even the implied warranty of
  12. # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
  13. # GNU Affero General Public License for more details.
  14. #
  15. # You should have received a copy of the GNU Affero General Public License
  16. # along with this program. If not, see <http://www.gnu.org/licenses/>.
  17. #
  18. ###################################################################################
  19. import json
  20. import logging
  21. import psycopg2
  22. import functools
  23. from contextlib import closing
  24. from contextlib import contextmanager
  25. from datetime import datetime, date
  26. from werkzeug.contrib.sessions import SessionStore
  27. from odoo.sql_db import db_connect
  28. from odoo.tools import config
  29. _logger = logging.getLogger(__name__)
  30. def retry_database(func):
  31. @functools.wraps(func)
  32. def wrapper(self, *args, **kwargs):
  33. for attempts in range(1, 6):
  34. try:
  35. return func(self, *args, **kwargs)
  36. except psycopg2.InterfaceError as error:
  37. _logger.warn("SessionStore connection failed! (%s/5)" % attempts)
  38. if attempts >= 5:
  39. raise error
  40. return wrapper
  41. class PostgresSessionStore(SessionStore):
  42. def __init__(self, *args, **kwargs):
  43. super(PostgresSessionStore, self).__init__(*args, **kwargs)
  44. self.dbname = config.get('session_store_dbname', 'session_store')
  45. self._setup_database(raise_exception=False)
  46. def _setup_database(self, raise_exception=True):
  47. try:
  48. with db_connect(self.dbname, allow_uri=True).cursor() as cursor:
  49. cursor.autocommit(True)
  50. self._create_table(cursor)
  51. except:
  52. self._create_database()
  53. self._setup_database()
  54. def _create_database(self):
  55. with db_connect("postgres").cursor() as cursor:
  56. cursor.autocommit(True)
  57. cursor.execute("""
  58. CREATE DATABASE {dbname}
  59. ENCODING 'unicode'
  60. TEMPLATE 'template0';
  61. """.format(dbname=self.dbname))
  62. def _create_table(self, cursor):
  63. cursor.execute("""
  64. CREATE TABLE IF NOT EXISTS sessions (
  65. sid varchar PRIMARY KEY,
  66. write_date timestamp without time zone NOT NULL,
  67. payload text NOT NULL
  68. );
  69. """)
  70. @contextmanager
  71. def open_cursor(self):
  72. connection = db_connect(self.dbname, allow_uri=True)
  73. cursor = connection.cursor()
  74. cursor.autocommit(True)
  75. yield cursor
  76. cursor.close()
  77. @retry_database
  78. def save(self, session):
  79. with self.open_cursor() as cursor:
  80. cursor.execute("""
  81. INSERT INTO sessions (sid, write_date, payload)
  82. VALUES (%(sid)s, now() at time zone 'UTC', %(payload)s)
  83. ON CONFLICT (sid)
  84. DO UPDATE SET payload = %(payload)s, write_date = now() at time zone 'UTC';
  85. """, dict(sid=session.sid, payload=json.dumps(dict(session))))
  86. @retry_database
  87. def delete(self, session):
  88. with self.open_cursor() as cursor:
  89. cursor.execute("DELETE FROM sessions WHERE sid=%s;", [session.sid])
  90. @retry_database
  91. def get(self, sid):
  92. if not self.is_valid_key(sid):
  93. return self.new()
  94. with self.open_cursor() as cursor:
  95. cursor.execute("""
  96. SELECT payload, write_date
  97. FROM sessions WHERE sid=%s;
  98. """, [sid])
  99. try:
  100. payload, write_date = cursor.fetchone()
  101. if write_date.date() != datetime.today().date():
  102. cursor.execute("""
  103. UPDATE sessions
  104. SET write_date = now() at time zone 'UTC'
  105. WHERE sid=%s;
  106. """, [sid])
  107. return self.session_class(json.loads(payload), sid, False)
  108. except Exception:
  109. return self.session_class({}, sid, False)
  110. @retry_database
  111. def list(self):
  112. with self.open_cursor() as cursor:
  113. cursor.execute("SELECT sid FROM sessions;")
  114. return [record[0] for record in cursor.fetchall()]
  115. @retry_database
  116. def clean(self):
  117. with self.open_cursor() as cursor:
  118. cursor.execute("""
  119. DELETE FROM sessions
  120. WHERE now() at time zone 'UTC' - write_date > '7 days';
  121. """)