From b4a634a443e08bba67430720ffbd8fb50e3e62f7 Mon Sep 17 00:00:00 2001
From: Sebastian Mohr <39738318+semohr@users.noreply.github.com>
Date: Tue, 15 Apr 2025 11:43:36 +0200
Subject: [PATCH] Allow to pickle db models by removing the current connection.
 (#5641)

## Description

This might be a quick one, depending on how you feel about it... It
allows you to pickle DB model objects. I don't think this is used
directly in Beets, but it might be useful in general. For instance, we
encountered an issue where we wanted to quickly pickle an Item or Album.
This sometimes worked and other times failed, which seemed quite
inconsistent.

Some DB model methods and properties have the side effect of attaching
an SQLite connection to self (._db), which prevents serialization. The
fix is quite straightforward, so I thought we might want to integrate
this into beets directly.

## To Do

- [x] Changelog
- [x] Tests
---
 beets/dbcore/db.py  |  9 +++++++++
 docs/changelog.rst  |  1 +
 test/test_dbcore.py | 14 ++++++++++++++
 3 files changed, 24 insertions(+)

diff --git a/beets/dbcore/db.py b/beets/dbcore/db.py
index 2aa0081d7..dd8401935 100755
--- a/beets/dbcore/db.py
+++ b/beets/dbcore/db.py
@@ -716,6 +716,15 @@ class Model(ABC, Generic[D]):
         """Set the object's key to a value represented by a string."""
         self[key] = self._parse(key, string)
 
+    def __getstate__(self):
+        """Return the state of the object for pickling.
+        Remove the database connection as sqlite connections are not
+        picklable.
+        """
+        state = self.__dict__.copy()
+        state["_db"] = None
+        return state
+
 
 # Database controller and supporting interfaces.
 
diff --git a/docs/changelog.rst b/docs/changelog.rst
index ca2bb50e3..77d237e6d 100644
--- a/docs/changelog.rst
+++ b/docs/changelog.rst
@@ -102,6 +102,7 @@ Other changes:
   EXTM3U playlists instead of JSON-encoding them.
 * :doc:`plugins/ftintitle`: Optimize the plugin by avoiding unnecessary writes
   to the database.
+* Database models are now serializable with pickle.
 
 2.2.0 (December 02, 2024)
 -------------------------
diff --git a/test/test_dbcore.py b/test/test_dbcore.py
index 2ff20c3a3..a4bae97c9 100644
--- a/test/test_dbcore.py
+++ b/test/test_dbcore.py
@@ -421,6 +421,20 @@ class ModelTest(unittest.TestCase):
         with pytest.raises(TypeError, match="must be a string"):
             dbcore.Model._parse(None, 42)
 
+    def test_pickle_dump(self):
+        """Tries to pickle an item. This tests the __getstate__ method
+        of the Model ABC"""
+        import pickle
+
+        model = ModelFixture1(self.db)
+        model.add(self.db)
+        model.field_one = 123
+
+        model.store()
+        assert model._db is not None
+
+        pickle.dumps(model)
+
 
 class FormatTest(unittest.TestCase):
     def test_format_fixed_field_integer(self):