Coverage for src/lumen/admin/loader.py: 88%

1""" 

2This module handles loading FHIR packages in TGZ format. 

3It extracts conformance resources and imports them into the resource repository. 

4""" 

5 

6from __future__ import annotations 

7 

8import io 

9import json 

10import tarfile 

11from typing import Any, TypedDict 

12from uuid import UUID 

13 

14import asyncpg 

15import asyncpg.pool 

16import structlog 

17 

18from lumen.db.repository import ResourceRepository 

19from lumen.search.registry import get_registry 

20 

21log = structlog.get_logger(__name__) 

22 

23# The FHIR conformance resources that we want to extract from packages 

24_CONFORMANCE_RESOURCES = { 

25 "StructureDefinition", 

26 "SearchParameter", 

27 "ValueSet", 

28 "CodeSystem", 

29 "ConceptMap", 

30 "OperationDefinition", 

31 "CapabilityStatement", 

32} 

33 

34_PACKAGE_TAG_SYSTEM = "urn:lumen:package" 

35 

36 

37class LoadPackageResult(TypedDict): 

38 name: str 

39 version: str 

40 imported: int 

41 skipped: int 

42 resource_types: dict[str, int] 

43 

44 

45def _extract_package_metadata(tar: tarfile.TarFile) -> tuple[str, str]: 

46 package_json_member = None 

47 for member in tar.getmembers(): 

48 if member.name.endswith("package.json") and "package/" in member.name: 

49 package_json_member = member 

50 break 

51 

52 if package_json_member: 

53 f = tar.extractfile(package_json_member) 

54 if f: 

55 try: 

56 pkg_data = json.loads(f.read().decode("utf-8")) 

57 return pkg_data.get("name", "unknown"), pkg_data.get("version", "unknown") 

58 except (json.JSONDecodeError, UnicodeDecodeError): 

59 pass 

60 return "unknown", "unknown" 

61 

62 

63async def _import_conformance_resource( 

64 repo: ResourceRepository, 

65 tenant_id: UUID, 

66 package_tag: str, 

67 res_dict: dict[str, Any], 

68 results: LoadPackageResult, 

69) -> None: 

70 res_type = res_dict.get("resourceType") 

71 res_id = res_dict.get("id") 

72 

73 if not res_type or not res_id: 

74 return 

75 

76 # We only care about importing structural / conformance resources 

77 if res_type not in _CONFORMANCE_RESOURCES: 

78 results["skipped"] += 1 

79 log.debug( 

80 "package.resource_skipped", 

81 package=package_tag, 

82 resource_type=res_type, 

83 resource_id=res_id, 

84 ) 

85 return 

86 

87 # Inject meta.tag to track which package this came from 

88 meta = res_dict.setdefault("meta", {}) 

89 tags = meta.setdefault("tag", []) 

90 

91 if not any( 

92 t.get("system") == _PACKAGE_TAG_SYSTEM and t.get("code") == package_tag for t in tags 

93 ): 

94 tags.append({"system": _PACKAGE_TAG_SYSTEM, "code": package_tag}) 

95 

96 # Upsert into database 

97 await repo.update(tenant_id, res_type, res_id, res_dict) 

98 

99 results["imported"] += 1 

100 results["resource_types"][res_type] = results["resource_types"].get(res_type, 0) + 1 

101 

102 

103async def load_package_tgz( 

104 conn: asyncpg.Connection[asyncpg.Record] | asyncpg.pool.PoolConnectionProxy[asyncpg.Record], 

105 tenant_id: UUID, 

106 file_bytes: bytes, 

107) -> LoadPackageResult: 

108 """Parse a .tgz FHIR package and store its conformance resources in the DB.""" 

109 registry = get_registry() 

110 repo = ResourceRepository(conn, registry) 

111 

112 results: LoadPackageResult = { 

113 "name": "unknown", 

114 "version": "unknown", 

115 "imported": 0, 

116 "skipped": 0, 

117 "resource_types": {}, 

118 } 

119 

120 with tarfile.open(fileobj=io.BytesIO(file_bytes), mode="r:gz") as tar: 

121 name, version = _extract_package_metadata(tar) 

122 results["name"] = name 

123 results["version"] = version 

124 

125 package_tag = f"{results['name']}@{results['version']}" 

126 log.info("package.upload_started", package=package_tag) 

127 

128 # 2. Iterate and import JSON resources 

129 for member in tar.getmembers(): 

130 if not member.name.endswith(".json"): 

131 continue 

132 # Skip package manifest or openapi definitions 

133 if "package.json" in member.name or "openapi" in member.name: 

134 continue 

135 

136 f = tar.extractfile(member) 

137 if not f: 

138 continue 

139 

140 try: 

141 res_dict = json.loads(f.read().decode("utf-8")) 

142 except (json.JSONDecodeError, UnicodeDecodeError): 

143 continue 

144 

145 await _import_conformance_resource(repo, tenant_id, package_tag, res_dict, results) 

146 

147 log.info( 

148 "package.upload_finished", 

149 package=package_tag, 

150 imported=results["imported"], 

151 skipped=results["skipped"], 

152 ) 

153 return results 

154 

155 

156async def get_installed_packages( 

157 conn: asyncpg.Connection[asyncpg.Record] | asyncpg.pool.PoolConnectionProxy[asyncpg.Record], 

158 tenant_id: UUID, 

159) -> list[str]: 

160 """Find all unique package tags in the tenant's resources.""" 

161 rows = await conn.fetch( 

162 """ 

163 SELECT DISTINCT code as package_tag 

164 FROM spidx_token 

165 WHERE tenant_id = $1 

166 AND param_name = '_tag' 

167 AND system = $2 

168 """, 

169 tenant_id, 

170 _PACKAGE_TAG_SYSTEM, 

171 ) 

172 return sorted([r["package_tag"] for r in rows if r["package_tag"]]) 

173 

174 

175async def delete_package( 

176 conn: asyncpg.Connection[asyncpg.Record] | asyncpg.pool.PoolConnectionProxy[asyncpg.Record], 

177 tenant_id: UUID, 

178 package_tag: str, 

179) -> int: 

180 """Soft-delete all resources associated with a package tag.""" 

181 registry = get_registry() 

182 repo = ResourceRepository(conn, registry) 

183 

184 rows = await conn.fetch( 

185 """ 

186 SELECT resource_type, resource_id 

187 FROM spidx_token 

188 WHERE tenant_id = $1 

189 AND param_name = '_tag' 

190 AND system = $2 

191 AND code = $3 

192 """, 

193 tenant_id, 

194 _PACKAGE_TAG_SYSTEM, 

195 package_tag, 

196 ) 

197 

198 deleted_count = 0 

199 for row in rows: 

200 success = await repo.delete(tenant_id, row["resource_type"], row["resource_id"]) 

201 if success: 

202 deleted_count += 1 

203 

204 log.info("package.deleted", package=package_tag, deleted_count=deleted_count) 

205 return deleted_count