Coverage for oc_meta / run / migration / extract_subset.py: 100%

1#!/usr/bin/env python 

2 

3# SPDX-FileCopyrightText: 2026 Arcangelo Massari <arcangelo.massari@unibo.it> 

4# 

5# SPDX-License-Identifier: ISC 

6 

7# -*- coding: utf-8 -*- 

8 

9import argparse 

10import gzip 

11import sys 

12from urllib.parse import urlparse 

13 

14import rdflib 

15from rdflib.term import Node 

16from rich_argparse import RichHelpFormatter 

17from oc_meta.lib.sparql import execute_sparql 

18 

19CHUNK_SIZE = 20 

20 

21 

22def get_subjects_of_class(endpoint: str, class_uri: str, limit: int) -> list[str]: 

23 query = f""" 

24 SELECT ?s 

25 WHERE {{ 

26 ?s a <{class_uri}> . 

27 }} 

28 LIMIT {limit} 

29 """ 

30 results = execute_sparql(endpoint, query) 

31 return [result["s"]["value"] for result in results["results"]["bindings"]] 

32 

33 

34def load_entities_from_file(entities_file: str) -> list[str]: 

35 with open(entities_file, "r") as f: 

36 return [line.strip() for line in f if line.strip()] 

37 

38 

39def parse_object( 

40 result: dict[str, dict[str, str]], 

41) -> rdflib.URIRef | rdflib.BNode | rdflib.Literal: 

42 o_value = result["o"]["value"] 

43 o_type = result["o"]["type"] 

44 if o_type == "uri": 

45 return rdflib.URIRef(o_value) 

46 if o_type == "bnode": 

47 return rdflib.BNode(o_value) 

48 if "datatype" in result["o"]: 

49 return rdflib.Literal(o_value, datatype=result["o"]["datatype"]) 

50 if "xml:lang" in result["o"]: 

51 return rdflib.Literal(o_value, lang=result["o"]["xml:lang"]) 

52 return rdflib.Literal(o_value) 

53 

54 

55def get_triples_for_entities( 

56 endpoint: str, 

57 entity_uris: list[str], 

58 use_graphs: bool, 

59) -> list[tuple[rdflib.URIRef, rdflib.URIRef, Node, rdflib.URIRef | None]]: 

60 quads: list[tuple[rdflib.URIRef, rdflib.URIRef, Node, rdflib.URIRef | None]] = [] 

61 

62 for i in range(0, len(entity_uris), CHUNK_SIZE): 

63 chunk = entity_uris[i : i + CHUNK_SIZE] 

64 values = " ".join(f"<{uri}>" for uri in chunk) 

65 

66 if use_graphs: 

67 query = f""" 

68 SELECT ?s ?p ?o ?g 

69 WHERE {{ 

70 GRAPH ?g {{ 

71 VALUES ?s {{ {values} }} 

72 ?s ?p ?o . 

73 }} 

74 }} 

75 """ 

76 else: 

77 query = f""" 

78 SELECT ?s ?p ?o 

79 WHERE {{ 

80 VALUES ?s {{ {values} }} 

81 ?s ?p ?o . 

82 }} 

83 """ 

84 

85 results = execute_sparql(endpoint, query) 

86 for result in results["results"]["bindings"]: 

87 s_term = rdflib.URIRef(result["s"]["value"]) 

88 p_term = rdflib.URIRef(result["p"]["value"]) 

89 o_term = parse_object(result) 

90 g_term = rdflib.URIRef(result["g"]["value"]) if "g" in result else None 

91 quads.append((s_term, p_term, o_term, g_term)) 

92 

93 return quads 

94 

95 

96def extract_subset( 

97 endpoint: str, 

98 limit: int, 

99 output_file: str, 

100 compress: bool, 

101 max_retries: int = 5, 

102 class_uri: str | None = None, 

103 entities_file: str | None = None, 

104 use_graphs: bool = True, 

105) -> tuple[int, str]: 

106 if entities_file: 

107 subjects = load_entities_from_file(entities_file) 

108 else: 

109 assert class_uri is not None 

110 subjects = get_subjects_of_class(endpoint, class_uri, limit) 

111 

112 processed_entities: set[str] = set() 

113 pending_entities = set(subjects) 

114 

115 dataset: rdflib.Dataset | None = None 

