Merge ~cjwatson/launchpad:refactor-pgsession-queries into launchpad:master

[Colin Watson]

Refactor PGSession queries to reduce manual SQL

Pushing the special handling for old session pickles saved by Python 2 down to the `SessionPkgData` model and adding some named function declarations allows us to make quite effective use of Storm's query compiler, eliminating some manually-constructed SQL and giving us better type annotations.

It's possible that at least the dump compatibility with Python 2 can now be removed, but I deliberately haven't attempted that in this commit: this should be a pure refactoring with no functional change.

Unmerged commits

1fe67cd... by Colin Watson on 2023-02-23

Refactor PGSession queries to reduce manual SQL

Pushing the special handling for old session pickles saved by Python 2
down to the `SessionPkgData` model and adding some named function
declarations allows us to make quite effective use of Storm's query
compiler, eliminating some manually-constructed SQL and giving us better
type annotations.

It's possible that at least the dump compatibility with Python 2 can now
be removed, but I deliberately haven't attempted that in this commit:
this should be a pure refactoring with no functional change.

	docs:0 (build)
	lint:0 (build)
	mypy:0 (build)

1 → 3 of 3 results

First • Previous • Next • Last

diff --git a/lib/lp/services/session/adapters.py b/lib/lp/services/session/adapters.py
index f40145c..ae1e342 100644
--- a/lib/lp/services/session/adapters.py
+++ b/lib/lp/services/session/adapters.py
@@ -3,8 +3,9 @@
 
 """Session adapters."""
 
-__all__ = []
+__all__ = []  # type: List[str]
 
+from typing import List
 
 from zope.component import adapter
 from zope.interface import implementer
diff --git a/lib/lp/services/session/model.py b/lib/lp/services/session/model.py
index b54d290..bd83e96 100644
--- a/lib/lp/services/session/model.py
+++ b/lib/lp/services/session/model.py
@@ -3,9 +3,21 @@
 
 """Session Storm database classes"""
 
-__all__ = ["SessionData", "SessionPkgData"]
+__all__ = [
+    "EnsureSessionClientID",
+    "SessionData",
+    "SessionPkgData",
+    "SetSessionPkgData",
+]
 
-from storm.locals import Pickle, Unicode
+import io
+import pickle
+from datetime import datetime
+from typing import Any
+
+from storm.expr import NamedFunc
+from storm.properties import SimpleProperty, Unicode
+from storm.variables import PickleVariable
 from zope.interface import implementer, provider
 
 from lp.services.database.datetimecol import UtcDateTimeCol
@@ -24,6 +36,46 @@ class SessionData(StormBase):
     last_accessed = UtcDateTimeCol()
 
 
+class Python2FriendlyUnpickler(pickle._Unpickler):
+    """An unpickler that handles Python 2 datetime objects.
+
+    Python 3 versions before 3.6 fail to unpickle Python 2 datetime objects
+    (https://bugs.python.org/issue22005); even in Python >= 3.6 they require
+    passing a different encoding to pickle.loads, which may have undesirable
+    effects on other objects being unpickled.  Work around this by instead
+    patching in a different encoding just for the argument to
+    datetime.datetime.
+    """
+
+    def find_class(self, module: str, name: str) -> Any:
+        if module == "datetime" and name == "datetime":
+            original_encoding = self.encoding  # type: ignore[has-type]
+            self.encoding = "bytes"
+
+            def datetime_factory(pickle_data):
+                self.encoding = original_encoding
+                return datetime(pickle_data)
+
+            return datetime_factory
+        else:
+            return super().find_class(module, name)
+
+
+class Python2FriendlyPickleVariable(PickleVariable):
+    def _loads(self, value: bytes) -> object:
+        return Python2FriendlyUnpickler(io.BytesIO(value)).load()
+
+
+class Python2FriendlyPickle(SimpleProperty):
+    """Like `storm.properties.Pickle`, but also handles data from Python 2.
+
+    We need this even though Launchpad no longer runs on Python 2, because
+    the session database may contain old data.
+    """
+
+    variable_class = Python2FriendlyPickleVariable
+
+
 @implementer(IUseSessionStore)
 @provider(IUseSessionStore)
 class SessionPkgData(StormBase):
@@ -35,4 +87,24 @@ class SessionPkgData(StormBase):
     client_id = Unicode()
     product_id = Unicode()
     key = Unicode()
