Coverage for core / src / sensorkit / cli / kv.py: 0%

1# SPDX-License-Identifier: Apache-2.0 

2import glob 

3import io 

4import json 

5import sys 

6from typing import Any, TextIO 

7 

8import aiofile 

9import asyncclick as click 

10from pydantic import BaseModel, ValidationError 

11from rich import print_json 

12 

13from sensorkit.cli.utils import console, entity_option, with_kit 

14 

15 

16@click.group("kv") 

17async def kv_group(): 

18 """Manage Key-Value storage operations for SensorKit.""" 

19 

20@kv_group.command("ls") 

21@entity_option() 

22@with_kit 

23async def list_kv(kit, entity: str | None): 

24 """ 

25 List all key-value pairs, optionally filtered by entity. 

26 

27 Args: 

28 entity: Optional entity name to filter the key-value context. 

29 """ 

30 from sensorkit.backend.base import Entity, KeyValueContext 

31 

32 context: KeyValueContext 

33 if entity: 

34 context = kit.backend.key_value(entity=Entity.at(entity)) 

35 else: 

36 context = kit.backend.key_value() 

37 

38 for item in await context.get_all(deep=True): 

39 console.print("-" * 80) 

40 console.print(f"[bold white]{item.key}[/bold white]\n") 

41 data = item.value.decode() 

42 try: 

43 json.loads(data) 

44 print_json(data) 

45 except json.JSONDecodeError: 

46 console.print(data) 

47 

48 console.print("-" * 80) 

49 

50@kv_group.command("get") 

51@entity_option(required=True) 

52@click.argument("key") 

53@with_kit 

54async def get_kv(kit, entity: str, key: str): 

55 """ 

56 Get the value of a single key from the key-value store. 

57 

58 Args: 

59 entity: The entity to read from. 

60 key: The key/property name to retrieve. 

61 """ 

62 from sensorkit.backend.base import Entity 

63 

64 context = kit.backend.key_value(entity=Entity.at(entity)) 

65 item = await context.get(key) 

66 

67 console.print("-" * 80) 

68 console.print(f'[bold white]{item.key}[/bold white]\n') 

69 data = item.value.decode() 

70 

71 try: 

72 json.loads(data) 

73 print_json(data) 

74 except json.JSONDecodeError: 

75 console.print(data) 

76 

77 console.print("-" * 80) 

78 

79@kv_group.command("put") 

80@entity_option(required=True) 

81@click.argument("key") 

82@click.argument("value") 

83@with_kit 

84async def put_kv(kit, entity: str, key: str, value: str): 

85 """ 

86 Put a new key-value pair into the store. 

87 

88 Args: 

89 entity: The entity to write to. 

90 key: The key/property name to update. 

91 value: The value to store as a string. 

92 """ 

93 from sensorkit.backend.base import Entity 

94 

95 context = kit.backend.key_value(entity=Entity.at(entity)) 

96 await context.update(key, value.encode()) 

97 console.print(f'[bold white]SUCCESS: kv put for {entity=} {key=} {value=}[/bold white]') 

98 

99 

100@kv_group.command("delete") 

101@entity_option(required=True) 

102@click.argument("key", required=False) 

103@with_kit 

104async def delete_kv(kit, entity: str, key: str | None): 

105 """ 

106 Delete one or more key-value entries from the SensorKit KV store. 

107 

108 This command deletes a single key or all keys under a given entity in the key-value store. 

109 If a key is specified, only that key is deleted. 

110 If no key is provided, all keys under the specified entity are deleted. 

111 

112 Args: 

113 entity (str): The entity namespace to operate on (required). 

114 key (str | None): The specific key to delete. If omitted, all keys under the entity will be deleted. 

115 """ 

116 from sensorkit.backend.base import Entity 

117 

118 context = kit.backend.key_value(entity=Entity.at(entity)) 

119 

120 if not key: 

121 entries = await context.get_all(deep=True) 

122 for e in entries: 

123 await context.delete(e.key.prop, revision=e.revision) 

124 console.print(f"[bold white]SUCCESS: kv delete on {e.key}[/bold white]") 

125 else: 

126 entry = await context.get(key) 

127 await context.delete(key, revision=entry.revision) 

