commit:     d121ea57ed5310d84328be27f10a8556b0a7d7ba
Author:     Zac Medico <zmedico <AT> gentoo <DOT> org>
AuthorDate: Fri May  8 23:32:49 2020 +0000
Commit:     Zac Medico <zmedico <AT> gentoo <DOT> org>
CommitDate: Wed Feb 24 15:27:47 2021 +0000
URL:        https://gitweb.gentoo.org/proj/portage.git/commit/?id=d121ea57

Add emirrordist shelve dump/restore (bug 721680)

Bug: https://bugs.gentoo.org/721680
Signed-off-by: Zac Medico <zmedico <AT> gentoo.org>

 bin/shelve-utils                      | 32 +++++++++++++++++++
 lib/portage/tests/util/test_shelve.py | 60 +++++++++++++++++++++++++++++++++++
 lib/portage/util/shelve.py            | 58 +++++++++++++++++++++++++++++++++
 3 files changed, 150 insertions(+)

diff --git a/bin/shelve-utils b/bin/shelve-utils
new file mode 100755
index 000000000..5088ab5eb
--- /dev/null
+++ b/bin/shelve-utils
@@ -0,0 +1,32 @@
+#!/usr/bin/python -b
+# Copyright 2020-2021 Gentoo Authors
+# Distributed under the terms of the GNU General Public License v2
+
+import argparse
+import sys
+
+import portage
+from portage.util.shelve import dump, restore
+
+
+def main(argv=None):
+       parser = argparse.ArgumentParser(prog="shelve-utils")
+       subparsers = parser.add_subparsers(help="sub-command help")
+
+       dump_command = subparsers.add_parser("dump", help="dump shelve 
database")
+       dump_command.add_argument("src", help="input shelve file")
+       dump_command.add_argument("dest", help="output pickle file")
+       dump_command.set_defaults(func=dump)
+
+       restore_command = subparsers.add_parser("restore", help="restore shelve 
database")
+       restore_command.add_argument("src", help="input pickle file")
+       restore_command.add_argument("dest", help="output shelve file")
+       restore_command.set_defaults(func=restore)
+
+       args = parser.parse_args(args=portage._decode_argv(argv or 
sys.argv)[1:])
+       args.func(args)
+
+
+if __name__ == "__main__":
+       portage.util.initialize_logger()
+       main(argv=sys.argv)

diff --git a/lib/portage/tests/util/test_shelve.py 
b/lib/portage/tests/util/test_shelve.py
new file mode 100644
index 000000000..60592c6fb
--- /dev/null
+++ b/lib/portage/tests/util/test_shelve.py
@@ -0,0 +1,60 @@
+# Copyright 2020-2021 Gentoo Authors
+# Distributed under the terms of the GNU General Public License v2
+
+import argparse
+import os
+import shutil
+import tempfile
+import time
+
+from portage.tests import TestCase
+from portage.util.shelve import dump, open_shelve, restore
+
+
+class ShelveUtilsTestCase(TestCase):
+
+       TEST_DATA = (
+               # distfiles_db
+               {
+                       "portage-2.3.89.tar.bz2": "sys-apps/portage-2.3.89",
+                       "portage-2.3.99.tar.bz2": "sys-apps/portage-2.3.99",
+               },
+               # deletion_db
+               {
+                       "portage-2.3.89.tar.bz2": time.time(),
+                       "portage-2.3.99.tar.bz2": time.time(),
+               },
+               # recycle_db
+               {
+                       "portage-2.3.89.tar.bz2": (0, time.time()),
+                       "portage-2.3.99.tar.bz2": (0, time.time()),
+               },
+       )
+
+       def test_dump_restore(self):
+               for data in self.TEST_DATA:
+                       tmpdir = tempfile.mkdtemp()
+                       try:
+                               dump_args = argparse.Namespace(
+                                       src=os.path.join(tmpdir, "shelve_file"),
+                                       dest=os.path.join(tmpdir, 
"pickle_file"),
+                               )
+                               db = open_shelve(dump_args.src, flag="c")
+                               for k, v in data.items():
+                                       db[k] = v
+                               db.close()
+                               dump(dump_args)
+
+                               os.unlink(dump_args.src)
+                               restore_args = argparse.Namespace(
+                                       dest=dump_args.src,
+                                       src=dump_args.dest,
+                               )
+                               restore(restore_args)
+
+                               db = open_shelve(restore_args.dest, flag="r")
+                               for k, v in data.items():
+                                       self.assertEqual(db[k], v)
+                               db.close()
+                       finally:
+                               shutil.rmtree(tmpdir)

diff --git a/lib/portage/util/shelve.py b/lib/portage/util/shelve.py
new file mode 100644
index 000000000..f070ee753
--- /dev/null
+++ b/lib/portage/util/shelve.py
@@ -0,0 +1,58 @@
+# Copyright 2020-2021 Gentoo Authors
+# Distributed under the terms of the GNU General Public License v2
+
+import logging
+import pickle
+import shelve
+
+
+def open_shelve(db_file, flag="r"):
+       """
+       The optional flag parameter has the same interpretation as the flag
+       parameter of dbm.open()
+       """
+       try:
+               db = shelve.open(db_file, flag=flag)
+       except ImportError as e:
+               # ImportError has different attributes for python2 vs. python3
+               if getattr(e, "name", None) == "bsddb" or getattr(e, "message", 
None) in (
+                       "No module named bsddb",
+                       "No module named _bsddb",
+               ):
+                       from bsddb3 import dbshelve
+
+                       db = dbshelve.open(db_file)
+               else:
+                       raise
+
+       return db
+
+
+def dump(args):
+       src = open_shelve(args.src, flag="r")
+       try:
+               with open(args.dest, "wb") as dest:
+                       for key in src:
+                               try:
+                                       value = src[key]
+                               except KeyError:
+                                       logging.exception(key)
+                                       continue
+                               pickle.dump((key, value), dest)
+       finally:
+               src.close()
+
+
+def restore(args):
+       dest = open_shelve(args.dest, flag="c")
+       try:
+               with open(args.src, "rb") as src:
+                       while True:
+                               try:
+                                       k, v = pickle.load(src)
+                               except EOFError:
+                                       break
+                               else:
+                                       dest[k] = v
+       finally:
+               dest.close()

Reply via email to