From: Michael Tremer <michael.tremer@ipfire.org>
Date: Sat, 20 May 2023 09:38:03 +0000 (+0000)
Subject: base: Add new functions to create objetcs easier
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=3817afa6f41b5dfed3651d43ba75cd2f10d50caa;p=pbs.git

base: Add new functions to create objetcs easier

Signed-off-by: Michael Tremer <michael.tremer@ipfire.org>
---

diff --git a/src/buildservice/__init__.py b/src/buildservice/__init__.py
index 4ac4770b..e6045f39 100644
--- a/src/buildservice/__init__.py
+++ b/src/buildservice/__init__.py
@@ -114,7 +114,7 @@ class Backend(object):
 
 		log.debug("Connecting to database %s @ %s" % (name, hostname))
 
-		return database.Connection(hostname, name, user=user, password=password)
+		return database.Connection(self, hostname, name, user=user, password=password)
 
 	def _create_tmp_path(self):
 		"""
diff --git a/src/buildservice/base.py b/src/buildservice/base.py
index da7a04de..f2a920e8 100644
--- a/src/buildservice/base.py
+++ b/src/buildservice/base.py
@@ -48,12 +48,16 @@ class DataObject(Object):
 	def __hash__(self):
 		return hash(self.id)
 
-	def init(self, id, data=None):
+	def init(self, id, data=None, **kwargs):
 		self.id = id
 
 		if data:
 			self.data = data
 
+		# Set any extra arguments (to populate the cache)
+		for arg in kwargs:
+			setattr(self, arg, kwargs[arg])
+
 	@lazy_property
 	def data(self):
 		assert self.table, "Table name is not set"
diff --git a/src/buildservice/database.py b/src/buildservice/database.py
index 27db3806..1ba98a54 100644
--- a/src/buildservice/database.py
+++ b/src/buildservice/database.py
@@ -33,7 +33,8 @@ class Connection(object):
 		We explicitly set the timezone to UTC and the character encoding to
 		UTF-8 on all connections to avoid time zone and encoding errors.
 	"""
-	def __init__(self, host, database, user=None, password=None):
+	def __init__(self, backend, host, database, user=None, password=None):
+		self.backend = backend
 		self.host = host
 		self.database = database
 
@@ -194,6 +195,25 @@ class Connection(object):
 	def transaction(self):
 		return Transaction(self)
 
+	def fetch_one(self, cls, query, *args, **kwargs):
+		"""
+			Takes a class and a query and will return one object of that class
+		"""
+		# Execute the query
+		res = self.get(query, *args)
+
+		# Return an object (if possible)
+		if res:
+			return cls(self.backend, res.id, res, **kwargs)
+
+	def fetch_many(self, cls, query, *args, **kwargs):
+		# Execute the query
+		res = self.query(query, *args)
+
+		# Return a generator with objects
+		for row in res:
+			yield cls(self.backend, row.id, row, **kwargs)
+
 
 class Row(dict):
 	"""A dict that allows for object-like property access syntax."""