diff options
author | Joshua Watt <JPEWhacker@gmail.com> | 2023-11-03 08:26:33 -0600 |
---|---|---|
committer | Richard Purdie <richard.purdie@linuxfoundation.org> | 2023-11-09 17:21:15 +0000 |
commit | c9c1224447e147e0de92953bc85cea75670b898c (patch) | |
tree | dde05ce0a1b1c17bcad5830f39ed52090f53b435 /lib/hashserv | |
parent | 71e2f5b52b686f34df364ae1f2fc058f45cd5e18 (diff) | |
download | bitbake-c9c1224447e147e0de92953bc85cea75670b898c.tar.gz |
hashserv: Add db-usage API
Adds an API to query the server for the usage of the database (e.g. how
many rows are present in each table)
Signed-off-by: Joshua Watt <JPEWhacker@gmail.com>
Signed-off-by: Richard Purdie <richard.purdie@linuxfoundation.org>
Diffstat (limited to 'lib/hashserv')
-rw-r--r-- | lib/hashserv/client.py | 5 | ||||
-rw-r--r-- | lib/hashserv/server.py | 5 | ||||
-rw-r--r-- | lib/hashserv/sqlalchemy.py | 14 | ||||
-rw-r--r-- | lib/hashserv/sqlite.py | 37 | ||||
-rw-r--r-- | lib/hashserv/tests.py | 9 |
5 files changed, 70 insertions, 0 deletions
diff --git a/lib/hashserv/client.py b/lib/hashserv/client.py index 4457f8e50..5e0a462b1 100644 --- a/lib/hashserv/client.py +++ b/lib/hashserv/client.py @@ -186,6 +186,10 @@ class AsyncClient(bb.asyncrpc.AsyncClient): self.saved_become_user = username return result + async def get_db_usage(self): + await self._set_mode(self.MODE_NORMAL) + return (await self.invoke({"get-db-usage": {}}))["usage"] + class Client(bb.asyncrpc.Client): def __init__(self, username=None, password=None): @@ -214,6 +218,7 @@ class Client(bb.asyncrpc.Client): "new_user", "delete_user", "become_user", + "get_db_usage", ) def _get_async_client(self): diff --git a/lib/hashserv/server.py b/lib/hashserv/server.py index ca419a1ab..c5b9797e4 100644 --- a/lib/hashserv/server.py +++ b/lib/hashserv/server.py @@ -249,6 +249,7 @@ class ServerClient(bb.asyncrpc.AsyncServerConnection): "get-outhash": self.handle_get_outhash, "get-stream": self.handle_get_stream, "get-stats": self.handle_get_stats, + "get-db-usage": self.handle_get_db_usage, # Not always read-only, but internally checks if the server is # read-only "report": self.handle_report, @@ -567,6 +568,10 @@ class ServerClient(bb.asyncrpc.AsyncServerConnection): oldest = datetime.now() - timedelta(seconds=-max_age) return {"count": await self.db.clean_unused(oldest)} + @permissions(DB_ADMIN_PERM) + async def handle_get_db_usage(self, request): + return {"usage": await self.db.get_usage()} + # The authentication API is always allowed async def handle_auth(self, request): username = str(request["username"]) diff --git a/lib/hashserv/sqlalchemy.py b/lib/hashserv/sqlalchemy.py index bfd8a8446..818b51951 100644 --- a/lib/hashserv/sqlalchemy.py +++ b/lib/hashserv/sqlalchemy.py @@ -27,6 +27,7 @@ from sqlalchemy import ( and_, delete, update, + func, ) import sqlalchemy.engine from sqlalchemy.orm import declarative_base @@ -401,3 +402,16 @@ class Database(object): async with self.db.begin(): result = await self.db.execute(statement) return result.rowcount != 0 + + async def get_usage(self): + usage = {} + async with self.db.begin() as session: + for name, table in Base.metadata.tables.items(): + statement = select(func.count()).select_from(table) + self.logger.debug("%s", statement) + result = await self.db.execute(statement) + usage[name] = { + "rows": result.scalar(), + } + + return usage diff --git a/lib/hashserv/sqlite.py b/lib/hashserv/sqlite.py index 414ee8ffb..dfdccbbaa 100644 --- a/lib/hashserv/sqlite.py +++ b/lib/hashserv/sqlite.py @@ -120,6 +120,18 @@ class Database(object): self.db = sqlite3.connect(self.dbname) self.db.row_factory = sqlite3.Row + with closing(self.db.cursor()) as cursor: + cursor.execute("SELECT sqlite_version()") + + version = [] + for v in cursor.fetchone()[0].split("."): + try: + version.append(int(v)) + except ValueError: + version.append(v) + + self.sqlite_version = tuple(version) + async def __aenter__(self): return self @@ -362,3 +374,28 @@ class Database(object): ) self.db.commit() return cursor.rowcount != 0 + + async def get_usage(self): + usage = {} + with closing(self.db.cursor()) as cursor: + if self.sqlite_version >= (3, 33): + table_name = "sqlite_schema" + else: + table_name = "sqlite_master" + + cursor.execute( + f""" + SELECT name FROM {table_name} WHERE type = 'table' AND name NOT LIKE 'sqlite_%' + """ + ) + for row in cursor.fetchall(): + cursor.execute( + """ + SELECT COUNT() FROM %s + """ + % row["name"], + ) + usage[row["name"]] = { + "rows": cursor.fetchone()[0], + } + return usage diff --git a/lib/hashserv/tests.py b/lib/hashserv/tests.py index 311b7b777..9d5bec245 100644 --- a/lib/hashserv/tests.py +++ b/lib/hashserv/tests.py @@ -767,6 +767,15 @@ class HashEquivalenceCommonTests(object): with self.auth_perms("@user-admin") as client: become = client.become_user(client.username) + def test_get_db_usage(self): + usage = self.client.get_db_usage() + + self.assertTrue(isinstance(usage, dict)) + for name in usage.keys(): + self.assertTrue(isinstance(usage[name], dict)) + self.assertIn("rows", usage[name]) + self.assertTrue(isinstance(usage[name]["rows"], int)) + class TestHashEquivalenceUnixServer(HashEquivalenceTestSetup, HashEquivalenceCommonTests, unittest.TestCase): def get_server_addr(self, server_idx): |