Coverage for src / server_list / spec / db.py: 79%

1#!/usr/bin/env python3 

2""" 

3Shared database utilities for server-list. 

4 

5Provides common path definitions and database connection utilities 

6using my_lib.sqlite_util for Kubernetes-optimized SQLite operations. 

7""" 

8 

9from __future__ import annotations 

10 

11from contextlib import contextmanager 

12from pathlib import Path 

13from typing import TYPE_CHECKING 

14 

15import my_lib.sqlite_util 

16 

17if TYPE_CHECKING: 

18 from server_list.config import Config 

19 

20# Base directory paths (defaults, can be overridden by config) 

21BASE_DIR = Path(__file__).resolve().parent.parent.parent.parent 

22SCHEMA_DIR = BASE_DIR / "schema" 

23 

24# Data directory (default, can be overridden by init_from_config) 

25DATA_DIR = BASE_DIR / "data" 

26 

27# Database paths (default, updated by init_from_config) 

28SERVER_DATA_DB = DATA_DIR / "server_data.db" 

29CACHE_DB = DATA_DIR / "cache.db" 

30CPU_SPEC_DB = DATA_DIR / "cpu_spec.db" 

31 

32# Image directory (default, can be overridden by init_from_config) 

33IMAGE_DIR = BASE_DIR / "img" 

34 

35# Schema paths 

36SQLITE_SCHEMA_PATH = SCHEMA_DIR / "sqlite.schema" 

37SECRET_SCHEMA_PATH = SCHEMA_DIR / "secret.schema" 

38CONFIG_SCHEMA_PATH = SCHEMA_DIR / "config.schema" 

39 

40# Config file paths 

41CONFIG_PATH = BASE_DIR / "config.yaml" 

42SECRET_PATH = BASE_DIR / "secret.yaml" 

43 

44 

45def init_from_config(config: Config) -> None: 

46 """Initialize paths from Config object. 

47 

48 Updates module-level path variables based on config values. 

49 

50 Args: 

51 config: Config object with data and webapp settings. 

52 """ 

53 global DATA_DIR, SERVER_DATA_DB, CACHE_DB, CPU_SPEC_DB, IMAGE_DIR 

54 

55 # Update data directory from config 

56 DATA_DIR = config.data.get_cache_dir(BASE_DIR) 

57 SERVER_DATA_DB = DATA_DIR / "server_data.db" 

58 CACHE_DB = DATA_DIR / "cache.db" 

59 CPU_SPEC_DB = DATA_DIR / "cpu_spec.db" 

60 

61 # Update image directory from config 

62 IMAGE_DIR = config.webapp.get_image_dir(BASE_DIR) 

63 

64 

65def ensure_data_dir() -> None: 

66 """Ensure the data directory exists.""" 

67 DATA_DIR.mkdir(parents=True, exist_ok=True) 

68 

69 

70@contextmanager 

71def get_connection(db_path: Path, timeout: float = 10.0): 

72 """ 

73 Get a database connection with proper cleanup. 

74 

75 Uses my_lib.sqlite_util for Kubernetes-optimized SQLite settings. 

76 

77 Args: 

78 db_path: Path to the database file 

79 timeout: Connection timeout in seconds 

80 

81 Yields: 

82 sqlite3.Connection: Database connection 

83 """ 

84 ensure_data_dir() 

85 with my_lib.sqlite_util.connect(db_path, timeout=timeout) as conn: 

86 yield conn 

87 

88 

89def init_schema(db_path: Path, schema_sql: str) -> None: 

90 """ 

91 Initialize database with given schema SQL. 

92 

93 Args: 

94 db_path: Path to the database file 

95 schema_sql: SQL script to execute for schema creation 

96 """ 

97 with get_connection(db_path) as conn: 

98 my_lib.sqlite_util.exec_schema(conn, schema_sql) 

99 conn.commit() 

100 

101 

102def init_schema_from_file(db_path: Path, schema_path: Path) -> None: 

103 """ 

104 Initialize database with schema from a file. 

105 

106 Args: 

107 db_path: Path to the database file 

108 schema_path: Path to the schema SQL file 

109 """ 

110 with get_connection(db_path) as conn: 

111 my_lib.sqlite_util.exec_schema_from_file(conn, schema_path) 

112 conn.commit()