Skip to main content

pelicode_db/
lib.rs

1//! SQLite storage of function features (specification ยง5).
2
3#![warn(missing_docs)]
4
5use std::collections::{BTreeMap, HashMap, HashSet};
6use std::path::Path;
7use std::str::FromStr;
8
9use pelicode_core::{Config, FileFeatures, Level, NORMALIZER_VERSION};
10use rusqlite::{Connection, OptionalExtension, Transaction, params};
11use serde::Serialize;
12
13/// Errors of the database.
14#[derive(Debug, thiserror::Error)]
15pub enum Error {
16    /// An error of SQLite.
17    #[error(transparent)]
18    Sqlite(#[from] rusqlite::Error),
19    /// The database was built with other parameters or another table layout than the current ones.
20    #[error("database was built with {key} = {stored}, but the current setting is {current}")]
21    Incompatible {
22        /// Name of the parameter.
23        key: String,
24        /// Value recorded in the database.
25        stored: String,
26        /// Value of the current setting.
27        current: String,
28    },
29    /// A parameter recorded in the database cannot be read.
30    #[error("invalid value in database meta: {0}")]
31    InvalidMeta(String),
32    /// The database is a snapshot, which is for checking only.
33    #[error("database is a snapshot for checking only; index into the database it was made from")]
34    Snapshot,
35}
36
37mod clean;
38mod query;
39mod snapshot;
40mod stats;
41
42pub use clean::CleanSummary;
43pub use query::{FunctionInfo, Occurrence, StoredFingerprint};
44pub use snapshot::{SnapshotOptions, SnapshotSummary};
45pub use stats::{CommonHash, HashStatsInfo};
46
47/// Result of the database operations.
48pub type Result<T> = std::result::Result<T, Error>;
49
50/// Version of the table layout; databases with another version are rejected.
51pub const SCHEMA_VERSION: u32 = 5;
52
53/// Oldest version of the table layout that is upgraded on open; older ones
54/// are rejected.
55pub const OLDEST_SCHEMA_VERSION: u32 = 4;
56
57const SCHEMA: &str = r#"
58CREATE TABLE IF NOT EXISTS meta (
59    key   TEXT PRIMARY KEY,
60    value TEXT NOT NULL
61);
62CREATE TABLE IF NOT EXISTS ecosystem (
63    id   INTEGER PRIMARY KEY,
64    name TEXT NOT NULL UNIQUE
65);
66CREATE TABLE IF NOT EXISTS package (
67    id             INTEGER PRIMARY KEY,
68    ecosystem_id   INTEGER NOT NULL REFERENCES ecosystem(id),
69    name           TEXT NOT NULL,
70    repository_url TEXT,
71    UNIQUE (ecosystem_id, name)
72);
73CREATE TABLE IF NOT EXISTS package_version (
74    id             INTEGER PRIMARY KEY,
75    package_id     INTEGER NOT NULL REFERENCES package(id),
76    version        TEXT NOT NULL,
77    published_at   TEXT,
78    license_spdx   TEXT,
79    archive_url    TEXT,
80    archive_sha256 TEXT,
81    crawled_at     TEXT,
82    status         TEXT NOT NULL,
83    error          TEXT,
84    UNIQUE (package_id, version)
85);
86CREATE TABLE IF NOT EXISTS source_file (
87    id             INTEGER PRIMARY KEY,
88    content_sha256 TEXT NOT NULL UNIQUE,
89    language       TEXT NOT NULL,
90    function_count INTEGER NOT NULL
91);
92CREATE TABLE IF NOT EXISTS version_file (
93    package_version_id INTEGER NOT NULL REFERENCES package_version(id),
94    source_file_id     INTEGER NOT NULL REFERENCES source_file(id),
95    path               TEXT NOT NULL,
96    PRIMARY KEY (package_version_id, path)
97);
98CREATE INDEX IF NOT EXISTS version_file_source ON version_file(source_file_id);
99CREATE TABLE IF NOT EXISTS function (
100    id                 INTEGER PRIMARY KEY,
101    source_file_id     INTEGER NOT NULL REFERENCES source_file(id),
102    parent_function_id INTEGER REFERENCES function(id),
103    name               TEXT NOT NULL,
104    kind               TEXT NOT NULL,
105    start_line         INTEGER NOT NULL,
106    end_line           INTEGER NOT NULL,
107    start_column       INTEGER NOT NULL,
108    end_column         INTEGER NOT NULL,
109    start_byte         INTEGER NOT NULL,
110    end_byte           INTEGER NOT NULL,
111    token_count        INTEGER NOT NULL,
112    node_count         INTEGER NOT NULL,
113    depth              INTEGER NOT NULL
114);
115CREATE INDEX IF NOT EXISTS function_source ON function(source_file_id);
116CREATE TABLE IF NOT EXISTS exact_hash (
117    function_id INTEGER NOT NULL REFERENCES function(id),
118    level       INTEGER NOT NULL,
119    hash        BLOB NOT NULL,
120    PRIMARY KEY (function_id, level)
121);
122CREATE INDEX IF NOT EXISTS exact_hash_lookup ON exact_hash(level, hash);
123CREATE TABLE IF NOT EXISTS fingerprint (
124    function_id INTEGER NOT NULL REFERENCES function(id),
125    level       INTEGER NOT NULL,
126    hash        INTEGER NOT NULL,
127    position    INTEGER NOT NULL,
128    start_line  INTEGER NOT NULL,
129    end_line    INTEGER NOT NULL
130);
131CREATE INDEX IF NOT EXISTS fingerprint_lookup ON fingerprint(level, hash);
132CREATE INDEX IF NOT EXISTS fingerprint_function ON fingerprint(function_id, level);
133CREATE TABLE IF NOT EXISTS minhash (
134    function_id INTEGER NOT NULL REFERENCES function(id),
135    level       INTEGER NOT NULL,
136    signature   BLOB NOT NULL,
137    PRIMARY KEY (function_id, level)
138);
139CREATE TABLE IF NOT EXISTS lsh_bucket (
140    level       INTEGER NOT NULL,
141    band        INTEGER NOT NULL,
142    bucket_hash INTEGER NOT NULL,
143    function_id INTEGER NOT NULL REFERENCES function(id)
144);
145CREATE INDEX IF NOT EXISTS lsh_bucket_lookup ON lsh_bucket(level, band, bucket_hash);
146-- Document frequencies, rebuilt by `refresh_hash_stats`. Only hashes found
147-- in two or more packages are stored; a missing hash counts as one package.
148CREATE TABLE IF NOT EXISTS exact_stats (
149    level          INTEGER NOT NULL,
150    hash           BLOB NOT NULL,
151    package_count  INTEGER NOT NULL,
152    function_count INTEGER NOT NULL,
153    PRIMARY KEY (level, hash)
154);
155CREATE TABLE IF NOT EXISTS fingerprint_stats (
156    level          INTEGER NOT NULL,
157    hash           INTEGER NOT NULL,
158    package_count  INTEGER NOT NULL,
159    function_count INTEGER NOT NULL,
160    PRIMARY KEY (level, hash)
161);
162-- Hashes whose package counts may have changed since the last statistics:
163-- those of the files linked to or unlinked from a version since then.
164CREATE TABLE IF NOT EXISTS stats_pending_exact (
165    level INTEGER NOT NULL,
166    hash  BLOB NOT NULL,
167    PRIMARY KEY (level, hash)
168) WITHOUT ROWID;
169CREATE TABLE IF NOT EXISTS stats_pending_fingerprint (
170    level INTEGER NOT NULL,
171    hash  INTEGER NOT NULL,
172    PRIMARY KEY (level, hash)
173) WITHOUT ROWID;
174"#;
175
176/// Where indexed source files come from.
177#[derive(Debug, Clone, Default)]
178pub struct PackageInfo {
179    /// Ecosystem of the package: `pypi`, `maven`, `go`, `rubygems`, `crates` or `local`.
180    pub ecosystem: String,
181    /// Name of the package within its ecosystem.
182    pub name: String,
183    /// Version of the package.
184    pub version: String,
185    /// License, as an SPDX expression when it could be identified.
186    pub license_spdx: Option<String>,
187    /// URL of the source repository.
188    pub repository_url: Option<String>,
189    /// Publication date and time (RFC 3339).
190    pub published_at: Option<String>,
191    /// URL the source archive was downloaded from.
192    pub archive_url: Option<String>,
193    /// SHA-256 of the downloaded archive (hex).
194    pub archive_sha256: Option<String>,
195}
196
197/// Result of [`Database::index`].
198#[derive(Debug, Clone, Default, Serialize)]
199pub struct IndexSummary {
200    /// Files given to [`Database::index`].
201    pub files: usize,
202    /// Files whose content was not in the database yet.
203    pub new_files: usize,
204    /// Functions stored from new files.
205    pub functions: usize,
206    /// Of `functions`, those with fingerprints.
207    pub indexed_functions: usize,
208}
209
210/// Number of rows of the main tables.
211#[derive(Debug, Clone, Default, Serialize)]
212pub struct Stats {
213    /// Ecosystems.
214    pub ecosystems: usize,
215    /// Packages.
216    pub packages: usize,
217    /// Package versions, failed ones included.
218    pub versions: usize,
219    /// Distinct source files (by content).
220    pub source_files: usize,
221    /// Stored functions.
222    pub functions: usize,
223    /// Stored functions with fingerprints.
224    pub indexed_functions: usize,
225}
226
227/// One occurrence of a stored function, for export.
228#[derive(Debug, Clone, Serialize)]
229pub struct FunctionRow {
230    /// Identifier of the function in the database.
231    pub id: i64,
232    /// Ecosystem of the package.
233    pub ecosystem: String,
234    /// Name of the package.
235    pub package: String,
236    /// Version of the package.
237    pub version: String,
238    /// License of the package version.
239    pub license: Option<String>,
240    /// Path of the file in the package.
241    pub path: String,
242    /// Language of the file.
243    pub language: String,
244    /// Name of the function (for display).
245    pub name: String,
246    /// Tree-sitter node kind of the function.
247    pub kind: String,
248    /// First line (1-based).
249    pub start_line: usize,
250    /// Last line (1-based, inclusive).
251    pub end_line: usize,
252    /// 1-based byte column of the first character.
253    pub start_column: usize,
254    /// 1-based byte column of the last character.
255    pub end_column: usize,
256    /// Start of the byte range in the file.
257    pub start_byte: usize,
258    /// End of the byte range in the file (exclusive).
259    pub end_byte: usize,
260    /// Identifier of the enclosing function, if any.
261    pub parent_id: Option<i64>,
262    /// Number of tokens.
263    pub token_count: usize,
264    /// Number of named syntax tree nodes.
265    pub node_count: usize,
266    /// Depth of the syntax tree.
267    pub depth: usize,
268    /// Whether the function has fingerprints.
269    pub indexed: bool,
270    /// Exact hashes per normalization level (hex).
271    pub exact_hashes: BTreeMap<Level, String>,
272}
273
274/// A pelicode database: an SQLite file holding the features of functions and
275/// where they were found, but no source code.
276///
277/// # Examples
278///
279/// ```
280/// use pelicode_core::{Config, Language, extract_bytes};
281/// use pelicode_db::{Database, PackageInfo};
282///
283/// let dir = tempfile::tempdir().unwrap();
284/// let config = Config::default();
285/// let mut db = Database::open(dir.path().join("pelicode.db"), &config).unwrap();
286///
287/// let source = b"package demo\nfunc Add(a, b int) int { return a + b }\n";
288/// let file = extract_bytes("add.go".as_ref(), Language::Go, source, &config).unwrap();
289/// let package = PackageInfo {
290///     ecosystem: "local".into(),
291///     name: "demo".into(),
292///     version: "1.0".into(),
293///     ..PackageInfo::default()
294/// };
295/// db.index(&package, &[file]).unwrap();
296/// assert_eq!(db.stats().unwrap().functions, 1);
297/// ```
298pub struct Database {
299    conn: Connection,
300}
301
302impl Database {
303    /// Opens (or creates) a database for features computed with `config`.
304    /// Fails when the database was built with incompatible parameters.
305    ///
306    /// # Errors
307    ///
308    /// [`Error::Incompatible`] when the database was built with other feature
309    /// parameters or another table layout, and [`Error::Sqlite`] when the file cannot
310    /// be opened.
311    ///
312    /// # Examples
313    ///
314    /// ```
315    /// use pelicode_core::Config;
316    /// use pelicode_db::{Database, Error};
317    ///
318    /// let dir = tempfile::tempdir().unwrap();
319    /// let path = dir.path().join("pelicode.db");
320    /// Database::open(&path, &Config::default()).unwrap();
321    ///
322    /// let other = Config { kgram: 5, ..Config::default() };
323    /// assert!(matches!(Database::open(&path, &other), Err(Error::Incompatible { .. })));
324    /// ```
325    pub fn open(path: impl AsRef<Path>, config: &Config) -> Result<Database> {
326        let db = Database::open_unchecked(path)?;
327        db.check_meta(config)?;
328        Ok(db)
329    }
330
331    /// Opens a database without checking its parameters (for reading).
332    ///
333    /// The table layout is still checked; use [`Database::config_from_meta`] to featurize
334    /// code like the stored functions.
335    ///
336    /// # Errors
337    ///
338    /// [`Error::Incompatible`] when the table layout differs, and [`Error::Sqlite`]
339    /// when the file cannot be opened.
340    ///
341    /// # Examples
342    ///
343    /// ```
344    /// use pelicode_core::Config;
345    /// use pelicode_db::Database;
346    ///
347    /// let dir = tempfile::tempdir().unwrap();
348    /// let path = dir.path().join("pelicode.db");
349    /// Database::open(&path, &Config { kgram: 5, ..Config::default() }).unwrap();
350    ///
351    /// let db = Database::open_unchecked(&path).unwrap();
352    /// assert_eq!(db.meta().unwrap()["kgram"], "5");
353    /// ```
354    pub fn open_unchecked(path: impl AsRef<Path>) -> Result<Database> {
355        let conn = Connection::open(path)?;
356        conn.pragma_update(None, "journal_mode", "WAL")?;
357        conn.pragma_update(None, "foreign_keys", "ON")?;
358        conn.execute_batch(SCHEMA)?;
359        let db = Database { conn };
360        db.check_schema()?;
361        Ok(db)
362    }
363
364    fn check_schema(&self) -> Result<()> {
365        let stored = self.meta()?;
366        match stored.get("schema_version") {
367            Some(v) if *v == SCHEMA_VERSION.to_string() => Ok(()),
368            // Version 5 only adds the tables of pending hashes, created above.
369            Some(v) if v == "4" => {
370                // The statistics can be updated incrementally from now on if
371                // nothing was indexed since they were made.
372                let current = stored.get("index_generation").map_or("0", String::as_str);
373                if stored.get("stats_generation").map(String::as_str) == Some(current) {
374                    stats::set_pending_complete(&self.conn, true)?;
375                }
376                self.conn.execute(
377                    "UPDATE meta SET value = ?1 WHERE key = 'schema_version'",
378                    [SCHEMA_VERSION.to_string()],
379                )?;
380                Ok(())
381            }
382            None if stored.is_empty() => {
383                self.conn.execute(
384                    "INSERT INTO meta (key, value) VALUES ('schema_version', ?1)",
385                    [SCHEMA_VERSION.to_string()],
386                )?;
387                Ok(())
388            }
389            other => Err(Error::Incompatible {
390                key: "schema_version".into(),
391                stored: other.cloned().unwrap_or_else(|| "(none)".into()),
392                current: SCHEMA_VERSION.to_string(),
393            }),
394        }
395    }
396
397    /// The parameters recorded in the database.
398    ///
399    /// The keys are `schema_version`, `normalizer_version`, `ident_mode`, `kgram`,
400    /// `window`, `minhash_size` and `bands`, and the bookkeeping of the hash
401    /// statistics.
402    ///
403    /// # Errors
404    ///
405    /// [`Error::Sqlite`] when the table cannot be read.
406    ///
407    /// # Examples
408    ///
409    /// ```
410    /// # use pelicode_core::Config;
411    /// # use pelicode_db::Database;
412    /// # let dir = tempfile::tempdir().unwrap();
413    /// let db = Database::open(dir.path().join("pelicode.db"), &Config::default()).unwrap();
414    /// let meta = db.meta().unwrap();
415    /// assert_eq!(meta["kgram"], "10");
416    /// assert_eq!(meta["window"], "4");
417    /// ```
418    pub fn meta(&self) -> Result<BTreeMap<String, String>> {
419        let mut stmt = self.conn.prepare("SELECT key, value FROM meta")?;
420        let rows = stmt.query_map([], |r| Ok((r.get(0)?, r.get(1)?)))?;
421        Ok(rows.collect::<rusqlite::Result<_>>()?)
422    }
423
424    /// `base` with the parameters recorded in the database, so that queries are
425    /// featurized exactly like the stored functions. Fingerprint levels become
426    /// the levels that actually have fingerprints in the database.
427    ///
428    /// # Errors
429    ///
430    /// [`Error::Incompatible`] when the database was built with other normalization
431    /// rules, and [`Error::InvalidMeta`] when a recorded parameter cannot be parsed.
432    ///
433    /// # Examples
434    ///
435    /// ```
436    /// use pelicode_core::Config;
437    /// use pelicode_db::Database;
438    ///
439    /// let dir = tempfile::tempdir().unwrap();
440    /// let path = dir.path().join("pelicode.db");
441    /// Database::open(&path, &Config { kgram: 7, ..Config::default() }).unwrap();
442    ///
443    /// let db = Database::open_unchecked(&path).unwrap();
444    /// let config = db.config_from_meta(&Config::default()).unwrap();
445    /// assert_eq!(config.kgram, 7);
446    /// ```
447    pub fn config_from_meta(&self, base: &Config) -> Result<Config> {
448        let stored = self.meta()?;
449        fn parse<T: FromStr>(
450            stored: &BTreeMap<String, String>,
451            key: &str,
452            value: &mut T,
453        ) -> Result<()> {
454            if let Some(s) = stored.get(key) {
455                *value = s
456                    .parse()
457                    .map_err(|_| Error::InvalidMeta(format!("{key} = {s}")))?;
458            }
459            Ok(())
460        }
461        let mut version = NORMALIZER_VERSION;
462        parse(&stored, "normalizer_version", &mut version)?;
463        if version != NORMALIZER_VERSION {
464            return Err(Error::Incompatible {
465                key: "normalizer_version".into(),
466                stored: version.to_string(),
467                current: NORMALIZER_VERSION.to_string(),
468            });
469        }
470        let mut config = base.clone();
471        parse(&stored, "ident_mode", &mut config.ident_mode)?;
472        parse(&stored, "kgram", &mut config.kgram)?;
473        parse(&stored, "window", &mut config.window)?;
474        parse(&stored, "minhash_size", &mut config.minhash_size)?;
475        parse(&stored, "bands", &mut config.bands)?;
476
477        let levels: Vec<u8> = match stored.get(snapshot::LEVELS_KEY) {
478            Some(levels) => levels
479                .split(',')
480                .filter(|l| !l.is_empty())
481                .map(|l| l.parse())
482                .collect::<std::result::Result<_, _>>()
483                .map_err(|_| Error::InvalidMeta(format!("{} = {levels}", snapshot::LEVELS_KEY)))?,
484            None => {
485                let mut stmt = self
486                    .conn
487                    .prepare("SELECT DISTINCT level FROM minhash ORDER BY level")?;
488                stmt.query_map([], |r| r.get(0))?
489                    .collect::<rusqlite::Result<_>>()?
490            }
491        };
492        if !levels.is_empty() {
493            config.fingerprint_levels = levels.into_iter().filter_map(Level::from_u8).collect();
494        }
495        Ok(config)
496    }
497
498    fn check_meta(&self, config: &Config) -> Result<()> {
499        let current = [
500            ("normalizer_version", NORMALIZER_VERSION.to_string()),
501            ("ident_mode", config.ident_mode.to_string()),
502            ("kgram", config.kgram.to_string()),
503            ("window", config.window.to_string()),
504            ("minhash_size", config.minhash_size.to_string()),
505            ("bands", config.bands.to_string()),
506        ];
507        let stored = self.meta()?;
508        for (key, value) in &current {
509            match stored.get(*key) {
510                Some(s) if s != value => {
511                    return Err(Error::Incompatible {
512                        key: key.to_string(),
513                        stored: s.clone(),
514                        current: value.clone(),
515                    });
516                }
517                Some(_) => {}
518                None => {
519                    self.conn.execute(
520                        "INSERT INTO meta (key, value) VALUES (?1, ?2)",
521                        params![key, value],
522                    )?;
523                }
524            }
525        }
526        Ok(())
527    }
528
529    /// Stores `files` as the content of package version `pkg`, replacing what
530    /// was previously recorded for that version. Files already in the
531    /// database (same content hash) are linked without storing them again.
532    ///
533    /// # Errors
534    ///
535    /// [`Error::Snapshot`] when this database is a snapshot, and [`Error::Sqlite`]
536    /// when the files cannot be stored; nothing is stored then.
537    ///
538    /// # Examples
539    ///
540    /// ```
541    /// use pelicode_core::{Config, Language, extract_bytes};
542    /// use pelicode_db::{Database, PackageInfo};
543    ///
544    /// let dir = tempfile::tempdir().unwrap();
545    /// let config = Config::default();
546    /// let mut db = Database::open(dir.path().join("pelicode.db"), &config).unwrap();
547    /// let file = extract_bytes("a.rb".as_ref(), Language::Ruby, b"def one\n  1\nend\n", &config).unwrap();
548    /// let package = |version: &str| PackageInfo {
549    ///     ecosystem: "rubygems".into(),
550    ///     name: "one".into(),
551    ///     version: version.into(),
552    ///     ..PackageInfo::default()
553    /// };
554    ///
555    /// let first = db.index(&package("1.0"), &[file.clone()]).unwrap();
556    /// assert_eq!((first.files, first.new_files, first.functions), (1, 1, 1));
557    ///
558    /// // The same file in another version is linked, not stored again.
559    /// let second = db.index(&package("1.1"), &[file]).unwrap();
560    /// assert_eq!(second.new_files, 0);
561    /// ```
562    pub fn index(&mut self, pkg: &PackageInfo, files: &[FileFeatures]) -> Result<IndexSummary> {
563        self.refuse_snapshot()?;
564        let tx = self.conn.transaction()?;
565        let version_id = upsert_version(&tx, pkg, "indexed", None)?;
566        // The hashes of the files unlinked here and linked below may change
567        // their package counts; remember them for an incremental refresh.
568        let track = stats::pending_complete(&tx)?;
569        if track {
570            stats::mark_pending(&tx, version_id)?;
571        }
572        tx.execute(
573            "DELETE FROM version_file WHERE package_version_id = ?1",
574            [version_id],
575        )?;
576
577        let mut summary = IndexSummary {
578            files: files.len(),
579            ..IndexSummary::default()
580        };
581        for file in files {
582            let existing: Option<i64> = tx
583                .query_row(
584                    "SELECT id FROM source_file WHERE content_sha256 = ?1",
585                    [&file.content_sha256],
586                    |r| r.get(0),
587                )
588                .optional()?;
589            let file_id = match existing {
590                Some(id) => id,
591                None => {
592                    summary.new_files += 1;
593                    insert_file(&tx, file, &mut summary)?
594                }
595            };
596            tx.execute(
597                "INSERT OR REPLACE INTO version_file (package_version_id, source_file_id, path)
598                 VALUES (?1, ?2, ?3)",
599                params![version_id, file_id, file.path.to_string_lossy()],
600            )?;
601        }
602        if track {
603            stats::mark_pending(&tx, version_id)?;
604        }
605        // Lets `hash_stats_info` tell whether the statistics are out of date.
606        tx.execute(
607            "INSERT INTO meta (key, value) VALUES ('index_generation', '1')
608             ON CONFLICT (key) DO UPDATE SET value = CAST(value AS INTEGER) + 1",
609            [],
610        )?;
611        tx.commit()?;
612        Ok(summary)
613    }
614
615    /// Status of a package version (`indexed` or `failed`), if recorded.
616    ///
617    /// # Errors
618    ///
619    /// [`Error::Sqlite`] when the database cannot be read.
620    ///
621    /// # Examples
622    ///
623    /// ```
624    /// # use pelicode_core::{Config, Language, extract_bytes};
625    /// # use pelicode_db::{Database, PackageInfo};
626    /// # let dir = tempfile::tempdir().unwrap();
627    /// # let config = Config { min_tokens: 0, ..Config::default() };
628    /// # let mut db = Database::open(dir.path().join("pelicode.db"), &config).unwrap();
629    /// # let source = b"def total(xs):\n    s = 0\n    for x in xs:\n        s += x * 2\n    return s\n";
630    /// # let file = extract_bytes("src/calc.py".as_ref(), Language::Python, source, &config).unwrap();
631    /// # let package = PackageInfo {
632    /// #     ecosystem: "pypi".into(), name: "calc".into(), version: "1.0".into(),
633    /// #     license_spdx: Some("MIT".into()), ..PackageInfo::default()
634    /// # };
635    /// # db.index(&package, &[file]).unwrap();
636    /// assert_eq!(db.version_status("pypi", "calc", "1.0").unwrap().as_deref(), Some("indexed"));
637    /// assert_eq!(db.version_status("pypi", "calc", "2.0").unwrap(), None);
638    /// ```
639    pub fn version_status(
640        &self,
641        ecosystem: &str,
642        name: &str,
643        version: &str,
644    ) -> Result<Option<String>> {
645        Ok(self
646            .conn
647            .query_row(
648                "SELECT pv.status FROM package_version pv
649                 JOIN package p ON p.id = pv.package_id
650                 JOIN ecosystem e ON e.id = p.ecosystem_id
651                 WHERE e.name = ?1 AND p.name = ?2 AND pv.version = ?3",
652                params![ecosystem, name, version],
653                |r| r.get(0),
654            )
655            .optional()?)
656    }
657
658    /// Records that fetching or indexing `pkg` failed, so that it can be retried later.
659    /// Files already linked to the version are kept.
660    ///
661    /// # Errors
662    ///
663    /// [`Error::Snapshot`] when this database is a snapshot, and [`Error::Sqlite`]
664    /// when the failure cannot be stored.
665    ///
666    /// # Examples
667    ///
668    /// ```
669    /// use pelicode_core::Config;
670    /// use pelicode_db::{Database, PackageInfo};
671    ///
672    /// let dir = tempfile::tempdir().unwrap();
673    /// let mut db = Database::open(dir.path().join("pelicode.db"), &Config::default()).unwrap();
674    /// let package = PackageInfo {
675    ///     ecosystem: "pypi".into(),
676    ///     name: "broken".into(),
677    ///     version: "0.1".into(),
678    ///     ..PackageInfo::default()
679    /// };
680    /// db.record_failure(&package, "digest mismatch").unwrap();
681    /// assert_eq!(db.version_status("pypi", "broken", "0.1").unwrap().as_deref(), Some("failed"));
682    /// ```
683    pub fn record_failure(&mut self, pkg: &PackageInfo, error: &str) -> Result<()> {
684        self.refuse_snapshot()?;
685        let tx = self.conn.transaction()?;
686        upsert_version(&tx, pkg, "failed", Some(error))?;
687        tx.commit()?;
688        Ok(())
689    }
690
691    /// Number of rows of the main tables.
692    ///
693    /// # Errors
694    ///
695    /// [`Error::Sqlite`] when the database cannot be read.
696    ///
697    /// # Examples
698    ///
699    /// ```
700    /// # use pelicode_core::{Config, Language, extract_bytes};
701    /// # use pelicode_db::{Database, PackageInfo};
702    /// # let dir = tempfile::tempdir().unwrap();
703    /// # let config = Config { min_tokens: 0, ..Config::default() };
704    /// # let mut db = Database::open(dir.path().join("pelicode.db"), &config).unwrap();
705    /// # let source = b"def total(xs):\n    s = 0\n    for x in xs:\n        s += x * 2\n    return s\n";
706    /// # let file = extract_bytes("src/calc.py".as_ref(), Language::Python, source, &config).unwrap();
707    /// # let package = PackageInfo {
708    /// #     ecosystem: "pypi".into(), name: "calc".into(), version: "1.0".into(),
709    /// #     license_spdx: Some("MIT".into()), ..PackageInfo::default()
710    /// # };
711    /// # db.index(&package, &[file]).unwrap();
712    /// let stats = db.stats().unwrap();
713    /// assert_eq!((stats.packages, stats.versions, stats.functions), (1, 1, 1));
714    /// ```
715    pub fn stats(&self) -> Result<Stats> {
716        let count = |table: &str| -> Result<usize> {
717            let n: i64 =
718                self.conn
719                    .query_row(&format!("SELECT COUNT(*) FROM {table}"), [], |r| r.get(0))?;
720            Ok(n as usize)
721        };
722        let indexed: i64 =
723            self.conn
724                .query_row("SELECT COUNT(DISTINCT function_id) FROM minhash", [], |r| {
725                    r.get(0)
726                })?;
727        Ok(Stats {
728            ecosystems: count("ecosystem")?,
729            packages: count("package")?,
730            versions: count("package_version")?,
731            source_files: count("source_file")?,
732            functions: count("function")?,
733            indexed_functions: indexed as usize,
734        })
735    }
736
737    /// Names of the packages of `ecosystem` with at least one version, crawled
738    /// or failed, sorted: the packages to check for new versions.
739    ///
740    /// # Errors
741    ///
742    /// [`Error::Sqlite`] when the database cannot be read.
743    ///
744    /// # Examples
745    ///
746    /// ```
747    /// use pelicode_core::Config;
748    /// use pelicode_db::{Database, PackageInfo};
749    ///
750    /// let dir = tempfile::tempdir().unwrap();
751    /// let mut db = Database::open(dir.path().join("pelicode.db"), &Config::default()).unwrap();
752    /// for name in ["six", "requests"] {
753    ///     let package = PackageInfo {
754    ///         ecosystem: "pypi".into(), name: name.into(), version: "1.0".into(), ..PackageInfo::default()
755    ///     };
756    ///     db.index(&package, &[]).unwrap();
757    /// }
758    /// assert_eq!(db.package_names("pypi").unwrap(), ["requests", "six"]);
759    /// assert!(db.package_names("crates").unwrap().is_empty());
760    /// ```
761    pub fn package_names(&self, ecosystem: &str) -> Result<Vec<String>> {
762        let mut stmt = self.conn.prepare(
763            "SELECT p.name FROM package p JOIN ecosystem e ON e.id = p.ecosystem_id
764             WHERE e.name = ?1 AND p.id IN (SELECT package_id FROM package_version)
765             ORDER BY p.name",
766        )?;
767        let names = stmt.query_map([ecosystem], |r| r.get(0))?;
768        Ok(names.collect::<rusqlite::Result<_>>()?)
769    }
770
771    /// When `ecosystem` was last checked for new versions (RFC 3339), if ever.
772    ///
773    /// # Errors
774    ///
775    /// [`Error::Sqlite`] when the database cannot be read.
776    ///
777    /// # Examples
778    ///
779    /// ```
780    /// use pelicode_core::Config;
781    /// use pelicode_db::Database;
782    ///
783    /// let dir = tempfile::tempdir().unwrap();
784    /// let db = Database::open(dir.path().join("pelicode.db"), &Config::default()).unwrap();
785    /// assert_eq!(db.updates_checked_at("pypi").unwrap(), None);
786    /// db.record_updates_check("pypi").unwrap();
787    /// assert!(db.updates_checked_at("pypi").unwrap().is_some());
788    /// ```
789    pub fn updates_checked_at(&self, ecosystem: &str) -> Result<Option<String>> {
790        Ok(self
791            .meta()?
792            .remove(&format!("updates_checked_at.{ecosystem}")))
793    }
794
795    /// Records that `ecosystem` has been checked for new versions now.
796    ///
797    /// # Errors
798    ///
799    /// [`Error::Snapshot`] when this database is a snapshot, and [`Error::Sqlite`]
800    /// when it cannot be written.
801    ///
802    /// # Examples
803    ///
804    /// See [`Database::updates_checked_at`].
805    pub fn record_updates_check(&self, ecosystem: &str) -> Result<()> {
806        self.refuse_snapshot()?;
807        self.conn.execute(
808            "INSERT INTO meta (key, value) VALUES (?1, strftime('%Y-%m-%dT%H:%M:%SZ', 'now'))
809             ON CONFLICT (key) DO UPDATE SET value = excluded.value",
810            [format!("updates_checked_at.{ecosystem}")],
811        )?;
812        Ok(())
813    }
814
815    /// Names of the ecosystems that have packages, sorted.
816    ///
817    /// # Errors
818    ///
819    /// [`Error::Sqlite`] when the database cannot be read.
820    ///
821    /// # Examples
822    ///
823    /// ```
824    /// use pelicode_core::{Config, Language, extract_bytes};
825    /// use pelicode_db::{Database, PackageInfo};
826    ///
827    /// let dir = tempfile::tempdir().unwrap();
828    /// let config = Config::default();
829    /// let mut db = Database::open(dir.path().join("pelicode.db"), &config).unwrap();
830    /// let file = extract_bytes("a.py".as_ref(), Language::Python, b"def f():\n    return 1\n", &config).unwrap();
831    /// for ecosystem in ["pypi", "crates"] {
832    ///     let package = PackageInfo {
833    ///         ecosystem: ecosystem.into(), name: "demo".into(), version: "1".into(), ..PackageInfo::default()
834    ///     };
835    ///     db.index(&package, &[file.clone()]).unwrap();
836    /// }
837    /// assert_eq!(db.ecosystems().unwrap(), ["crates", "pypi"]);
838    /// ```
839    pub fn ecosystems(&self) -> Result<Vec<String>> {
840        let mut stmt = self.conn.prepare(
841            "SELECT name FROM ecosystem
842             WHERE id IN (SELECT ecosystem_id FROM package) ORDER BY name",
843        )?;
844        let names = stmt.query_map([], |r| r.get(0))?;
845        Ok(names.collect::<rusqlite::Result<_>>()?)
846    }
847
848    /// Every occurrence of every stored function, ordered by package and location.
849    ///
850    /// # Errors
851    ///
852    /// [`Error::Sqlite`] when the database cannot be read.
853    ///
854    /// # Examples
855    ///
856    /// ```
857    /// # use pelicode_core::{Config, Language, extract_bytes};
858    /// # use pelicode_db::{Database, PackageInfo};
859    /// # let dir = tempfile::tempdir().unwrap();
860    /// # let config = Config { min_tokens: 0, ..Config::default() };
861    /// # let mut db = Database::open(dir.path().join("pelicode.db"), &config).unwrap();
862    /// # let source = b"def total(xs):\n    s = 0\n    for x in xs:\n        s += x * 2\n    return s\n";
863    /// # let file = extract_bytes("src/calc.py".as_ref(), Language::Python, source, &config).unwrap();
864    /// # let package = PackageInfo {
865    /// #     ecosystem: "pypi".into(), name: "calc".into(), version: "1.0".into(),
866    /// #     license_spdx: Some("MIT".into()), ..PackageInfo::default()
867    /// # };
868    /// # db.index(&package, &[file]).unwrap();
869    /// let rows = db.functions().unwrap();
870    /// assert_eq!(rows[0].name, "total");
871    /// assert_eq!(rows[0].path, "src/calc.py");
872    /// assert_eq!(rows[0].license.as_deref(), Some("MIT"));
873    /// ```
874    pub fn functions(&self) -> Result<Vec<FunctionRow>> {
875        let mut hashes: HashMap<i64, BTreeMap<Level, String>> = HashMap::new();
876        let mut stmt = self
877            .conn
878            .prepare("SELECT function_id, level, hash FROM exact_hash")?;
879        let mut rows = stmt.query([])?;
880        while let Some(row) = rows.next()? {
881            let level: u8 = row.get(1)?;
882            let hash: Vec<u8> = row.get(2)?;
883            if let Some(level) = Level::from_u8(level) {
884                hashes
885                    .entry(row.get(0)?)
886                    .or_default()
887                    .insert(level, pelicode_core::hash::to_hex(&hash));
888            }
889        }
890
891        let mut stmt = self
892            .conn
893            .prepare("SELECT DISTINCT function_id FROM minhash")?;
894        let indexed: HashSet<i64> = stmt
895            .query_map([], |r| r.get(0))?
896            .collect::<rusqlite::Result<_>>()?;
897
898        let mut stmt = self.conn.prepare(
899            "SELECT f.id, e.name, p.name, pv.version, pv.license_spdx, vf.path, sf.language,
900                    f.name, f.kind, f.start_line, f.end_line, f.start_column, f.end_column,
901                    f.start_byte, f.end_byte, f.parent_function_id,
902                    f.token_count, f.node_count, f.depth
903             FROM function f
904             JOIN source_file sf ON sf.id = f.source_file_id
905             JOIN version_file vf ON vf.source_file_id = sf.id
906             JOIN package_version pv ON pv.id = vf.package_version_id
907             JOIN package p ON p.id = pv.package_id
908             JOIN ecosystem e ON e.id = p.ecosystem_id
909             ORDER BY e.name, p.name, pv.version, vf.path, f.start_line, f.id",
910        )?;
911        let rows = stmt.query_map([], |r| {
912            let id: i64 = r.get(0)?;
913            Ok(FunctionRow {
914                id,
915                ecosystem: r.get(1)?,
916                package: r.get(2)?,
917                version: r.get(3)?,
918                license: r.get(4)?,
919                path: r.get(5)?,
920                language: r.get(6)?,
921                name: r.get(7)?,
922                kind: r.get(8)?,
923                start_line: r.get::<_, i64>(9)? as usize,
924                end_line: r.get::<_, i64>(10)? as usize,
925                start_column: r.get::<_, i64>(11)? as usize,
926                end_column: r.get::<_, i64>(12)? as usize,
927                start_byte: r.get::<_, i64>(13)? as usize,
928                end_byte: r.get::<_, i64>(14)? as usize,
929                parent_id: r.get(15)?,
930                token_count: r.get::<_, i64>(16)? as usize,
931                node_count: r.get::<_, i64>(17)? as usize,
932                depth: r.get::<_, i64>(18)? as usize,
933                indexed: indexed.contains(&id),
934                exact_hashes: hashes.get(&id).cloned().unwrap_or_default(),
935            })
936        })?;
937        Ok(rows.collect::<rusqlite::Result<_>>()?)
938    }
939}
940
941fn upsert_version(
942    tx: &Transaction<'_>,
943    pkg: &PackageInfo,
944    status: &str,
945    error: Option<&str>,
946) -> Result<i64> {
947    tx.execute(
948        "INSERT OR IGNORE INTO ecosystem (name) VALUES (?1)",
949        [&pkg.ecosystem],
950    )?;
951    let ecosystem_id: i64 = tx.query_row(
952        "SELECT id FROM ecosystem WHERE name = ?1",
953        [&pkg.ecosystem],
954        |r| r.get(0),
955    )?;
956    let package_id: i64 = tx.query_row(
957        "INSERT INTO package (ecosystem_id, name, repository_url) VALUES (?1, ?2, ?3)
958         ON CONFLICT (ecosystem_id, name)
959         DO UPDATE SET repository_url = COALESCE(excluded.repository_url, repository_url)
960         RETURNING id",
961        params![ecosystem_id, pkg.name, pkg.repository_url],
962        |r| r.get(0),
963    )?;
964    let version_id: i64 = tx.query_row(
965        "INSERT INTO package_version
966             (package_id, version, published_at, license_spdx, archive_url, archive_sha256,
967              crawled_at, status, error)
968         VALUES (?1, ?2, ?3, ?4, ?5, ?6, strftime('%Y-%m-%dT%H:%M:%SZ', 'now'), ?7, ?8)
969         ON CONFLICT (package_id, version) DO UPDATE SET
970             published_at   = COALESCE(excluded.published_at, published_at),
971             license_spdx   = COALESCE(excluded.license_spdx, license_spdx),
972             archive_url    = COALESCE(excluded.archive_url, archive_url),
973             archive_sha256 = COALESCE(excluded.archive_sha256, archive_sha256),
974             crawled_at     = excluded.crawled_at,
975             status         = excluded.status,
976             error          = excluded.error
977         RETURNING id",
978        params![
979            package_id,
980            pkg.version,
981            pkg.published_at,
982            pkg.license_spdx,
983            pkg.archive_url,
984            pkg.archive_sha256,
985            status,
986            error
987        ],
988        |r| r.get(0),
989    )?;
990    Ok(version_id)
991}
992
993fn insert_file(
994    tx: &Transaction<'_>,
995    file: &FileFeatures,
996    summary: &mut IndexSummary,
997) -> Result<i64> {
998    tx.execute(
999        "INSERT INTO source_file (content_sha256, language, function_count) VALUES (?1, ?2, ?3)",
1000        params![
1001            file.content_sha256,
1002            file.language.name(),
1003            file.functions.len() as i64
1004        ],
1005    )?;
1006    let file_id = tx.last_insert_rowid();
1007
1008    let mut insert_function = tx.prepare_cached(
1009        "INSERT INTO function (source_file_id, parent_function_id, name, kind, start_line,
1010                               end_line, start_column, end_column, start_byte, end_byte,
1011                               token_count, node_count, depth)
1012         VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13)",
1013    )?;
1014    let mut insert_exact =
1015        tx.prepare_cached("INSERT INTO exact_hash (function_id, level, hash) VALUES (?1, ?2, ?3)")?;
1016    let mut insert_fingerprint = tx.prepare_cached(
1017        "INSERT INTO fingerprint (function_id, level, hash, position, start_line, end_line)
1018         VALUES (?1, ?2, ?3, ?4, ?5, ?6)",
1019    )?;
1020    let mut insert_minhash = tx.prepare_cached(
1021        "INSERT INTO minhash (function_id, level, signature) VALUES (?1, ?2, ?3)",
1022    )?;
1023    let mut insert_bucket = tx.prepare_cached(
1024        "INSERT INTO lsh_bucket (level, band, bucket_hash, function_id) VALUES (?1, ?2, ?3, ?4)",
1025    )?;
1026
1027    // Parents precede their children, so their ids are always known.
1028    let mut ids: Vec<i64> = Vec::with_capacity(file.functions.len());
1029    for function in &file.functions {
1030        let parent_id = function.parent.map(|i| ids[i]);
1031        insert_function.execute(params![
1032            file_id,
1033            parent_id,
1034            function.name,
1035            function.kind,
1036            function.start_line as i64,
1037            function.end_line as i64,
1038            function.start_column as i64,
1039            function.end_column as i64,
1040            function.start_byte as i64,
1041            function.end_byte as i64,
1042            function.token_count as i64,
1043            function.node_count as i64,
1044            function.depth as i64,
1045        ])?;
1046        let function_id = tx.last_insert_rowid();
1047        ids.push(function_id);
1048        summary.functions += 1;
1049        if function.is_indexed() {
1050            summary.indexed_functions += 1;
1051        }
1052
1053        for (level, hash) in &function.exact_hashes {
1054            insert_exact.execute(params![function_id, level.as_u8(), &hash.0[..]])?;
1055        }
1056        for (level, features) in &function.features {
1057            let level = level.as_u8();
1058            for fp in &features.fingerprints {
1059                insert_fingerprint.execute(params![
1060                    function_id,
1061                    level,
1062                    fp.hash as i64,
1063                    fp.position,
1064                    fp.start_line,
1065                    fp.end_line
1066                ])?;
1067            }
1068            let signature: Vec<u8> = features
1069                .minhash
1070                .iter()
1071                .flat_map(|h| h.to_le_bytes())
1072                .collect();
1073            insert_minhash.execute(params![function_id, level, signature])?;
1074            for (band, bucket) in features.lsh_buckets.iter().enumerate() {
1075                insert_bucket.execute(params![level, band as i64, *bucket as i64, function_id])?;
1076            }
1077        }
1078    }
1079    Ok(file_id)
1080}
1081
1082#[cfg(test)]
1083mod tests {
1084    use super::*;
1085    use pelicode_core::{Language, extract_source};
1086
1087    const SRC: &str = r#"
1088class A {
1089    int sum(int[] values) {
1090        int total = 0;
1091        for (int v : values) { total += v; }
1092        return total;
1093    }
1094    Runnable r() { return () -> System.out.println("x"); }
1095}
1096"#;
1097
1098    fn file(path: &str) -> FileFeatures {
1099        let config = Config {
1100            min_tokens: 10,
1101            ..Config::default()
1102        };
1103        FileFeatures {
1104            path: path.into(),
1105            language: Language::Java,
1106            content_sha256: pelicode_core::extract::sha256_hex(SRC.as_bytes()),
1107            functions: extract_source(Language::Java, SRC.as_bytes(), &config).unwrap(),
1108        }
1109    }
1110
1111    fn pkg(version: &str) -> PackageInfo {
1112        PackageInfo {
1113            ecosystem: "local".into(),
1114            name: "demo".into(),
1115            version: version.into(),
1116            license_spdx: Some("MIT".into()),
1117            ..PackageInfo::default()
1118        }
1119    }
1120
1121    #[test]
1122    fn indexes_and_deduplicates_files() {
1123        let dir = tempfile::tempdir().unwrap();
1124        let mut db = Database::open(dir.path().join("t.db"), &Config::default()).unwrap();
1125
1126        let first = db.index(&pkg("1.0"), &[file("A.java")]).unwrap();
1127        assert_eq!(first.new_files, 1);
1128        assert_eq!(first.functions, 3);
1129
1130        // Same content in another version and path: stored once, linked twice.
1131        let second = db.index(&pkg("2.0"), &[file("src/A.java")]).unwrap();
1132        assert_eq!(second.new_files, 0);
1133
1134        // Re-indexing a version replaces its file list.
1135        db.index(&pkg("2.0"), &[file("src/A.java")]).unwrap();
1136
1137        let stats = db.stats().unwrap();
1138        assert_eq!(stats.versions, 2);
1139        assert_eq!(stats.source_files, 1);
1140        assert_eq!(stats.functions, 3);
1141
1142        let rows = db.functions().unwrap();
1143        assert_eq!(rows.len(), 6);
1144        let lambda = rows.iter().find(|r| r.kind == "lambda_expression").unwrap();
1145        assert!(lambda.parent_id.is_some());
1146        assert_eq!(rows[0].exact_hashes.len(), 4);
1147        let sum = rows.iter().find(|r| r.name == "sum").unwrap();
1148        assert_eq!(
1149            &SRC[sum.start_byte..sum.end_byte],
1150            "int sum(int[] values) {\n        int total = 0;\n        for (int v : values) { total += v; }\n        return total;\n    }"
1151        );
1152        assert_eq!((sum.start_column, sum.end_column), (5, 5));
1153        let lines: (i64, i64) = db
1154            .conn
1155            .query_row(
1156                "SELECT MIN(start_line), MAX(end_line) FROM fingerprint WHERE function_id = ?1",
1157                [sum.id],
1158                |r| Ok((r.get(0)?, r.get(1)?)),
1159            )
1160            .unwrap();
1161        // Winnowing need not select the last k-gram, so the range may end before the function.
1162        assert!(
1163            lines.0 >= 3 && lines.1 <= 7 && lines.0 < lines.1,
1164            "{lines:?}"
1165        );
1166        assert_eq!(rows[0].license.as_deref(), Some("MIT"));
1167    }
1168
1169    #[test]
1170    fn records_crawl_status() {
1171        let dir = tempfile::tempdir().unwrap();
1172        let mut db = Database::open(dir.path().join("t.db"), &Config::default()).unwrap();
1173        assert_eq!(db.version_status("local", "demo", "1.0").unwrap(), None);
1174        db.record_failure(&pkg("1.0"), "network down").unwrap();
1175        assert_eq!(
1176            db.version_status("local", "demo", "1.0")
1177                .unwrap()
1178                .as_deref(),
1179            Some("failed")
1180        );
1181        db.index(&pkg("1.0"), &[file("A.java")]).unwrap();
1182        assert_eq!(
1183            db.version_status("local", "demo", "1.0")
1184                .unwrap()
1185                .as_deref(),
1186            Some("indexed")
1187        );
1188        let error: Option<String> = db
1189            .conn
1190            .query_row("SELECT error FROM package_version", [], |r| r.get(0))
1191            .unwrap();
1192        assert_eq!(error, None);
1193    }
1194
1195    #[test]
1196    fn counts_packages_per_hash() {
1197        let dir = tempfile::tempdir().unwrap();
1198        let mut db = Database::open(dir.path().join("t.db"), &Config::default()).unwrap();
1199        let other = |name: &str| PackageInfo {
1200            name: name.into(),
1201            ..pkg("1.0")
1202        };
1203        // Two versions of one package count once.
1204        db.index(&pkg("1.0"), &[file("A.java")]).unwrap();
1205        db.index(&pkg("2.0"), &[file("A.java")]).unwrap();
1206        assert!(db.hash_stats_info().unwrap().stale);
1207
1208        let info = db.refresh_hash_stats().unwrap();
1209        assert!(!info.stale);
1210        assert_eq!(info.packages, 1);
1211        assert_eq!(info.common_exact_hashes, 0);
1212
1213        db.index(&other("copy"), &[file("B.java")]).unwrap();
1214        assert!(db.hash_stats_info().unwrap().stale);
1215        let info = db.refresh_hash_stats().unwrap();
1216        assert_eq!(info.packages, 2);
1217        assert!(info.common_exact_hashes > 0);
1218        assert!(info.common_fingerprints > 0);
1219
1220        let sum = db
1221            .functions()
1222            .unwrap()
1223            .into_iter()
1224            .find(|r| r.name == "sum")
1225            .unwrap();
1226        let hash = pelicode_core::Hash128::from_hex(&sum.exact_hashes[&Level::N1]).unwrap();
1227        assert_eq!(db.exact_package_count(Level::N1, &hash).unwrap(), 2);
1228        assert_eq!(
1229            db.exact_package_count(Level::N1, &pelicode_core::Hash128([0; 16]))
1230                .unwrap(),
1231            1
1232        );
1233        let common = db.common_exact_hashes(Level::N1, 0, 10).unwrap();
1234        assert!(
1235            common
1236                .iter()
1237                .all(|c| c.package_count == 2 && c.sample.is_some())
1238        );
1239        assert!(!db.common_fingerprints(Level::N1, 10).unwrap().is_empty());
1240    }
1241
1242    #[test]
1243    fn clean_keeps_files_linked_elsewhere() {
1244        let dir = tempfile::tempdir().unwrap();
1245        let mut db = Database::open(dir.path().join("t.db"), &Config::default()).unwrap();
1246        db.index(&pkg("1.0"), &[file("A.java")]).unwrap();
1247        db.index(&pkg("2.0"), &[file("A.java")]).unwrap();
1248        // 1.0 no longer has the file, but 2.0 still does.
1249        db.index(&pkg("1.0"), &[]).unwrap();
1250        assert_eq!(db.unreferenced().unwrap(), CleanSummary::default());
1251        assert_eq!(db.clean().unwrap(), CleanSummary::default());
1252        assert_eq!(db.stats().unwrap().functions, 3);
1253
1254        db.index(&pkg("2.0"), &[]).unwrap();
1255        let removed = db.clean().unwrap();
1256        assert_eq!((removed.files, removed.functions), (1, 3));
1257        let left: i64 = db
1258            .conn
1259            .query_row(
1260                "SELECT (SELECT COUNT(*) FROM exact_hash) + (SELECT COUNT(*) FROM fingerprint)
1261                      + (SELECT COUNT(*) FROM minhash) + (SELECT COUNT(*) FROM lsh_bucket)",
1262                [],
1263                |r| r.get(0),
1264            )
1265            .unwrap();
1266        assert_eq!(left, 0);
1267        db.vacuum().unwrap();
1268    }
1269
1270    fn stats_tables(db: &Database) -> Vec<(String, u8, Vec<u8>, i64, i64)> {
1271        let mut rows = Vec::new();
1272        for (table, hash) in [
1273            ("exact_stats", "hash"),
1274            ("fingerprint_stats", "CAST(hash AS BLOB)"),
1275        ] {
1276            let mut stmt = db
1277                .conn
1278                .prepare(&format!(
1279                    "SELECT level, {hash}, package_count, function_count FROM {table}
1280                     ORDER BY level, hash"
1281                ))
1282                .unwrap();
1283            let found = stmt
1284                .query_map([], |r| {
1285                    Ok((
1286                        table.to_string(),
1287                        r.get(0)?,
1288                        r.get(1)?,
1289                        r.get(2)?,
1290                        r.get(3)?,
1291                    ))
1292                })
1293                .unwrap()
1294                .collect::<rusqlite::Result<Vec<_>>>()
1295                .unwrap();
1296            rows.extend(found);
1297        }
1298        rows
1299    }
1300
1301    #[test]
1302    fn incremental_statistics_equal_a_rebuild() {
1303        let dir = tempfile::tempdir().unwrap();
1304        let mut db = Database::open(dir.path().join("t.db"), &Config::default()).unwrap();
1305        let config = Config {
1306            min_tokens: 10,
1307            ..Config::default()
1308        };
1309        let source = |n: usize| {
1310            let body: String = (0..n)
1311                .map(|i| format!("        total += values[{i}] * {i};\n"))
1312                .collect();
1313            format!(
1314                "class C{n} {{\n    int f(int[] values) {{\n        int total = 0;\n{body}        return total;\n    }}\n}}\n"
1315            )
1316        };
1317        let file = |name: &str, n: usize| {
1318            let src = source(n);
1319            FileFeatures {
1320                path: name.into(),
1321                language: Language::Java,
1322                content_sha256: pelicode_core::extract::sha256_hex(src.as_bytes()),
1323                functions: extract_source(Language::Java, src.as_bytes(), &config).unwrap(),
1324            }
1325        };
1326        let package = |name: &str, version: &str| PackageInfo {
1327            ecosystem: "maven".into(),
1328            name: name.into(),
1329            version: version.into(),
1330            ..PackageInfo::default()
1331        };
1332
1333        db.index(&package("a", "1"), &[file("A.java", 3), file("B.java", 4)])
1334            .unwrap();
1335        db.index(&package("b", "1"), &[file("A.java", 3)]).unwrap();
1336        db.index(&package("e", "1"), &[file("E.java", 7)]).unwrap();
1337        db.index(&package("f", "1"), &[file("E.java", 7)]).unwrap();
1338        let info = db.refresh_hash_stats().unwrap();
1339        assert!(!info.incremental || info.pending_hashes == 0);
1340
1341        // New packages, shared files, a version whose files are replaced
1342        // (unlinking A from package a), and a cleanup in between.
1343        db.index(&package("c", "1"), &[file("A.java", 3), file("C.java", 5)])
1344            .unwrap();
1345        db.index(&package("a", "1"), &[file("B.java", 4)]).unwrap();
1346        db.index(&package("a", "2"), &[file("C.java", 5)]).unwrap();
1347        db.index(&package("d", "1"), &[file("B.java", 4), file("D.java", 6)])
1348            .unwrap();
1349        // E is no longer shared: its hashes must leave the statistics.
1350        db.index(&package("e", "1"), &[]).unwrap();
1351        db.clean().unwrap();
1352
1353        let info = db.hash_stats_info().unwrap();
1354        assert!(info.incremental && info.stale && info.pending_hashes > 0);
1355        db.refresh_hash_stats().unwrap();
1356        let incremental = stats_tables(&db);
1357        assert_eq!(db.hash_stats_info().unwrap().pending_hashes, 0);
1358
1359        db.rebuild_hash_stats().unwrap();
1360        assert_eq!(incremental, stats_tables(&db));
1361        assert!(!incremental.is_empty());
1362        let shared = |db: &Database, path: &str| -> i64 {
1363            db.conn
1364                .query_row(
1365                    "SELECT COUNT(*) FROM exact_stats s
1366                     JOIN exact_hash h ON h.level = s.level AND h.hash = s.hash
1367                     JOIN function f ON f.id = h.function_id
1368                     JOIN version_file vf ON vf.source_file_id = f.source_file_id
1369                     WHERE vf.path = ?1",
1370                    [path],
1371                    |r| r.get(0),
1372                )
1373                .unwrap()
1374        };
1375        assert_eq!(shared(&db, "E.java"), 0);
1376        assert!(shared(&db, "A.java") > 0);
1377    }
1378
1379    #[test]
1380    fn migrates_version_4_databases() {
1381        let dir = tempfile::tempdir().unwrap();
1382        let path = dir.path().join("t.db");
1383        let downgrade = |stale: bool| {
1384            let mut db = Database::open(&path, &Config::default()).unwrap();
1385            db.index(&pkg("1.0"), &[file("A.java")]).unwrap();
1386            db.rebuild_hash_stats().unwrap();
1387            if stale {
1388                db.index(&pkg("2.0"), &[file("A.java")]).unwrap();
1389            }
1390            db.conn
1391                .execute_batch(
1392                    "UPDATE meta SET value = '4' WHERE key = 'schema_version';
1393                     DELETE FROM meta WHERE key = 'stats_pending_complete';
1394                     DROP TABLE stats_pending_exact;
1395                     DROP TABLE stats_pending_fingerprint;",
1396                )
1397                .unwrap();
1398        };
1399
1400        downgrade(false);
1401        let db = Database::open(&path, &Config::default()).unwrap();
1402        assert_eq!(
1403            db.meta().unwrap()["schema_version"],
1404            SCHEMA_VERSION.to_string()
1405        );
1406        assert!(db.hash_stats_info().unwrap().incremental);
1407        drop(db);
1408
1409        std::fs::remove_file(&path).unwrap();
1410        downgrade(true);
1411        let mut db = Database::open(&path, &Config::default()).unwrap();
1412        // Something was indexed after the statistics without being tracked.
1413        assert!(!db.hash_stats_info().unwrap().incremental);
1414        db.refresh_hash_stats().unwrap();
1415        assert!(db.hash_stats_info().unwrap().incremental);
1416    }
1417
1418    #[test]
1419    fn rejects_incompatible_parameters() {
1420        let dir = tempfile::tempdir().unwrap();
1421        let path = dir.path().join("t.db");
1422        Database::open(&path, &Config::default()).unwrap();
1423        let other = Config {
1424            kgram: 7,
1425            ..Config::default()
1426        };
1427        assert!(matches!(
1428            Database::open(&path, &other),
1429            Err(Error::Incompatible { .. })
1430        ));
1431    }
1432}