From 38cc67385fb1b36aa0881bc5982bc58d75dac464 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Luk=C3=A1=C5=A1=20Hr=C3=A1zk=C3=BD?= <lhrazky@redhat.com>
Date: Wed, 11 Nov 2020 18:45:11 +0100
Subject: [PATCH] Package: add a get_header() method (RhBug:1876606)

Adds get_header() method to the Package class, which returns the rpm
header of an installed package.

= changelog =
msg: Add get_header() method to the Package class
type: enhancement
resolves: https://bugzilla.redhat.com/show_bug.cgi?id=1876606
---
 dnf/package.py        | 24 ++++++++++++++++++++++++
 tests/test_package.py | 12 ++++++++++++
 2 files changed, 36 insertions(+)

diff --git a/dnf/package.py b/dnf/package.py
index baef04fa5b..836e0e4989 100644
--- a/dnf/package.py
+++ b/dnf/package.py
@@ -26,11 +26,13 @@
 from dnf.i18n import _
 
 import binascii
+import dnf.exceptions
 import dnf.rpm
 import dnf.yum.misc
 import hawkey
 import logging
 import os
+import rpm
 
 logger = logging.getLogger("dnf")
 
@@ -95,6 +97,11 @@ def from_repo(self):
 
     @property
     def _header(self):
+        """
+        Returns the header of a locally present rpm package file. As opposed to
+        self.get_header(), which retrieves the header of an installed package
+        from rpmdb.
+        """
         return dnf.rpm._header(self.localPkg())
 
     @property
@@ -164,6 +171,23 @@ def debugsource_name(self):
         src_name = self.source_name if self.source_name is not None else self.name
         return src_name + self.DEBUGSOURCE_SUFFIX
 
+    def get_header(self):
+        """
+        Returns the rpm header of the package if it is installed. If not
+        installed, returns None. The header is not cached, it is retrieved from
+        rpmdb on every call. In case of a failure (e.g. when the rpmdb changes
+        between loading the data and calling this method), raises an instance
+        of PackageNotFoundError.
+        """
+        if not self._from_system:
+            return None
+
+        try:
+            # RPMDBI_PACKAGES stands for the header of the package
+            return next(self.base._ts.dbMatch(rpm.RPMDBI_PACKAGES, self.rpmdbid))
+        except StopIteration:
+            raise dnf.exceptions.PackageNotFoundError("Package not found when attempting to retrieve header", str(self))
+
     @property
     def source_debug_name(self):
         # :api
diff --git a/tests/test_package.py b/tests/test_package.py
index cd4872e631..514e5bf099 100644
--- a/tests/test_package.py
+++ b/tests/test_package.py
@@ -68,6 +68,18 @@ def fn_getter():
             with self.assertRaises(IOError):
                 pkg._header
 
+    # rpm.hdr() is not easy to construct with custom data, we just return a string
+    # instead, as we don't actually need an instance of rpm.hdr for the test
+    @mock.patch("rpm.TransactionSet.dbMatch", lambda self, a, b: iter(["package_header_test_data"]))
+    def test_get_header(self):
+        pkg = self.sack.query().installed().filter(name="pepper")[0]
+        header = pkg.get_header()
+        self.assertEqual(header, "package_header_test_data")
+
+        pkg = self.sack.query().available().filter(name="pepper")[0]
+        header = pkg.get_header()
+        self.assertEqual(header, None)
+
     @mock.patch("dnf.package.Package.rpmdbid", long(3))
     def test_idx(self):
         """ pkg.idx is an int. """