-    pickle = Pickle()
+    pickle = Python2FriendlyPickle()
+
+
+class EnsureSessionClientID(NamedFunc):
+    __slots__ = ()
+    name = "ensure_session_client_id"
+
+    def __init__(self, hashed_client_id: str) -> None:
+        super().__init__(hashed_client_id)
+
+
+class SetSessionPkgData(NamedFunc):
+    __slots__ = ()
+    name = "set_session_pkg_data"
+
+    def __init__(
+        self, hashed_client_id: str, product_id: str, key: str, value: object
+    ) -> None:
+        # Use protocol 2 for Python 2 compatibility.
+        pickled_value = pickle.dumps(value, protocol=2)
+        super().__init__(hashed_client_id, product_id, key, pickled_value)
diff --git a/lib/lp/services/webapp/pgsession.py b/lib/lp/services/webapp/pgsession.py
index 1f25642..5a0fa5d 100644
--- a/lib/lp/services/webapp/pgsession.py
+++ b/lib/lp/services/webapp/pgsession.py
@@ -4,18 +4,24 @@
 """PostgreSQL server side session storage for Zope3."""
 
 import hashlib
-import io
-import pickle
 from collections.abc import MutableMapping
-from datetime import datetime
+from datetime import timedelta
 
 import six
 from lazr.restful.utils import get_current_browser_request
+from storm.expr import Cast, Select
 from storm.zope.interfaces import IZStorm
 from zope.authentication.interfaces import IUnauthenticatedPrincipal
 from zope.component import getUtility
 from zope.interface import implementer
 
+from lp.services.database.constants import UTC_NOW
+from lp.services.session.model import (
+    EnsureSessionClientID,
+    SessionData,
+    SessionPkgData,
+    SetSessionPkgData,
+)
 from lp.services.webapp.interfaces import (
     IClientIdManager,
     ISessionData,
@@ -29,31 +35,6 @@ HOURS = 60 * MINUTES
 DAYS = 24 * HOURS
 
 
-class Python2FriendlyUnpickler(pickle._Unpickler):
-    """An unpickler that handles Python 2 datetime objects.
-
-    Python 3 versions before 3.6 fail to unpickle Python 2 datetime objects
-    (https://bugs.python.org/issue22005); even in Python >= 3.6 they require
-    passing a different encoding to pickle.loads, which may have undesirable
-    effects on other objects being unpickled.  Work around this by instead
-    patching in a different encoding just for the argument to
-    datetime.datetime.
-    """
-
-    def find_class(self, module, name):
-        if module == "datetime" and name == "datetime":
-            original_encoding = self.encoding
-            self.encoding = "bytes"
-
-            def datetime_factory(pickle_data):
-                self.encoding = original_encoding
-                return datetime(pickle_data)
-
-            return datetime_factory
-        else:
-            return super().find_class(module, name)
-
-
 class PGSessionBase:
     store_name = "session"
 
@@ -90,9 +71,6 @@ class PGSessionDataContainer(PGSessionBase):
     # using the session data.
     resolution = 9 * MINUTES
 
-    session_data_table_name = "SessionData"
-    session_pkg_data_table_name = "SessionPkgData"
-
     def __getitem__(self, client_id):
         """See `ISessionDataContainer`."""
         return PGSessionData(self, client_id)
@@ -119,16 +97,16 @@ class PGSessionData(PGSessionBase):
         ).hexdigest()
 
         # Update the last access time in the db if it is out of date
-        table_name = session_data_container.session_data_table_name
-        query = """
-            UPDATE %s SET last_accessed = CURRENT_TIMESTAMP
-            WHERE client_id = ?
-                AND last_accessed < CURRENT_TIMESTAMP - '%d seconds'::interval
-            """ % (
-            table_name,
-            session_data_container.resolution,
-        )
-        self.store.execute(query, (self.hashed_client_id,), noresult=True)
+        self.store.find(
+            SessionData,
+            SessionData.client_id == self.hashed_client_id,
+            SessionData.last_accessed
+            < UTC_NOW
+            - Cast(
+                timedelta(seconds=session_data_container.resolution),
+                "interval",
+            ),
+        ).set(last_accessed=UTC_NOW)
 
     def _ensureClientId(self):
         if self._have_ensured_client_id:
@@ -137,8 +115,7 @@ class PGSessionData(PGSessionBase):
         # about our client id. We're doing it lazily to try and keep anonymous
         # users from having a session.
         self.store.execute(
-            "SELECT ensure_session_client_id(?)",
-            (self.hashed_client_id,),
+            Select(EnsureSessionClientID(self.hashed_client_id)),
             noresult=True,
         )
         request = get_current_browser_request()