116 graph: rdflib.Graph | None = None 

117 if use_graphs: 

118 dataset = rdflib.Dataset() 

119 else: 

120 graph = rdflib.Graph() 

121 

122 while pending_entities: 

123 batch = sorted(pending_entities - processed_entities) 

124 if not batch: 

125 break # pragma: no cover 

126 

127 processed_entities.update(batch) 

128 pending_entities.clear() 

129 

130 quads = get_triples_for_entities(endpoint, batch, use_graphs) 

131 

132 for s_term, p_term, o_term, g_term in quads: 

133 if dataset is not None: 

134 named_graph = dataset.graph(g_term) 

135 named_graph.add((s_term, p_term, o_term)) 

136 elif graph is not None: 

137 graph.add((s_term, p_term, o_term)) 

138 

139 if isinstance(o_term, rdflib.URIRef): 

140 o_str = str(o_term) 

141 if o_str not in processed_entities: 

142 pending_entities.add(o_str) 

143 

144 store = dataset if dataset is not None else graph 

145 assert store is not None 

146 output_format = "nquads" if use_graphs else "nt" 

147 if compress: 

148 if not output_file.endswith(".gz"): 

149 output_file = output_file + ".gz" 

150 with gzip.open(output_file, "wb") as f: 

151 store.serialize(destination=f, format=output_format) # type: ignore[arg-type] 

152 else: 

153 store.serialize(destination=output_file, format=output_format) 

154 

155 return len(processed_entities), output_file 

156 

157 

158def main(): # pragma: no cover 

159 parser = argparse.ArgumentParser( 

160 description="Extract a subset of data from a SPARQL endpoint", 

161 formatter_class=RichHelpFormatter, 

162 ) 

163 parser.add_argument( 

164 "--endpoint", 

165 default="http://localhost:8890/sparql", 

166 help="SPARQL endpoint URL (default: http://localhost:8890/sparql)", 

167 ) 

168 

169 discovery = parser.add_mutually_exclusive_group() 

170 discovery.add_argument( 

171 "--class", 

172 dest="class_uri", 

173 help="Class URI to extract instances of (default: fabio:Expression)", 

174 ) 

175 discovery.add_argument( 

176 "--entities-file", 

177 dest="entities_file", 

178 help="File with entity URIs to extract (one per line)", 

179 ) 

180 

181 parser.add_argument( 

182 "--limit", 

183 type=int, 

184 default=1000, 

185 help="Maximum number of initial entities to process (default: 1000)", 

186 ) 

187 parser.add_argument( 

188 "--output", default="output.nq", help="Output file name (default: output.nq)" 

189 ) 

190 parser.add_argument( 

191 "--compress", action="store_true", help="Compress output file using gzip" 

192 ) 

193 parser.add_argument( 

194 "--retries", 

195 type=int, 

196 default=5, 

197 help="Maximum number of retries for failed queries (default: 5)", 

198 ) 

199 parser.add_argument( 

200 "--no-graphs", 

201 action="store_true", 

202 help="Disable named graph queries and output N-Triples instead of N-Quads", 

203 ) 

204 

205 args = parser.parse_args() 

206 

207 if not args.class_uri and not args.entities_file: 

208 args.class_uri = "http://purl.org/spar/fabio/Expression" 

209 

210 try: 

211 parsed_url = urlparse(args.endpoint) 

212 if not all([parsed_url.scheme, parsed_url.netloc]): 

213 raise ValueError("Invalid endpoint URL") 

214 except Exception: 

215 print(f"Error: Invalid endpoint URL: {args.endpoint}") 

216 return 1 

217 

218 try: 

219 entity_count, final_output_file = extract_subset( 

220 args.endpoint, 

221 args.limit, 

222 args.output, 

223 args.compress, 

224 args.retries, 

225 class_uri=args.class_uri, 

226 entities_file=args.entities_file, 

227 use_graphs=not args.no_graphs, 

228 ) 

229 

230 print(f"Extraction complete. Processed {entity_count} entities.") 

231 print(f"Output saved to {final_output_file}") 

232 

233 return 0 

234 except Exception as e: 

235 print(f"Error: {e}") 

236 return 1 

237 

238 

239if __name__ == "__main__": # pragma: no cover 

240 sys.exit(main())