@@ -161,6 +161,19 @@ def main():
r = client.delete_user(args.username)
print_user(r)
+ def handle_get_db_usage(args, client):
+ usage = client.get_db_usage()
+ print(usage)
+ tables = sorted(usage.keys())
+ print("{name:20}| {rows:20}".format(name="Table name", rows="Rows"))
+ print(("-" * 20) + "+" + ("-" * 20))
+ for t in tables:
+ print("{name:20}| {rows:<20}".format(name=t, rows=usage[t]["rows"]))
+ print()
+
+ total_rows = sum(t["rows"] for t in usage.values())
+ print(f"Total rows: {total_rows}")
+
parser = argparse.ArgumentParser(description='Hash Equivalence Client')
parser.add_argument('--address', default=DEFAULT_ADDRESS, help='Server address (default "%(default)s")')
parser.add_argument('--log', default='WARNING', help='Set logging level')
@@ -223,6 +236,9 @@ def main():
delete_user_parser.add_argument("--username", "-u", help="Username", required=True)
delete_user_parser.set_defaults(func=handle_delete_user)
+ db_usage_parser = subparsers.add_parser('get-db-usage', help="Database Usage")
+ db_usage_parser.set_defaults(func=handle_get_db_usage)
+
args = parser.parse_args()
logger = logging.getLogger('hashserv')
@@ -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):
@@ -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,
@@ -566,6 +567,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"])
@@ -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
@@ -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
@@ -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):
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> --- bin/bitbake-hashclient | 16 ++++++++++++++++ lib/hashserv/client.py | 5 +++++ lib/hashserv/server.py | 5 +++++ lib/hashserv/sqlalchemy.py | 14 ++++++++++++++ lib/hashserv/sqlite.py | 37 +++++++++++++++++++++++++++++++++++++ lib/hashserv/tests.py | 9 +++++++++ 6 files changed, 86 insertions(+)