@@ -198,29 +175,18 @@ class PGSessionPkgData(MutableMapping, PGSessionBase):
     def __init__(self, session_data, product_id):
         self.session_data = session_data
         self.product_id = six.ensure_text(product_id, "ascii")
-        self.table_name = (
-            session_data.session_data_container.session_pkg_data_table_name
-        )
         self._populate()
 
     _data_cache = None
 
     def _populate(self):
         self._data_cache = {}
-        query = (
-            """
-            SELECT key, pickle FROM %s WHERE client_id = ?
-                AND product_id = ?
-            """
-            % self.table_name
-        )
-        result = self.store.execute(
-            query, (self.session_data.hashed_client_id, self.product_id)
+        result = self.store.find(
+            (SessionPkgData.key, SessionPkgData.pickle),
+            SessionPkgData.client_id == self.session_data.hashed_client_id,
+            SessionPkgData.product_id == self.product_id,
         )
-        for key, pickled_value in result:
-            value = Python2FriendlyUnpickler(
-                io.BytesIO(bytes(pickled_value))
-            ).load()
+        for key, value in result:
             self._data_cache[key] = value
 
     def __getitem__(self, key):
@@ -228,17 +194,16 @@ class PGSessionPkgData(MutableMapping, PGSessionBase):
 
     def __setitem__(self, key, value):
         key = six.ensure_text(key, "ascii")
-        # Use protocol 2 for Python 2 compatibility.
-        pickled_value = pickle.dumps(value, protocol=2)
 
         self.session_data._ensureClientId()
         self.store.execute(
-            "SELECT set_session_pkg_data(?, ?, ?, ?)",
-            (
-                self.session_data.hashed_client_id,
-                self.product_id,
-                key,
-                pickled_value,
+            Select(
+                SetSessionPkgData(
+                    self.session_data.hashed_client_id,
+                    self.product_id,
+                    key,
+                    value,
+                )
             ),
             noresult=True,
         )
@@ -259,22 +224,12 @@ class PGSessionPkgData(MutableMapping, PGSessionBase):
             # another process has inserted it and we should keep our grubby
             # fingers out of it.
             return
-        query = (
-            """
-            DELETE FROM %s
-            WHERE client_id = ? AND product_id = ? AND key = ?
-            """
-            % self.table_name
-        )
-        self.store.execute(
-            query,
-            (
-                self.session_data.hashed_client_id,
-                self.product_id,
-                six.ensure_text(key, "ascii"),
-            ),
-            noresult=True,
-        )
+        self.store.find(
+            SessionPkgData,
+            SessionPkgData.client_id == self.session_data.hashed_client_id,
+            SessionPkgData.product_id == self.product_id,
+            SessionPkgData.key == six.ensure_text(key, "ascii"),
+        ).remove()
 
     def __iter__(self):
         return iter(self._data_cache)
diff --git a/tox.ini b/tox.ini
index 9054ea2..9319722 100644
--- a/tox.ini
+++ b/tox.ini
@@ -27,7 +27,7 @@ commands_pre =
     {toxinidir}/scripts/update-version-info.sh
 commands =
     mypy --follow-imports=silent \
-    {posargs:lib/lp/answers lib/lp/app lib/lp/archivepublisher lib/lp/archiveuploader lib/lp/buildmaster lib/lp/charms/model/charmrecipebuildbehaviour.py lib/lp/code/model/cibuildbehaviour.py lib/lp/code/model/recipebuilder.py lib/lp/oci/model/ocirecipebuildbehaviour.py lib/lp/snappy/model/snapbuildbehaviour.py lib/lp/soyuz/model/binarypackagebuildbehaviour.py lib/lp/soyuz/model/livefsbuildbehaviour.py lib/lp/testing lib/lp/translations/model/translationtemplatesbuildbehaviour.py}
+    {posargs:lib/lp/answers lib/lp/app lib/lp/archivepublisher lib/lp/archiveuploader lib/lp/buildmaster lib/lp/charms/model/charmrecipebuildbehaviour.py lib/lp/code/model/cibuildbehaviour.py lib/lp/code/model/recipebuilder.py lib/lp/oci/model/ocirecipebuildbehaviour.py lib/lp/services/session lib/lp/services/webapp/pgsession.py lib/lp/snappy/model/snapbuildbehaviour.py lib/lp/soyuz/model/binarypackagebuildbehaviour.py lib/lp/soyuz/model/livefsbuildbehaviour.py lib/lp/testing lib/lp/translations/model/translationtemplatesbuildbehaviour.py}
 
 [testenv:docs]
 basepython = python3