Update openspace/grounding/core/quality/store.py

openspace/grounding/core/quality/store.py

@@ -22,6 +22,98 @@ try:
 except ImportError:
     libsql = None
 
+class _LibsqlCursorProxy:
+    def __init__(self, cursor, conn_proxy):
+        self._cursor = cursor
+        self._conn_proxy = conn_proxy
+
+    def execute(self, *args, **kwargs):
+        self._cursor.execute(*args, **kwargs)
+        return self
+
+    def executescript(self, *args, **kwargs):
+        self._cursor.executescript(*args, **kwargs)
+        return self
+
+    def fetchone(self):
+        row = self._cursor.fetchone()
+        if row is not None and self._conn_proxy.row_factory:
+            return self._conn_proxy.row_factory(self, row)
+        return row
+
+    def fetchall(self):
+        rows = self._cursor.fetchall()
+        if self._conn_proxy.row_factory:
+            return [self._conn_proxy.row_factory(self, row) for row in rows]
+        return rows
+
+    @property
+    def description(self):
+        return getattr(self._cursor, "description", [])
+
+    @property
+    def rowcount(self):
+        return getattr(self._cursor, "rowcount", -1)
+
+    @property
+    def lastrowid(self):
+        return getattr(self._cursor, "lastrowid", None)
+
+class _LibsqlConnectionProxy:
+    def __init__(self, conn):
+        self._conn = conn
+        self.row_factory = None
+
+    def execute(self, *args, **kwargs):
+        cursor = self.cursor()
+        return cursor.execute(*args, **kwargs)
+
+    def executescript(self, *args, **kwargs):
+        cursor = self.cursor()
+        return cursor.executescript(*args, **kwargs)
+
+    def cursor(self):
+        return _LibsqlCursorProxy(self._conn.cursor(), self)
+
+    def commit(self):
+        if hasattr(self._conn, "commit"):
+            return self._conn.commit()
+
+    def rollback(self):
+        if hasattr(self._conn, "rollback"):
+            return self._conn.rollback()
+
+    def close(self):
+        if hasattr(self._conn, "close"):
+            return self._conn.close()
+
+class _RowProxy:
+    def __init__(self, row, description):
+        self._row = row
+        self._description = description
+        self._col_map = {col[0]: idx for idx, col in enumerate(description)}
+
+    def __getitem__(self, item):
+        if isinstance(item, int):
+            return self._row[item]
+        if item in self._col_map:
+            return self._row[self._col_map[item]]
+        raise KeyError(item)
+
+    def keys(self):
+        return self._col_map.keys()
+
+    def __iter__(self):
+        return iter(self._row)
+
+    def __len__(self):
+        return len(self._row)
+
+def _dict_factory(cursor, row):
+    if hasattr(cursor, "description") and cursor.description:
+        return _RowProxy(row, cursor.description)
+    return row
+
 from .types import ToolQualityRecord, ExecutionRecord, DescriptionQuality
 from openspace.utils.logging import Logger
 from openspace.config.constants import PROJECT_ROOT
@@ -90,8 +182,9 @@ class QualityStore:
         turso_token = os.environ.get("TURSO_AUTH_TOKEN")
 
         if turso_url and libsql is not None:
-            self._conn = libsql.connect(turso_url, auth_token=turso_token)
-            self._conn.row_factory = sqlite3.Row
+            raw_conn = libsql.connect(turso_url, auth_token=turso_token)
+            self._conn = _LibsqlConnectionProxy(raw_conn)
+            self._conn.row_factory = _dict_factory
             logger.debug(f"QualityStore ready (Turso) at {turso_url}")
         else:
             self._conn = sqlite3.connect(