From 8d7bc2d9546f7251e68cf279ff12117b862e994d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Fr=C3=A9d=C3=A9ric=20P=C3=A9ters?= <fpeters@entrouvert.com>
Date: Sat, 25 Jun 2022 11:16:02 +0200
Subject: [PATCH] sql: store transient data in a specific table (#66620)

---
 tests/form_pages/test_all.py |   9 ++-
 tests/test_sessions.py       |  50 +++++++++++-
 tests/utilities.py           |   1 +
 wcs/publisher.py             |   1 +
 wcs/sql.py                   | 153 ++++++++++++++++++++++++++++++++++-
 5 files changed, 208 insertions(+), 6 deletions(-)

diff --git a/tests/form_pages/test_all.py b/tests/form_pages/test_all.py
index 0919997d8..4b9c41f57 100644
--- a/tests/form_pages/test_all.py
+++ b/tests/form_pages/test_all.py
@@ -1521,9 +1521,12 @@ def test_form_submit_with_just_disabled_user(pub, emails):
     user.store()
     resp = resp.form.submit('submit')
     resp = resp.follow()
-    assert 'The form has been recorded' in resp
-    assert formdef.data_class().count() == 1
-    assert formdef.data_class().select()[0].user_id is None
+    if pub.is_using_postgresql():
+        assert 'Sorry, your session have been lost.' in resp
+    else:
+        assert 'The form has been recorded' in resp
+        assert formdef.data_class().count() == 1
+        assert formdef.data_class().select()[0].user_id is None
 
 
 def test_form_titles(pub):
diff --git a/tests/test_sessions.py b/tests/test_sessions.py
index c0a7114ec..44fe6842e 100644
--- a/tests/test_sessions.py
+++ b/tests/test_sessions.py
@@ -1,11 +1,12 @@
 import datetime
 import os
+import pickle
 import shutil
 import time
 
 import pytest
 
-from wcs import fields
+from wcs import fields, sql
 from wcs.formdef import FormDef
 from wcs.qommon.http_request import HTTPRequest
 from wcs.qommon.ident.password_accounts import PasswordAccount
@@ -225,3 +226,50 @@ def test_inactive_user(pub, user, app):
     user.is_active = False
     user.store()
     assert 'Logout' not in app.get('/')
+
+
+def test_magictoken_migration(pub, app):
+    if not pub.is_using_postgresql():
+        pytest.skip('this requires SQL')
+        return
+
+    pub.session_manager.session_class.wipe()
+    sql.TransientData.wipe()
+    resp = app.get('/')
+
+    formdef = FormDef()
+    formdef.name = 'foobar'
+    formdef.fields = [
+        fields.PageField(id='0', label='1st PAGE', type='page'),
+        fields.StringField(id='1', label='string', type='string'),
+        fields.PageField(id='2', label='2nd PAGE', type='page'),
+        fields.PageField(id='3', label='3rd PAGE', type='page'),
+    ]
+    formdef.store()
+    formdef.data_class().wipe()
+
+    resp = app.get('/foobar/')
+    resp.form['f1'] = 'test'
+    resp = resp.form.submit('submit')
+
+    # migrate back session to look like before transient data table
+    assert pub.session_manager.session_class.count() == 1
+    session = pub.session_manager.session_class.select()[0]
+    session.magictokens = {}
+    for transient_data in sql.TransientData.select():
+        session.magictokens[transient_data.id] = transient_data.data
+    sql.TransientData.wipe()
+
+    conn, cur = sql.get_connection_and_cursor()
+    sql_statement = 'UPDATE sessions SET session_data = %s WHERE id = %s'
+    cur.execute(sql_statement, (bytearray(pickle.dumps(session.__dict__, protocol=2)), session.id))
+    conn.commit()
+    cur.close()
+
+    # and get back to submitting form, it should run migration
+    resp = resp.form.submit('submit')  # -> 3rd page
+    resp = resp.form.submit('submit')  # -> validation page
+    resp = resp.form.submit('submit')  # -> submit
+    assert formdef.data_class().count() == 1
+    formdata = formdef.data_class().select()[0]
+    assert formdata.data['1'] == 'test'
diff --git a/tests/utilities.py b/tests/utilities.py
index d4336190a..021d5a737 100644
--- a/tests/utilities.py
+++ b/tests/utilities.py
@@ -175,6 +175,7 @@ def create_temporary_pub(sql_mode=False, legacy_theme_mode=False, lazy_mode=Fals
         sql.do_tokens_table()
         sql.do_tracking_code_table()
         sql.do_session_table()
+        sql.do_transient_data_table()
         sql.do_custom_views_table()
         sql.do_snapshots_table()
         sql.do_loggederrors_table()
diff --git a/wcs/publisher.py b/wcs/publisher.py
index 3bbc567b8..7e30f84fd 100644
--- a/wcs/publisher.py
+++ b/wcs/publisher.py
@@ -358,6 +358,7 @@ class WcsPublisher(QommonPublisher):
         sql.do_role_table()
         sql.do_tracking_code_table()
         sql.do_custom_views_table()
+        sql.do_transient_data_table()
         sql.do_snapshots_table()
         sql.do_loggederrors_table()
         sql.do_tokens_table()
diff --git a/wcs/sql.py b/wcs/sql.py
index cbca3b599..65f5a8c6f 100644
--- a/wcs/sql.py
+++ b/wcs/sql.py
@@ -33,6 +33,7 @@ except ImportError:
     import pickle
 
 from django.utils.encoding import force_bytes, force_text
+from django.utils.timezone import now
 from quixote import get_publisher
 
 import wcs.carddata
@@ -1250,6 +1251,42 @@ def do_session_table():
     cur.close()
 
 
+def do_transient_data_table():
+    conn, cur = get_connection_and_cursor()
+    table_name = TransientData._table_name
+
+    cur.execute(
+        '''SELECT COUNT(*) FROM information_schema.tables
+                    WHERE table_schema = 'public'
+                      AND table_name = %s''',
+        (table_name,),
+    )
+    if cur.fetchone()[0] == 0:
+        cur.execute(
+            '''CREATE TABLE %s (id varchar PRIMARY KEY,
+                                session_id varchar,
+                                data bytea,
+                                last_update_time timestamptz
+                                )'''
+            % table_name
+        )
+    cur.execute(
+        '''SELECT column_name FROM information_schema.columns
+                    WHERE table_schema = 'public'
+                      AND table_name = %s''',
+        (table_name,),
+    )
+    existing_fields = {x[0] for x in cur.fetchall()}
+    needed_fields = {x[0] for x in TransientData._table_static_fields}
+
+    # delete obsolete fields
+    for field in existing_fields - needed_fields:
+        cur.execute('''ALTER TABLE %s DROP COLUMN %s''' % (table_name, field))
+
+    conn.commit()
+    cur.close()
+
+
 def do_custom_views_table():
     conn, cur = get_connection_and_cursor()
     table_name = 'custom_views'
@@ -3275,6 +3312,72 @@ class Role(SqlMixin, wcs.roles.Role):
         return o
 
 
+class TransientData(SqlMixin):
+    # table to keep some transient submission data out of global session dictionary
+    _table_name = 'transient_data'
+    _table_static_fields = [
+        ('id', 'varchar'),
+        ('session_id', 'varchar'),
+        ('data', 'bytea'),
+        ('last_update_time', 'timestamptz'),
+    ]
+    _numerical_id = False
+
+    def __init__(self, id, session_id, data):
+        self.id = id
+        self.session_id = session_id
+        self.data = data
+
+    @guard_postgres
+    def store(self):
+        sql_dict = {
+            'id': self.id,
+            'session_id': self.session_id,
+            'data': bytearray(pickle.dumps(self.data, protocol=2)),
+            'last_update_time': now(),
+        }
+
+        conn, cur = get_connection_and_cursor()
+        column_names = sql_dict.keys()
+        sql_statement = '''UPDATE %s SET %s WHERE id = %%(id)s RETURNING id''' % (
+            self._table_name,
+            ', '.join(['%s = %%(%s)s' % (x, x) for x in column_names]),
+        )
+        cur.execute(sql_statement, sql_dict)
+        if cur.fetchone() is None:
+            sql_statement = '''INSERT INTO %s (%s) VALUES (%s)''' % (
+                self._table_name,
+                ', '.join(column_names),
+                ', '.join(['%%(%s)s' % x for x in column_names]),
+            )
+            cur.execute(sql_statement, sql_dict)
+
+        conn.commit()
+        cur.close()
+
+    @classmethod
+    def _row2ob(cls, row, **kwargs):
+        o = cls.__new__(cls)
+        o.id = str_encode(row[0])
+        o.session_id = row[1]
+        o.data = pickle_loads(row[2])
+        return o
+
+    @classmethod
+    def get_data_fields(cls):
+        return []
+
+    @classmethod
+    @guard_postgres
+    def remove_for_session(cls, session_id):
+        conn, cur = get_connection_and_cursor()
+        sql_statement = 'DELETE FROM %s WHERE ' % cls._table_name
+        sql_statement += 'session_id = %s'
+        cur.execute(sql_statement, (session_id,))
+        conn.commit()
+        cur.close()
+
+
 class Session(SqlMixin, wcs.sessions.BasicSession):
     _table_name = 'sessions'
     _table_static_fields = [
@@ -3296,9 +3399,18 @@ class Session(SqlMixin, wcs.sessions.BasicSession):
         if self.message:
             # escape lazy gettext
             self.message = (self.message[0], str(self.message[1]))
+
+        # store transient data
+        for v in (self.magictokens or {}).values():
+            v.store()
+
+        # force to be empty, to make sure there's no leftover direct usage
+        session_data = copy.copy(self.__dict__)
+        session_data['magictokens'] = None
+
         sql_dict = {
             'id': self.id,
-            'session_data': bytearray(pickle.dumps(self.__dict__, protocol=2)),
+            'session_data': bytearray(pickle.dumps(session_data, protocol=2)),
             # the other fields are stored to run optimized SELECT() against the
             # table, they are ignored when loading the data.
             'name_identifier': self.name_identifier,
@@ -3333,6 +3445,12 @@ class Session(SqlMixin, wcs.sessions.BasicSession):
         session_data = pickle_loads(row[1])
         for k, v in session_data.items():
             setattr(o, k, v)
+        if o.magictokens:
+            # migration, obsolete storage of magictokens in session
+            for k, v in o.magictokens.items():
+                o.add_magictoken(k, v)
+            o.magictokens = None
+            o.store()
         return o
 
     @classmethod
@@ -3375,6 +3493,34 @@ class Session(SqlMixin, wcs.sessions.BasicSession):
     def get_data_fields(cls):
         return []
 
+    def add_magictoken(self, token, data):
+        super().add_magictoken(token, data)
+        if not self.id:
+            self.store()
+        self.magictokens[token] = TransientData(id=token, session_id=self.id, data=data)
+        self.magictokens[token].store()
+
+    def get_by_magictoken(self, token, default=None):
+        if not self.magictokens:
+            self.magictokens = {}
+        try:
+            if token not in self.magictokens:
+                self.magictokens[token] = TransientData.select(
+                    [Equal('session_id', self.id), Equal('id', token)]
+                )[0]
+            return self.magictokens[token].data
+        except IndexError:
+            return default
+
+    def remove_magictoken(self, token):
+        super().remove_magictoken(token)
+        TransientData.remove_object(token)
+
+    @classmethod
+    def remove_object(cls, id):
+        TransientData.remove_for_session(id)
+        super().remove_object(id)
+
 
 class TrackingCode(SqlMixin, wcs.tracking_code.TrackingCode):
     _table_name = 'tracking_codes'
@@ -4141,7 +4287,7 @@ def get_period_total(
 # latest migration, number + description (description is not used
 # programmaticaly but will make sure git conflicts if two migrations are
 # separately added with the same number)
-SQL_LEVEL = (63, 'add index on snapshot table')
+SQL_LEVEL = (64, 'add transient data table')
 
 
 def migrate_global_views(conn, cur):
@@ -4306,6 +4452,9 @@ def migrate():
         # 37: create custom_views table
         # 44: add is_default column to custom_views table
         do_custom_views_table()
+    if sql_level < 64:
+        # 64: add transient data table
+        do_transient_data_table()
     if sql_level < 30:
         # 30: actually remove evo.who on anonymised formdatas
         from wcs.formdef import FormDef
-- 
2.36.1

