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 ( 

40 f"Tables ({len(tables)}): " + ", ".join(tables) 

41 if tables 

42 else "[database] No tables." 

43 ) 

44 

45 if action == "schema": 

46 if not query: 

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

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

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

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

51 

52 if action == "query": 

53 if not query: 

54 return "[database] SQL query required." 

55 cur = conn.execute(query) 

56 rows = cur.fetchall() 

57 if cur.description: 

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

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

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

61 return result 

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

63 

64 if action == "create_table": 

65 if not query: 

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

67 conn.execute(query) 

68 conn.commit() 

69 return "[database] Table created." 

70 

71 if action == "insert": 

72 if not query: 

73 return "[database] INSERT statement required." 

74 conn.execute(query) 

75 conn.commit() 

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

77 

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

79 

80 except Exception as e: 

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

82 finally: 

83 if db_path != ":memory:": 

84 conn.close() 

85 

86 

87__all__ = ["run"]