Add gitlab support

Author: GuillaumeGomez

README.md

@@ -191,7 +191,10 @@ docker-compose run -- database add-directory <DIRECTORY> [PREFIX]
 # You need to set CRATESFYI_GITHUB_USERNAME, CRATESFYI_GITHUB_ACCESSTOKEN
 # environment variables in order to run this command.
 # You can set this environment variables in ~/.cratesfyi.env file.
-docker-compose run -- database update-github-fields
+cargo run -- database update-github-fields
+# Updates github stats for crates.
+# You need to set CRATESFYI_GITLAB_ACCESSTOKEN environment variable.
+cargo run -- database update-gitlab-fields
 ```
 
 If you want to explore or edit database manually, you can connect to the database

src/bin/cratesfyi.rs

@@ -377,6 +377,12 @@ enum DatabaseSubcommand {
     /// Backfill GitHub stats for crates.
     BackfillGithubStats,
 
+    /// Updates gitlab stats for crates.
+    UpdateGitlabFields,
+
+    /// Backfill Gitlab stats for crates.
+    BackfillGitlabStats,
+
     /// Updates info for a crate from the registry's API
     UpdateCrateRegistryFields {
         #[structopt(name = "CRATE")]
@@ -436,6 +442,18 @@ impl DatabaseSubcommand {
                     .backfill_repositories()?;
             }
 
+            Self::UpdateGitlabFields => {
+                docs_rs::utils::GitlabUpdater::new(ctx.config()?, ctx.pool()?)?
+                    .ok_or_else(|| failure::format_err!("missing Gitlab token"))?
+                    .update_all_crates()?;
+            }
+
+            Self::BackfillGitlabStats => {
+                docs_rs::utils::GitlabUpdater::new(ctx.config()?, ctx.pool()?)?
+                    .ok_or_else(|| failure::format_err!("missing Gitlab token"))?
+                    .backfill_repositories()?;
+            }
+
             Self::UpdateCrateRegistryFields { name } => {
                 let index = ctx.index()?;
 

src/config.rs

@@ -30,6 +30,10 @@ pub struct Config {
     pub(crate) github_accesstoken: Option<String>,
     pub(crate) github_updater_min_rate_limit: u32,
 
+    // Gitlab authentication
+    pub(crate) gitlab_accesstoken: Option<String>,
+    pub(crate) gitlab_updater_min_rate_limit: u32,
+
     // Max size of the files served by the docs.rs frontend
     pub(crate) max_file_size: usize,
     pub(crate) max_file_size_html: usize,
@@ -77,6 +81,9 @@ impl Config {
             github_accesstoken: maybe_env("CRATESFYI_GITHUB_ACCESSTOKEN")?,
             github_updater_min_rate_limit: env("DOCSRS_GITHUB_UPDATER_MIN_RATE_LIMIT", 2500)?,
 
+            gitlab_accesstoken: maybe_env("CRATESFYI_GITLAB_ACCESSTOKEN")?,
+            gitlab_updater_min_rate_limit: env("DOCSRS_GITLAB_UPDATER_MIN_RATE_LIMIT", 2500)?,
+
             max_file_size: env("DOCSRS_MAX_FILE_SIZE", 50 * 1024 * 1024)?,
             max_file_size_html: env("DOCSRS_MAX_FILE_SIZE_HTML", 50 * 1024 * 1024)?,
             // LOL HTML only uses as much memory as the size of the start tag!

src/db/add_package.rs

@@ -37,7 +37,7 @@ pub(crate) fn add_package_into_database(
     has_docs: bool,
     has_examples: bool,
     compression_algorithms: std::collections::HashSet<CompressionAlgorithm>,
-    github_repo: Option<String>,
+    repository_id: Option<i32>,
 ) -> Result<i32> {
     debug!("Adding package into database");
     let crate_id = initialize_package_in_database(conn, metadata_pkg)?;
@@ -55,7 +55,8 @@ pub(crate) fn add_package_into_database(
             homepage_url, description, description_long, readme,
             authors, keywords, have_examples, downloads, files,
             doc_targets, is_library, doc_rustc_version,
-            documentation_url, default_target, features, github_repo
+            documentation_url, default_target, features,
+            repository
          )
          VALUES (
             $1,  $2,  $3,  $4,  $5,  $6,  $7,  $8,  $9,
@@ -87,7 +88,7 @@ pub(crate) fn add_package_into_database(
                 documentation_url = $24,
                 default_target = $25,
                 features = $26,
-                github_repo = $27
+                repository = $27
          RETURNING id",
         &[
             &crate_id,
@@ -116,7 +117,7 @@ pub(crate) fn add_package_into_database(
             &metadata_pkg.documentation,
             &default_target,
             &features,
-            &github_repo,
+            &repository_id,
         ],
     )?;
 

src/db/migrate.rs

@@ -618,6 +618,88 @@ pub fn migrate(version: Option<Version>, conn: &mut Client) -> CratesfyiResult<(
             DROP INDEX github_repos_stars_idx;
             ",
         ),
+        migration!(
+            context,
+            27,
+            // description
+            "Add gitlab handling: creation of the new repositories table which replaces and extend \
+             github_repos",
+            // upgrade query
+            "
+                CREATE TABLE repositories (
+                    id SERIAL PRIMARY KEY,
+                    host VARCHAR NOT NULL,
+                    host_id VARCHAR NOT NULL,
+                    name VARCHAR NOT NULL,
+                    description VARCHAR,
+                    last_commit TIMESTAMPTZ,
+                    stars INT NOT NULL,
+                    forks INT NOT NULL,
+                    issues INT NOT NULL,
+                    updated_at TIMESTAMPTZ NOT NULL,
+                    UNIQUE (host, host_id)
+                );
+
+                ALTER TABLE releases ADD COLUMN repository INTEGER
+                    REFERENCES repositories(id) ON DELETE SET NULL;
+
+                INSERT INTO repositories(host, host_id, name, description, last_commit, stars, forks, issues, updated_at)
+                    SELECT 'github', id, name, description, last_commit, stars, forks, issues, updated_at
+                    FROM github_repos;
+
+                UPDATE releases
+                    SET repository = repositories.id
+                FROM repositories
+                WHERE releases.github_repo IS NOT NULL AND repositories.host_id = releases.github_repo;
+
+                DROP INDEX releases_github_repo_idx;
+                DROP INDEX github_repos_stars_idx;
+
+                CREATE INDEX releases_github_repo_idx ON releases(repository);
+                CREATE INDEX github_repos_stars_idx ON repositories(stars DESC);
+
+                ALTER TABLE releases
+                    DROP COLUMN github_repo;
+
+                DROP TABLE github_repos;
+            ",
+            // downgrade query
+            "
+                CREATE TABLE github_repos (
+                    id VARCHAR PRIMARY KEY NOT NULL,
+                    name VARCHAR NOT NULL,
+                    description VARCHAR,
+                    last_commit TIMESTAMPTZ,
+                    stars INT NOT NULL,
+                    forks INT NOT NULL,
+                    issues INT NOT NULL,
+                    updated_at TIMESTAMPTZ NOT NULL
+                );
+
+                ALTER TABLE releases ADD COLUMN github_repo VARCHAR
+                    REFERENCES github_repos(id) ON DELETE SET NULL;
+
+                INSERT INTO github_repos(id, name, description, last_commit, stars, forks, issues, updated_at)
+                    SELECT host_id, name, description, last_commit, stars, forks, issues, updated_at
+                    FROM repositories WHERE repositories.host = 'github';
+
+                UPDATE releases
+                    SET github_repo = repositories.host_id
+                FROM repositories
+                WHERE repositories.host_id = releases.github_repo AND releases.repository IS NOT NULL AND repositories.host = 'github';
+
+                DROP INDEX releases_github_repo_idx;
+                DROP INDEX github_repos_stars_idx;
+
+                CREATE INDEX releases_github_repo_idx ON releases (github_repo);
+                CREATE INDEX github_repos_stars_idx ON github_repos(stars DESC);
+
+                ALTER TABLE releases
+                    DROP COLUMN repository;
+
+                DROP TABLE repository;
+            "
+        ),
     ];
 
     for migration in migrations {

src/docbuilder/rustwide_builder.rs

@@ -7,7 +7,9 @@ use crate::docbuilder::{crates::crates_from_path, Limits};
 use crate::error::Result;
 use crate::index::api::ReleaseData;
 use crate::storage::CompressionAlgorithms;
-use crate::utils::{copy_doc_dir, parse_rustc_version, CargoMetadata, GithubUpdater};
+use crate::utils::{
+    copy_doc_dir, parse_rustc_version, CargoMetadata, GithubUpdater, GitlabUpdater,
+};
 use crate::{db::blacklist::is_blacklisted, utils::MetadataPackage};
 use crate::{Config, Context, Index, Metrics, Storage};
 use docsrs_metadata::{Metadata, DEFAULT_TARGETS, HOST_TARGET};
@@ -393,7 +395,10 @@ impl RustwideBuilder {
                 };
 
                 let cargo_metadata = res.cargo_metadata.root();
-                let github_repo = self.get_github_repo(&mut conn, cargo_metadata)?;
+                let repository = match self.get_github_repo(&mut conn, cargo_metadata)? {
+                    Some(r) => Some(r),
+                    None => self.get_gitlab_repo(&mut conn, cargo_metadata)?,
+                };
 
                 let release_id = add_package_into_database(
                     &mut conn,
@@ -407,7 +412,7 @@ impl RustwideBuilder {
                     has_docs,
                     has_examples,
                     algs,
-                    github_repo,
+                    repository,
                 )?;
 
                 if let Some(doc_coverage) = res.doc_coverage {
@@ -678,7 +683,7 @@ impl RustwideBuilder {
         &self,
         conn: &mut Client,
         metadata: &MetadataPackage,
-    ) -> Result<Option<String>> {
+    ) -> Result<Option<i32>> {
         let updater = match GithubUpdater::new(self.config.clone(), self.db.clone())? {
             Some(updater) => updater,
             None => {
@@ -701,6 +706,34 @@ impl RustwideBuilder {
             }
         }
     }
+
+    fn get_gitlab_repo(
+        &self,
+        conn: &mut Client,
+        metadata: &MetadataPackage,
+    ) -> Result<Option<i32>> {
+        let updater = match GitlabUpdater::new(self.config.clone(), self.db.clone())? {
+            Some(updater) => updater,
+            None => {
+                warn!("did not collect Gitlab stats as no token was provided");
+                return Ok(None);
+            }
+        };
+        let repo = match &metadata.repository {
+            Some(url) => url,
+            None => {
+                debug!("did not collect Gitlab stats as no repository URL was present");
+                return Ok(None);
+            }
+        };
+        match updater.load_repository(conn, repo) {
+            Ok(repo) => Ok(repo),
+            Err(err) => {
+                warn!("failed to collect Gitlab stats: {}", err);
+                Ok(None)
+            }
+        }
+    }
 }
 
 struct FullBuildResult {

src/test/fakes.rs

@@ -316,7 +316,7 @@ impl<'a> FakeRelease<'a> {
             }
         }
 
-        let github_repo = match self.github_stats {
+        let repository = match self.github_stats {
             Some(stats) => Some(stats.create(&mut self.db.conn())?),
             None => None,
         };
@@ -338,7 +338,7 @@ impl<'a> FakeRelease<'a> {
             self.has_docs,
             self.has_examples,
             algs,
-            github_repo,
+            repository,
         )?;
         crate::db::update_crate_data_in_database(
             &mut db.conn(),
@@ -366,13 +366,13 @@ struct FakeGithubStats {
 impl FakeGithubStats {
     fn create(&self, conn: &mut Client) -> Result<String, Error> {
         let existing_count: i64 = conn
-            .query_one("SELECT COUNT(*) FROM github_repos;", &[])?
+            .query_one("SELECT COUNT(*) FROM repositories;", &[])?
             .get(0);
         let id = base64::encode(format!("FAKE ID {}", existing_count));
 
         conn.execute(
-            "INSERT INTO github_repos (id, name, description, last_commit, stars, forks, issues, updated_at)
-             VALUES ($1, $2, 'Fake description!', NOW(), $3, $4, $5, NOW());",
+            "INSERT INTO repositories (host, host_id, name, description, last_commit, stars, forks, issues, updated_at)
+             VALUES ('github', $1, $2, 'Fake description!', NOW(), $3, $4, $5, NOW());",
             &[&id, &self.repo, &self.stars, &self.forks, &self.issues],
         )?;
 

src/utils/daemon.rs

@@ -3,7 +3,7 @@
 //! This daemon will start web server, track new packages and build them
 
 use crate::{
-    utils::{queue_builder, update_release_activity, GithubUpdater},
+    utils::{queue_builder, update_release_activity, GithubUpdater, GitlabUpdater},
     Context, DocBuilder, RustwideBuilder,
 };
 use chrono::{Timelike, Utc};
@@ -93,7 +93,7 @@ pub fn start_daemon(context: &dyn Context, enable_registry_watcher: bool) -> Res
         },
     )?;
 
-    if let Some(github_updater) = GithubUpdater::new(config, context.pool()?)? {
+    if let Some(github_updater) = GithubUpdater::new(config.clone(), context.pool()?)? {
         cron(
             "github stats updater",
             Duration::from_secs(60 * 60),
@@ -106,6 +106,19 @@ pub fn start_daemon(context: &dyn Context, enable_registry_watcher: bool) -> Res
         log::warn!("GitHub stats updater not started as no token was provided");
     }
 
+    if let Some(gitlab_updater) = GitlabUpdater::new(config, context.pool()?)? {
+        cron(
+            "gitlab stats updater",
+            Duration::from_secs(60 * 60),
+            move || {
+                gitlab_updater.update_all_crates()?;
+                Ok(())
+            },
+        )?;
+    } else {
+        log::warn!("Gitlab stats updater not started as no token was provided");
+    }
+
     // Never returns; `server` blocks indefinitely when dropped
     // NOTE: if a failure occurred earlier in `start_daemon`, the server will _not_ be joined -
     // instead it will get killed when the process exits.