128 console.print(f"[bold white]SUCCESS: kv delete on {key}[/bold white]") 

129 

130class KVRecord(BaseModel): 

131 """ 

132 A single key-value record with an associated entity. 

133 

134 Attributes: 

135 entity (str): The name of the entity. 

136 key (str): The property/key name. 

137 value (Any): The value to associate with the key. 

138 """ 

139 entity: str 

140 key: str 

141 value: Any 

142 

143def load_kv_records(source: TextIO) -> list[KVRecord]: 

144 """ 

145 Load and validate a list of KVRecord entries from a YAML stream. 

146 

147 Accepts multiple YAML documents using `yaml.safe_load_all`. 

148 

149 Args: 

150 source: A readable stream (file or stdin) containing YAML records. 

151 

152 Returns: 

153 A list of validated KVRecord objects. 

154 """ 

155 import yaml 

156 

157 configurations: list[KVRecord] = [] 

158 

159 try: 

160 for doc in yaml.safe_load_all(source): 

161 if not doc: 

162 continue 

163 try: 

164 record = KVRecord.model_validate(doc) 

165 configurations.append(record) 

166 except ValidationError: 

167 console.print(f'[bold red]ERROR: {doc=} is not a KVRecord.[/bold red]') 

168 except yaml.YAMLError: 

169 console.print('[bold red]ERROR: could not load yaml.[/bold red]') 

170 

171 return configurations 

172 

173def expand_files(files: tuple[str, ...]) -> list[str]: 

174 """ 

175 Expand glob patterns into a list of matching file paths. 

176 

177 Args: 

178 files: Tuple of filenames or glob patterns. 

179 

180 Returns: 

181 List of matching file paths. 

182 """ 

183 if not files: 

184 return [] 

185 

186 expanded = [] 

187 for pattern in files: 

188 expanded.extend(glob.glob(pattern)) 

189 return expanded 

190 

191@kv_group.command("load") 

192@entity_option() 

193@click.option('-n', "--no-clobber", is_flag=True) 

194@click.argument("files", nargs=-1, required=False) 

195@with_kit 

196async def load_kv_command(kit, entity: str | None, no_clobber: bool, files: tuple[str, ...]): 

197 """ 

198 Load multiple key-value records from YAML files or stdin. 

199 

200 Supports reading from: 

201 - One or more file paths or glob patterns 

202 - Standard input if no files are provided 

203 

204 Each YAML document should conform to the KVRecord schema. 

205 

206 Args: 

207 entity: If set, filters updates to only records matching this entity. 

208 files: One or more YAML file paths or globs. 

209 """ 

210 from sensorkit.backend.base import Entity, KVError 

211 

212 configurations = await _read_configs(files) 

213 

214 for item in configurations: 

215 if entity and entity != item.entity: 

216 continue 

217 

218 context = kit.backend.key_value(entity=Entity.at(item.entity)) 

219 

220 if no_clobber: 

221 try: 

222 prev_entry = await context.get(item.key) 

223 prev_value = prev_entry.value.decode() 

224 try: 

225 prev_value = json.loads(prev_value) 

226 except json.JSONDecodeError: 

227 pass 

228 

229 if prev_value == item.value: 

230 console.print(f'[bold yellow]WARNING: kv put skipped for {item.entity}.{item.key}[/bold yellow]') 

231 continue 

232 except KVError: # only errors with kvget 

233 pass 

234 

235 await context.update(item.key, json.dumps(item.value).encode()) 

236 console.print(f'[bold white]SUCCESS: kv put for {item.entity=} {item.key=} {item.value=}[/bold white]') 

237 

238 

239async def _read_configs(files: tuple[str, ...]) -> list[KVRecord]: 

240 if len(files) == 0: 

241 return load_kv_records(sys.stdin) 

242 

243 configurations = [] 

244 

245 for path in expand_files(files): 

246 try: 

247 async with aiofile.async_open(path, "r") as f: 

248 configurations.extend( 

249 load_kv_records(io.StringIO(await f.read())) 

250 ) 

251 except FileNotFoundError: 

252 console.print(f"[bold red]ERROR: file not found {path}[/bold red]") 

253 

254 return configurations