Coverage for agentos/marketplace/skills/database/database.py: 68%

1""" 

2database — Local SQLite database operations. 

3 

4Actions: query, tables, schema, create_table, insert 

5Works with any .db/.sqlite file. :memory: in same process shares one connection. 

6""" 

7 

8import sqlite3 

9from typing import Any 

10 

11# Persist :memory: connections within same process 

12_mem_conn = None 

13 

14 

15def _get_conn(db_path: str): 

16 global _mem_conn 

17 if db_path == ":memory:": 

18 if _mem_conn is None: 

19 _mem_conn = sqlite3.connect(":memory:") 

20 _mem_conn.row_factory = sqlite3.Row 

21 return _mem_conn 

22 conn = sqlite3.connect(db_path) 

23 conn.row_factory = sqlite3.Row 

24 return conn 

25 

26 

27def run(action: str = "tables", db_path: str = ":memory:", query: str = "", **kwargs: Any) -> str: 

28 try: 

29 conn = _get_conn(db_path) 

30 except Exception as e: 

31 return f"[database] Connection error: {e}" 

32 

33 try: 

34 if action == "tables": 

35 rows = conn.execute( 

36 "SELECT name FROM sqlite_master WHERE type='table' ORDER BY name" 

37 ).fetchall() 

38 tables = [r[0] for r in rows] 

39 return f"Tables ({len(tables)}): " + ", ".join(tables) if tables else "[database] No tables." 

40 

41 if action == "schema": 

42 if not query: 

43 return "[database] Table name required for schema." 

44 rows = conn.execute(f"PRAGMA table_info({query})").fetchall() 

45 cols = [f"{r['name']} {r['type']}" for r in rows] 

46 return f"Schema for {query}:\n" + "\n".join(f" {c}" for c in cols) 

47 

48 if action == "query": 

49 if not query: 

50 return "[database] SQL query required." 

51 cur = conn.execute(query) 

52 rows = cur.fetchall() 

53 if cur.description: 

54 headers = [d[0] for d in cur.description] 

55 result = "\t".join(headers) + "\n" 

56 result += "\n".join("\t".join(str(v) for v in row) for row in rows[:50]) 

57 return result 

58 return f"Query executed. Affected rows: {cur.rowcount}" 

59 

60 if action == "create_table": 

61 if not query: 

62 return "[database] CREATE TABLE statement required." 

63 conn.execute(query) 

64 conn.commit() 

65 return "[database] Table created." 

66 

67 if action == "insert": 

68 if not query: 

69 return "[database] INSERT statement required." 

70 conn.execute(query) 

71 conn.commit() 

72 return f"[database] Row inserted. Last ID: {conn.execute('SELECT last_insert_rowid()').fetchone()[0]}" 

73 

74 return f"[database] Unknown action: {action}" 

75 

76 except Exception as e: 

77 return f"[database] Error: {e}" 

78 finally: 

79 if db_path != ":memory:": 

80 conn.close() 

81 

82 

83__all__ = ["run"]