1use std::collections::BTreeSet;
2use std::sync::Arc;
3
4use rpc::ExtensionProvides;
5
6use crate::test_both_dbs;
7use collab::db::Database;
8use collab::db::ExtensionVersionConstraints;
9use collab::db::{NewExtensionVersion, queries::extensions::convert_time_to_chrono};
10use rpc::ExtensionMetadata;
11test_both_dbs!(
12 test_extensions,
13 test_extensions_postgres,
14 test_extensions_sqlite
15);
16
17test_both_dbs!(
18 test_agent_servers_filter,
19 test_agent_servers_filter_postgres,
20 test_agent_servers_filter_sqlite
21);
22
23async fn test_agent_servers_filter(db: &Arc<Database>) {
24 // No extensions initially
25 let versions = db.get_known_extension_versions().await.unwrap();
26 assert!(versions.is_empty());
27
28 // Shared timestamp
29 let t0 = time::OffsetDateTime::from_unix_timestamp_nanos(0).unwrap();
30 let t0 = time::PrimitiveDateTime::new(t0.date(), t0.time());
31
32 // Insert two extensions, only one provides AgentServers
33 db.insert_extension_versions(
34 &[
35 (
36 "ext_agent_servers",
37 vec![NewExtensionVersion {
38 name: "Agent Servers Provider".into(),
39 version: semver::Version::parse("1.0.0").unwrap(),
40 description: "has agent servers".into(),
41 authors: vec!["author".into()],
42 repository: "org/agent-servers".into(),
43 schema_version: 1,
44 wasm_api_version: None,
45 provides: BTreeSet::from_iter([ExtensionProvides::AgentServers]),
46 published_at: t0,
47 }],
48 ),
49 (
50 "ext_plain",
51 vec![NewExtensionVersion {
52 name: "Plain Extension".into(),
53 version: semver::Version::parse("0.1.0").unwrap(),
54 description: "no agent servers".into(),
55 authors: vec!["author2".into()],
56 repository: "org/plain".into(),
57 schema_version: 1,
58 wasm_api_version: None,
59 provides: BTreeSet::default(),
60 published_at: t0,
61 }],
62 ),
63 ]
64 .into_iter()
65 .collect(),
66 )
67 .await
68 .unwrap();
69
70 // Filter by AgentServers provides
71 let provides_filter = BTreeSet::from_iter([ExtensionProvides::AgentServers]);
72
73 let filtered = db
74 .get_extensions(None, Some(&provides_filter), 1, 10)
75 .await
76 .unwrap();
77
78 // Expect only the extension that declared AgentServers
79 assert_eq!(filtered.len(), 1);
80 assert_eq!(filtered[0].id.as_ref(), "ext_agent_servers");
81}
82
83async fn test_extensions(db: &Arc<Database>) {
84 let versions = db.get_known_extension_versions().await.unwrap();
85 assert!(versions.is_empty());
86
87 let extensions = db.get_extensions(None, None, 1, 5).await.unwrap();
88 assert!(extensions.is_empty());
89
90 let t0 = time::OffsetDateTime::from_unix_timestamp_nanos(0).unwrap();
91 let t0 = time::PrimitiveDateTime::new(t0.date(), t0.time());
92
93 let t0_chrono = convert_time_to_chrono(t0);
94
95 db.insert_extension_versions(
96 &[
97 (
98 "ext1",
99 vec![
100 NewExtensionVersion {
101 name: "Extension 1".into(),
102 version: semver::Version::parse("0.0.1").unwrap(),
103 description: "an extension".into(),
104 authors: vec!["max".into()],
105 repository: "ext1/repo".into(),
106 schema_version: 1,
107 wasm_api_version: None,
108 provides: BTreeSet::default(),
109 published_at: t0,
110 },
111 NewExtensionVersion {
112 name: "Extension One".into(),
113 version: semver::Version::parse("0.0.2").unwrap(),
114 description: "a good extension".into(),
115 authors: vec!["max".into(), "marshall".into()],
116 repository: "ext1/repo".into(),
117 schema_version: 1,
118 wasm_api_version: None,
119 provides: BTreeSet::default(),
120 published_at: t0,
121 },
122 ],
123 ),
124 (
125 "ext2",
126 vec![NewExtensionVersion {
127 name: "Extension Two".into(),
128 version: semver::Version::parse("0.2.0").unwrap(),
129 description: "a great extension".into(),
130 authors: vec!["marshall".into()],
131 repository: "ext2/repo".into(),
132 schema_version: 0,
133 wasm_api_version: None,
134 provides: BTreeSet::default(),
135 published_at: t0,
136 }],
137 ),
138 ]
139 .into_iter()
140 .collect(),
141 )
142 .await
143 .unwrap();
144
145 let versions = db.get_known_extension_versions().await.unwrap();
146 assert_eq!(
147 versions,
148 [
149 ("ext1".into(), vec!["0.0.1".into(), "0.0.2".into()]),
150 ("ext2".into(), vec!["0.2.0".into()])
151 ]
152 .into_iter()
153 .collect()
154 );
155
156 // The latest version of each extension is returned.
157 let extensions = db.get_extensions(None, None, 1, 5).await.unwrap();
158 assert_eq!(
159 extensions,
160 &[
161 ExtensionMetadata {
162 id: "ext1".into(),
163 manifest: rpc::ExtensionApiManifest {
164 name: "Extension One".into(),
165 version: "0.0.2".into(),
166 authors: vec!["max".into(), "marshall".into()],
167 description: Some("a good extension".into()),
168 repository: "ext1/repo".into(),
169 schema_version: Some(1),
170 wasm_api_version: None,
171 provides: BTreeSet::default(),
172 },
173 published_at: t0_chrono,
174 download_count: 0,
175 },
176 ExtensionMetadata {
177 id: "ext2".into(),
178 manifest: rpc::ExtensionApiManifest {
179 name: "Extension Two".into(),
180 version: "0.2.0".into(),
181 authors: vec!["marshall".into()],
182 description: Some("a great extension".into()),
183 repository: "ext2/repo".into(),
184 schema_version: Some(0),
185 wasm_api_version: None,
186 provides: BTreeSet::default(),
187 },
188 published_at: t0_chrono,
189 download_count: 0
190 },
191 ]
192 );
193
194 // Extensions with too new of a schema version are excluded.
195 let extensions = db.get_extensions(None, None, 0, 5).await.unwrap();
196 assert_eq!(
197 extensions,
198 &[ExtensionMetadata {
199 id: "ext2".into(),
200 manifest: rpc::ExtensionApiManifest {
201 name: "Extension Two".into(),
202 version: "0.2.0".into(),
203 authors: vec!["marshall".into()],
204 description: Some("a great extension".into()),
205 repository: "ext2/repo".into(),
206 schema_version: Some(0),
207 wasm_api_version: None,
208 provides: BTreeSet::default(),
209 },
210 published_at: t0_chrono,
211 download_count: 0
212 },]
213 );
214
215 // Record extensions being downloaded.
216 for _ in 0..7 {
217 assert!(db.record_extension_download("ext2", "0.0.2").await.unwrap());
218 }
219
220 for _ in 0..3 {
221 assert!(db.record_extension_download("ext1", "0.0.1").await.unwrap());
222 }
223
224 for _ in 0..2 {
225 assert!(db.record_extension_download("ext1", "0.0.2").await.unwrap());
226 }
227
228 // Record download returns false if the extension does not exist.
229 assert!(
230 !db.record_extension_download("no-such-extension", "0.0.2")
231 .await
232 .unwrap()
233 );
234
235 // Extensions are returned in descending order of total downloads.
236 let extensions = db.get_extensions(None, None, 1, 5).await.unwrap();
237 assert_eq!(
238 extensions,
239 &[
240 ExtensionMetadata {
241 id: "ext2".into(),
242 manifest: rpc::ExtensionApiManifest {
243 name: "Extension Two".into(),
244 version: "0.2.0".into(),
245 authors: vec!["marshall".into()],
246 description: Some("a great extension".into()),
247 repository: "ext2/repo".into(),
248 schema_version: Some(0),
249 wasm_api_version: None,
250 provides: BTreeSet::default(),
251 },
252 published_at: t0_chrono,
253 download_count: 7
254 },
255 ExtensionMetadata {
256 id: "ext1".into(),
257 manifest: rpc::ExtensionApiManifest {
258 name: "Extension One".into(),
259 version: "0.0.2".into(),
260 authors: vec!["max".into(), "marshall".into()],
261 description: Some("a good extension".into()),
262 repository: "ext1/repo".into(),
263 schema_version: Some(1),
264 wasm_api_version: None,
265 provides: BTreeSet::default(),
266 },
267 published_at: t0_chrono,
268 download_count: 5,
269 },
270 ]
271 );
272
273 // Add more extensions, including a new version of `ext1`, and backfilling
274 // an older version of `ext2`.
275 db.insert_extension_versions(
276 &[
277 (
278 "ext1",
279 vec![NewExtensionVersion {
280 name: "Extension One".into(),
281 version: semver::Version::parse("0.0.3").unwrap(),
282 description: "a real good extension".into(),
283 authors: vec!["max".into(), "marshall".into()],
284 repository: "ext1/repo".into(),
285 schema_version: 1,
286 wasm_api_version: None,
287 provides: BTreeSet::default(),
288 published_at: t0,
289 }],
290 ),
291 (
292 "ext2",
293 vec![NewExtensionVersion {
294 name: "Extension Two".into(),
295 version: semver::Version::parse("0.1.0").unwrap(),
296 description: "an old extension".into(),
297 authors: vec!["marshall".into()],
298 repository: "ext2/repo".into(),
299 schema_version: 0,
300 wasm_api_version: None,
301 provides: BTreeSet::default(),
302 published_at: t0,
303 }],
304 ),
305 ]
306 .into_iter()
307 .collect(),
308 )
309 .await
310 .unwrap();
311
312 let versions = db.get_known_extension_versions().await.unwrap();
313 assert_eq!(
314 versions,
315 [
316 (
317 "ext1".into(),
318 vec!["0.0.1".into(), "0.0.2".into(), "0.0.3".into()]
319 ),
320 ("ext2".into(), vec!["0.1.0".into(), "0.2.0".into()])
321 ]
322 .into_iter()
323 .collect()
324 );
325
326 let extensions = db.get_extensions(None, None, 1, 5).await.unwrap();
327 assert_eq!(
328 extensions,
329 &[
330 ExtensionMetadata {
331 id: "ext2".into(),
332 manifest: rpc::ExtensionApiManifest {
333 name: "Extension Two".into(),
334 version: "0.2.0".into(),
335 authors: vec!["marshall".into()],
336 description: Some("a great extension".into()),
337 repository: "ext2/repo".into(),
338 schema_version: Some(0),
339 wasm_api_version: None,
340 provides: BTreeSet::default(),
341 },
342 published_at: t0_chrono,
343 download_count: 7
344 },
345 ExtensionMetadata {
346 id: "ext1".into(),
347 manifest: rpc::ExtensionApiManifest {
348 name: "Extension One".into(),
349 version: "0.0.3".into(),
350 authors: vec!["max".into(), "marshall".into()],
351 description: Some("a real good extension".into()),
352 repository: "ext1/repo".into(),
353 schema_version: Some(1),
354 wasm_api_version: None,
355 provides: BTreeSet::default(),
356 },
357 published_at: t0_chrono,
358 download_count: 5,
359 },
360 ]
361 );
362}
363
364test_both_dbs!(
365 test_extensions_by_id,
366 test_extensions_by_id_postgres,
367 test_extensions_by_id_sqlite
368);
369
370async fn test_extensions_by_id(db: &Arc<Database>) {
371 let versions = db.get_known_extension_versions().await.unwrap();
372 assert!(versions.is_empty());
373
374 let extensions = db.get_extensions(None, None, 1, 5).await.unwrap();
375 assert!(extensions.is_empty());
376
377 let t0 = time::OffsetDateTime::from_unix_timestamp_nanos(0).unwrap();
378 let t0 = time::PrimitiveDateTime::new(t0.date(), t0.time());
379
380 let t0_chrono = convert_time_to_chrono(t0);
381
382 db.insert_extension_versions(
383 &[
384 (
385 "ext1",
386 vec![
387 NewExtensionVersion {
388 name: "Extension 1".into(),
389 version: semver::Version::parse("0.0.1").unwrap(),
390 description: "an extension".into(),
391 authors: vec!["max".into()],
392 repository: "ext1/repo".into(),
393 schema_version: 1,
394 wasm_api_version: Some("0.0.4".into()),
395 provides: BTreeSet::from_iter([
396 ExtensionProvides::Grammars,
397 ExtensionProvides::Languages,
398 ]),
399 published_at: t0,
400 },
401 NewExtensionVersion {
402 name: "Extension 1".into(),
403 version: semver::Version::parse("0.0.2").unwrap(),
404 description: "a good extension".into(),
405 authors: vec!["max".into()],
406 repository: "ext1/repo".into(),
407 schema_version: 1,
408 wasm_api_version: Some("0.0.4".into()),
409 provides: BTreeSet::from_iter([
410 ExtensionProvides::Grammars,
411 ExtensionProvides::Languages,
412 ExtensionProvides::LanguageServers,
413 ]),
414 published_at: t0,
415 },
416 NewExtensionVersion {
417 name: "Extension 1".into(),
418 version: semver::Version::parse("0.0.3").unwrap(),
419 description: "a real good extension".into(),
420 authors: vec!["max".into(), "marshall".into()],
421 repository: "ext1/repo".into(),
422 schema_version: 1,
423 wasm_api_version: Some("0.0.5".into()),
424 provides: BTreeSet::from_iter([
425 ExtensionProvides::Grammars,
426 ExtensionProvides::Languages,
427 ExtensionProvides::LanguageServers,
428 ]),
429 published_at: t0,
430 },
431 ],
432 ),
433 (
434 "ext2",
435 vec![NewExtensionVersion {
436 name: "Extension 2".into(),
437 version: semver::Version::parse("0.2.0").unwrap(),
438 description: "a great extension".into(),
439 authors: vec!["marshall".into()],
440 repository: "ext2/repo".into(),
441 schema_version: 0,
442 wasm_api_version: None,
443 provides: BTreeSet::default(),
444 published_at: t0,
445 }],
446 ),
447 ]
448 .into_iter()
449 .collect(),
450 )
451 .await
452 .unwrap();
453
454 let extensions = db
455 .get_extensions_by_ids(
456 &["ext1"],
457 Some(&ExtensionVersionConstraints {
458 schema_versions: 1..=1,
459 wasm_api_versions: "0.0.1".parse().unwrap()..="0.0.4".parse().unwrap(),
460 }),
461 )
462 .await
463 .unwrap();
464
465 assert_eq!(
466 extensions,
467 &[ExtensionMetadata {
468 id: "ext1".into(),
469 manifest: rpc::ExtensionApiManifest {
470 name: "Extension 1".into(),
471 version: "0.0.2".into(),
472 authors: vec!["max".into()],
473 description: Some("a good extension".into()),
474 repository: "ext1/repo".into(),
475 schema_version: Some(1),
476 wasm_api_version: Some("0.0.4".into()),
477 provides: BTreeSet::from_iter([
478 ExtensionProvides::Grammars,
479 ExtensionProvides::Languages,
480 ExtensionProvides::LanguageServers,
481 ]),
482 },
483 published_at: t0_chrono,
484 download_count: 0,
485 }]
486 );
487}