Updating Partition() information after mount and unmount. (#1508)

* Updating Partition() information after mount and unmount.

* Cleaned up raw_parted() to gracefully output relevant partition error information.

* Adding timestmap to cmd_history.txt as it's impossible to debug properly otherwise

* Adding more verbose debugging information

* Reinstating the lsblk retry code for PARTUUID and UUID on Partition()'s information

* Added error handling for JSON parsing

3 files changed, 67 insertions, 15 deletions

diff --git a/archinstall/lib/disk/blockdevice.py b/archinstall/lib/disk/blockdevice.py
index 736bacbc..cd0e14a6 100644
--- a/archinstall/lib/disk/blockdevice.py
+++ b/archinstall/lib/disk/blockdevice.py
@@ -274,12 +274,16 @@ class BlockDevice:
 		if not uuid and not partuuid:
 			raise ValueError(f"BlockDevice.get_partition() requires either a UUID or a PARTUUID for lookups.")
 
+		log(f"Retrieving partition PARTUUID={partuuid} or UUID={uuid}", level=logging.INFO, fg="teal")
+
 		for count in range(storage.get('DISK_RETRY_ATTEMPTS', 5)):
 			for partition_index, partition in self.partitions.items():
 				try:
 					if uuid and partition.uuid and partition.uuid.lower() == uuid.lower():
+						log(f"Matched UUID={uuid} against {partition.uuid}", level=logging.INFO, fg="teal")
 						return partition
 					elif partuuid and partition.part_uuid and partition.part_uuid.lower() == partuuid.lower():
+						log(f"Matched PARTUUID={partuuid} against {partition.part_uuid}", level=logging.INFO, fg="teal")
 						return partition
 				except DiskError as error:
 					# Most likely a blockdevice that doesn't support or use UUID's
diff --git a/archinstall/lib/disk/filesystem.py b/archinstall/lib/disk/filesystem.py
index 5d5952a0..c8eaf0be 100644
--- a/archinstall/lib/disk/filesystem.py
+++ b/archinstall/lib/disk/filesystem.py
@@ -189,10 +189,13 @@ class Filesystem:
 		return True
 
 	def raw_parted(self, string: str) -> SysCommand:
-		if (cmd_handle := SysCommand(f'/usr/bin/parted -s {string}')).exit_code != 0:
-			log(f"Parted ended with a bad exit code: {cmd_handle}", level=logging.ERROR, fg="red")
-		time.sleep(0.5)
-		return cmd_handle
+		try:
+			cmd_handle = SysCommand(f'/usr/bin/parted -s {string}')
+			time.sleep(0.5)
+			return cmd_handle
+		except SysCallError as error:
+			log(f"Parted ended with a bad exit code: {error.exit_code} ({error})", level=logging.ERROR, fg="red")
+			return error
 
 	def parted(self, string: str) -> bool:
 		"""
@@ -258,6 +261,9 @@ class Filesystem:
 				new_partition_uuids = [partition.part_uuid for partition in self.blockdevice.partitions.values()]
 				new_partuuid_set = (set(previous_partuuids) ^ set(new_partition_uuids))
 
+				log(f'Old partition set: {previous_partuuids}', level=logging.INFO, fg="teal")
+				log(f'New partition set: {new_partition_uuids}', level=logging.INFO, fg="teal")
+
 				if len(new_partuuid_set) and (new_partuuid := new_partuuid_set.pop()):
 					try:
 						return self.blockdevice.get_partition(partuuid=new_partuuid)
@@ -282,6 +288,7 @@ class Filesystem:
 		log(f"Could not find the new PARTUUID after adding the partition.", level=logging.ERROR, fg="red")
 		log(f"Previous partitions: {previous_partuuids}", level=logging.ERROR, fg="red")
 		log(f"New partitions: {total_partitions}", level=logging.ERROR, fg="red")
+
 		raise DiskError(f"Could not add partition using: {parted_string}")
 
 	def set_name(self, partition: int, name: str) -> bool:
diff --git a/archinstall/lib/disk/partition.py b/archinstall/lib/disk/partition.py
index 56a7d436..04d33453 100644
--- a/archinstall/lib/disk/partition.py
+++ b/archinstall/lib/disk/partition.py
@@ -5,7 +5,7 @@ import json
 import os
 import hashlib
 import typing
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from pathlib import Path
 from typing import Optional, Dict, Any, List, Union, Iterator
 
@@ -18,19 +18,51 @@ from ..general import SysCommand
 from .btrfs.btrfs_helpers import subvolume_info_from_path
 from .btrfs.btrfssubvolumeinfo import BtrfsSubvolumeInfo
 
-
 @dataclass
 class PartitionInfo:
-	pttype: str
-	partuuid: str
-	uuid: str
-	start: Optional[int]
-	end: Optional[int]
+	partition_object: 'Partition'
+	device_path: str # This would be /dev/sda1 for instance
 	bootable: bool
 	size: float
 	sector_size: int
-	filesystem_type: str
-	mountpoints: List[Path]
+	start: Optional[int]
+	end: Optional[int]
+	pttype: Optional[str]
+	filesystem_type: Optional[str]
+	partuuid: Optional[str]
+	uuid: Optional[str]
+	mountpoints: List[Path] = field(default_factory=list)
+
+	def __post_init__(self):
+		if not all([self.partuuid, self.uuid]):
+			for i in range(storage['DISK_RETRY_ATTEMPTS']):
+				lsblk_info = SysCommand(f"lsblk --json -b -o+LOG-SEC,SIZE,PTTYPE,PARTUUID,UUID,FSTYPE {self.device_path}").decode('UTF-8')
+				try:
+					lsblk_info = json.loads(lsblk_info)
+				except json.decoder.JSONDecodeError:
+					log(f"Could not decode JSON: {lsblk_info}", fg="red", level=logging.ERROR)
+					raise DiskError(f'Failed to retrieve information for "{self.device_path}" with lsblk')
+
+				if not (device := lsblk_info.get('blockdevices', [None])[0]):
+					raise DiskError(f'Failed to retrieve information for "{self.device_path}" with lsblk')
+
+				self.partuuid = device.get('partuuid')
+				self.uuid = device.get('uuid')
+
+				# Lets build a list of requirements that we would like
+				# to retry and build (stuff that can take time between partprobes)
+				requirements = []
+				requirements.append(self.partuuid)
+
+				# Unformatted partitions won't have a UUID
+				if lsblk_info.get('fstype') is not None:
+					requirements.append(self.uuid)
+
+				if all(requirements):
+					break
+
+				self.partition_object.partprobe()
+				time.sleep(max(0.1, storage['DISK_TIMEOUTS'] * i))
 
 	def get_first_mountpoint(self) -> Optional[Path]:
 		if len(self.mountpoints) > 0:
@@ -154,8 +186,11 @@ class Partition:
 			output = error.worker.decode('UTF-8')
 
 		if output:
-			lsblk_info = json.loads(output)
-			return lsblk_info
+			try:
+				lsblk_info = json.loads(output)
+				return lsblk_info
+			except json.decoder.JSONDecodeError:
+				log(f"Could not decode JSON: {output}", fg="red", level=logging.ERROR)
 
 		raise DiskError(f'Failed to read disk "{self.device_path}" with lsblk')
 
@@ -185,6 +220,8 @@ class Partition:
 		bootable = sfdisk_info.get('bootable', False) or sfdisk_info.get('type', '') == 'C12A7328-F81F-11D2-BA4B-00A0C93EC93B'
 
 		return PartitionInfo(
+			partition_object=self,
+			device_path=self._path,
 			pttype=device['pttype'],
 			partuuid=device['partuuid'],
 			uuid=device['uuid'],
@@ -568,6 +605,8 @@ class Partition:
 			except SysCallError as err:
 				raise err
 
+			# Update the partition info since the mount info has changed after this call.
+			self._partition_info = self._fetch_information()
 			return True
 
 		return False
@@ -582,6 +621,8 @@ class Partition:
 		if exit_code and 0 < exit_code < 8000:
 			raise SysCallError(f"Could not unmount {self._path} properly: {worker}", exit_code=exit_code)
 
+		# Update the partition info since the mount info has changed after this call.
+		self._partition_info = self._fetch_information()
 		return True
 
 	def filesystem_supported(self) -> bool: