diff --git a/classes/config.py b/classes/config.py index 6d2d9b3..284f592 100644 --- a/classes/config.py +++ b/classes/config.py @@ -5,15 +5,18 @@ from typing import Union from classes.vessel import Vessel from classes.directory import Directory +from classes.database import Database class MonsterConfig: - def readFile(self, path: Union[str, Path]) -> None: + def readFile(self, path: Union[str, Path], dbclass: type = Database) -> None: """Read .ini file into MonsterConfig object Args: path (str, pathlib.Path): Location of the .ini file to read (absolute or relative to the working directory) + dbclass (type): Class to use for database connections. Defaults to + built-in Database using sqlite3. Raises: ValueError: Raised if the passed file is not a ContentMonster .ini @@ -43,12 +46,13 @@ class MonsterConfig: # Read Vessels from the config file elif section.startswith("Vessel"): - self.vessels.append(Vessel.fromConfig(parser[section])) + self.vessels.append( + Vessel.fromConfig(parser[section], dbclass)) def __init__(self) -> None: """Initialize a new (empty) MonsterConfig object """ self.directories = [] self.vessels = [] - self.chunksize = 10485760 # Default: 10 MiB - self.database = None # Default: "database.sqlite3" in base directory + self.chunksize = 10485760 # Default: 10 MiB + self.database = None # Default: "database.sqlite3" in base directory diff --git a/classes/file.py b/classes/file.py index 395eee5..69ac1a0 100644 --- a/classes/file.py +++ b/classes/file.py @@ -11,7 +11,7 @@ class File: """Object representing a file found in a local Directory """ - def __init__(self, name: str, directory, uuid: Optional[str] = None) -> None: + def __init__(self, name: str, directory, uuid: Optional[str] = None, dbclass: type = Database) -> None: """Initialize new File object Args: @@ -20,12 +20,15 @@ class File: is located within uuid (str, optional): Unique identifier of this File object. Will be retrieved from database if None. Defaults to None. + dbclass (type): Class to use for database connections. Defaults to + built-in Database using sqlite3. Raises: FileNotFoundError: Raised if the specified File does not exist """ self.name = name self.directory = directory + self.dbclass = dbclass if not self.exists(): raise FileNotFoundError(f"File {self.name} does not exist in {self.directory.name}!") @@ -47,9 +50,9 @@ class File: """Return unique identifier for this File object Returns: - str: File object's UUID retrieved from Database + str: File object's UUID retrieved from database """ - db = Database() + db = self.dbclass() return db.getFileUUID(self) def getFullPath(self) -> str: diff --git a/classes/shorethread.py b/classes/shorethread.py index e18ee86..2ae34c1 100644 --- a/classes/shorethread.py +++ b/classes/shorethread.py @@ -16,7 +16,7 @@ import os.path class ShoreThread(Process): """Thread handling the discovery of shore-side file changes """ - def __init__(self, state: dict) -> None: + def __init__(self, state: dict, dbclass: type = Database) -> None: """Create a new ShoreThread object Args: @@ -27,6 +27,7 @@ class ShoreThread(Process): self._state = state self.queue = Queue() self._logger = Logger() + self._dbclass = dbclass def getAllFiles(self) -> list: """Return File objects for all files in all Directories @@ -107,7 +108,7 @@ class ShoreThread(Process): # Remove file from database self._logger.debug(f"Purging file {name} from database") - db = Database() + db = self._dbclass() db.removeFile(directory, name) def addFile(self, fileobj): @@ -142,7 +143,7 @@ class ShoreThread(Process): self._state["files"].append(f) def checkFileCompletion(self, fileobj: File) -> bool: - db = Database() + db = self._dbclass() complete = db.getCompletionByFileUUID(fileobj.uuid) del(db) diff --git a/classes/vessel.py b/classes/vessel.py index b3ddb73..5225d64 100644 --- a/classes/vessel.py +++ b/classes/vessel.py @@ -14,12 +14,14 @@ class Vessel: """Class describing a Vessel (= a replication destination) """ @classmethod - def fromConfig(cls, config: SectionProxy): + def fromConfig(cls, config: SectionProxy, dbclass: type = Database): """Create Vessel object from a Vessel section in the Config file Args: config (configparser.SectionProxy): Vessel section defining a Vessel + dbclass (type): Class to use for database connections. Defaults to + built-in Database using sqlite3. Raises: ValueError: Raised if section does not contain Address parameter @@ -60,7 +62,7 @@ class Vessel: if "Address" in config.keys(): return cls(config.name.split()[1], config["Address"], username, - password, passphrase, port, timeout, tempdir, ignoredirs) + password, passphrase, port, timeout, tempdir, ignoredirs, dbclass) else: raise ValueError("Definition for Vessel " + config.name.split()[1] + " does not contain Address!") @@ -69,7 +71,7 @@ class Vessel: password: Optional[str] = None, passphrase: Optional[str] = None, port: Optional[int] = None, timeout: Optional[int] = None, tempdir: Optional[Union[str, pathlib.Path]] = None, - ignoredirs: list[Optional[str]] = []) -> None: + ignoredirs: list[Optional[str]] = [], dbclass: type = Database) -> None: """Initialize new Vessel object Args: @@ -89,6 +91,7 @@ class Vessel: self._connection = None self._uploaded = self.getUploadedFromDB() # Files already uploaded self._ignoredirs = ignoredirs # Directories not replicated to this vessel + self._dbclass = dbclass @property def connection(self) -> Connection: @@ -116,7 +119,7 @@ class Vessel: Returns: list: List of UUIDs of Files that have been successfully uploaded """ - db = Database() + db = self._dbclass() return db.getCompletionForVessel(self) def currentUpload(self) -> Optional[tuple[str, str, str]]: @@ -129,9 +132,9 @@ class Vessel: checksum is the SHA256 hash of the file at the time of insertion into the database. None is returned if no such record is found. """ - self.assertTempDirectory() # After a reboot, the tempdir may be gone + self.assertTempDirectory() # After a reboot, the tempdir may be gone - db = Database() + db = self._dbclass() output = db.getFileByUUID(self.connection.getCurrentUploadUUID()) del db diff --git a/classes/vesselthread.py b/classes/vesselthread.py index 771ea8a..9478a1b 100644 --- a/classes/vesselthread.py +++ b/classes/vesselthread.py @@ -16,7 +16,7 @@ class VesselThread(Process): """Thread processing uploads to a single vessel """ - def __init__(self, vessel: Vessel, state: dict) -> None: + def __init__(self, vessel: Vessel, state: dict, dbclass: type = Database) -> None: """Initialize a new VesselThread Args: @@ -27,6 +27,7 @@ class VesselThread(Process): self.vessel = vessel self._state = state self._logger = Logger() + self._dbclass = dbclass def run(self) -> NoReturn: """Run thread and process uploads to the vessel @@ -95,7 +96,7 @@ class VesselThread(Process): f"Start processing file {fileobj.name} in directory {fileobj.directory.name} on vessel {self.vessel.name}") while True: - db = Database() + db = self._dbclass() if not db.getFileByUUID(fileobj.uuid): self._logger.debug( f"File {fileobj.name} in directory {fileobj.directory.name} does not exist anymore - deleting from {self.vessel.name}") @@ -111,7 +112,7 @@ class VesselThread(Process): f"File {fileobj.name} uploaded to vessel {self.vessel.name} completely - finalizing") remotefile.finalizeUpload() - db = Database() + db = self._dbclass() db.logCompletion(fileobj, self.vessel) del(db) @@ -143,7 +144,7 @@ class VesselThread(Process): self.vessel.compileComplete(remotefile) def checkFileCompletion(self, fileobj: File) -> None: - db = Database() + db = self._dbclass() complete = db.getCompletionByFileUUID(fileobj.uuid) del(db)