/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use anyhow::Result;
use async_trait::async_trait;
use blobstore::Blobstore;
use blobstore::BlobstoreGetData;
use blobstore::BlobstoreIsPresent;
use blobstore::BlobstorePutOps;
use blobstore::OverwriteStatus;
use blobstore::PutBehaviour;
use context::CoreContext;
use mononoke_types::BlobstoreBytes;
mod errors;
pub use crate::errors::ErrorKind;

/// A layer over an existing blobstore that prevents writes.
#[derive(Debug)]
pub struct ReadOnlyBlobstore<T> {
    blobstore: T,
}

impl<T: std::fmt::Display> std::fmt::Display for ReadOnlyBlobstore<T> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "ReadOnlyBlobstore<{}>", &self.blobstore)
    }
}

impl<T> ReadOnlyBlobstore<T> {
    pub fn new(blobstore: T) -> Self {
        Self { blobstore }
    }
}

#[async_trait]
impl<T: Blobstore> Blobstore for ReadOnlyBlobstore<T> {
    #[inline]
    async fn get<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: &'a str,
    ) -> Result<Option<BlobstoreGetData>> {
        self.blobstore.get(ctx, key).await
    }

    #[inline]
    async fn put<'a>(
        &'a self,
        _ctx: &'a CoreContext,
        key: String,
        _value: BlobstoreBytes,
    ) -> Result<()> {
        Err(ErrorKind::ReadOnlyPut(key).into())
    }

    #[inline]
    async fn is_present<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: &'a str,
    ) -> Result<BlobstoreIsPresent> {
        self.blobstore.is_present(ctx, key).await
    }
}

#[async_trait]
impl<T: BlobstorePutOps> BlobstorePutOps for ReadOnlyBlobstore<T> {
    async fn put_explicit<'a>(
        &'a self,
        _ctx: &'a CoreContext,
        key: String,
        _value: BlobstoreBytes,
        _put_behaviour: PutBehaviour,
    ) -> Result<OverwriteStatus> {
        Err(ErrorKind::ReadOnlyPut(key).into())
    }

    async fn put_with_status<'a>(
        &'a self,
        _ctx: &'a CoreContext,
        key: String,
        _value: BlobstoreBytes,
    ) -> Result<OverwriteStatus> {
        Err(ErrorKind::ReadOnlyPut(key).into())
    }
}

#[cfg(test)]
mod test {
    use borrowed::borrowed;
    use fbinit::FacebookInit;
    use memblob::Memblob;

    use super::*;

    #[fbinit::test]
    async fn test_error_on_write(fb: FacebookInit) {
        let ctx = CoreContext::test_mock(fb);
        borrowed!(ctx);
        let base = Memblob::default();
        let wrapper = ReadOnlyBlobstore::new(base.clone());
        let key = "foobar";

        let r = wrapper
            .put(
                ctx,
                key.to_owned(),
                BlobstoreBytes::from_bytes("test foobar"),
            )
            .await;
        assert!(r.is_err());
        let base_present = base
            .is_present(ctx, key)
            .await
            .unwrap()
            .assume_not_found_if_unsure();
        assert!(!base_present);
    }

    #[fbinit::test]
    async fn test_error_on_put_with_status(fb: FacebookInit) {
        let ctx = CoreContext::test_mock(fb);
        borrowed!(ctx);
        let base = Memblob::default();
        let wrapper = ReadOnlyBlobstore::new(base.clone());
        let key = "foobar";

        let r = wrapper
            .put_with_status(
                ctx,
                key.to_owned(),
                BlobstoreBytes::from_bytes("test foobar"),
            )
            .await;
        assert!(r.is_err());
        let base_present = base
            .is_present(ctx, key)
            .await
            .unwrap()
            .assume_not_found_if_unsure();
        assert!(!base_present);
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use thiserror::Error;

#[derive(Debug, Error)]
pub enum ErrorKind {
    #[error("The blob {0} is censored.\n Task/Sev: {1}")]
    Censored(String, String),
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

mod errors;
mod redaction_config_blobstore;
pub mod store;

use std::num::NonZeroU64;
use std::ops::Deref;
use std::sync::Arc;

use anyhow::Error;
use anyhow::Result;
use async_trait::async_trait;
use blobstore::Blobstore;
use blobstore::BlobstoreGetData;
use blobstore::BlobstoreIsPresent;
use context::CoreContext;
use mononoke_types::BlobstoreBytes;
use scuba_ext::MononokeScubaSampleBuilder;
use slog::debug;
use tunables::tunables;

pub use crate::errors::ErrorKind;
pub use crate::redaction_config_blobstore::ArcRedactionConfigBlobstore;
pub use crate::redaction_config_blobstore::RedactionConfigBlobstore;
pub use crate::store::RedactedBlobs;
pub use crate::store::RedactedMetadata;
pub use crate::store::SqlRedactedContentStore;

pub mod config {
    pub const GET_OPERATION: &str = "GET";
    pub const PUT_OPERATION: &str = "PUT";
}

#[derive(Debug, Clone)]
pub struct RedactedBlobstoreConfigInner {
    redacted: Option<Arc<RedactedBlobs>>,
    scuba_builder: MononokeScubaSampleBuilder,
}

#[derive(Debug, Clone)]
pub struct RedactedBlobstoreConfig {
    inner: Arc<RedactedBlobstoreConfigInner>,
}

impl Deref for RedactedBlobstoreConfig {
    type Target = RedactedBlobstoreConfigInner;

    fn deref(&self) -> &Self::Target {
        &self.inner
    }
}

impl RedactedBlobstoreConfig {
    pub fn new(
        redacted: Option<Arc<RedactedBlobs>>,
        scuba_builder: MononokeScubaSampleBuilder,
    ) -> Self {
        Self {
            inner: Arc::new(RedactedBlobstoreConfigInner {
                redacted,
                scuba_builder,
            }),
        }
    }
}

#[derive(Debug)]
pub struct RedactedBlobstoreInner<T> {
    blobstore: T,
    config: RedactedBlobstoreConfig,
}

impl<T: std::fmt::Display> std::fmt::Display for RedactedBlobstoreInner<T> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "RedactedBlobstoreInner<{}>", &self.blobstore)
    }
}

// A wrapper for any blobstore, which provides a verification layer for the redacted blobs.
// The goal is to deny access to fetch sensitive data from the repository.
#[derive(Debug, Clone)]
pub struct RedactedBlobstore<T> {
    inner: Arc<RedactedBlobstoreInner<T>>,
}

impl<T: std::fmt::Display> std::fmt::Display for RedactedBlobstore<T> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "RedactedBlobstore<{}>", &self.inner.blobstore)
    }
}

impl<T> Deref for RedactedBlobstore<T> {
    type Target = RedactedBlobstoreInner<T>;

    #[inline]
    fn deref(&self) -> &Self::Target {
        &self.inner
    }
}

impl<T: Blobstore> RedactedBlobstore<T> {
    pub fn new(blobstore: T, config: RedactedBlobstoreConfig) -> Self {
        Self {
            inner: Arc::new(RedactedBlobstoreInner::new(blobstore, config)),
        }
    }

    pub fn boxed(&self) -> Arc<dyn Blobstore>
    where
        T: 'static,
    {
        self.inner.clone()
    }

    pub fn as_parts(&self) -> (T, RedactedBlobstoreConfig)
    where
        T: Clone,
    {
        (self.blobstore.clone(), self.config.clone())
    }

    pub fn as_inner_unredacted(&self) -> &T {
        &self.blobstore
    }
}

impl<T: Blobstore> RedactedBlobstoreInner<T> {
    pub fn new(blobstore: T, config: RedactedBlobstoreConfig) -> Self {
        Self { blobstore, config }
    }

    // Checks for access to this key, then yields the blobstore if access is allowed.
    pub fn access_blobstore<'s: 'a, 'a>(
        &'s self,
        ctx: &'a CoreContext,
        key: &'a str,
        operation: &'static str,
    ) -> Result<&'s T> {
        match &self.config.redacted {
            Some(redacted) => {
                redacted
                    .redacted()
                    .get(key)
                    .map_or(Ok(&self.blobstore), |metadata| {
                        debug!(
                            ctx.logger(),
                            "{} operation with redacted blobstore with key {:?}", operation, key
                        );
                        self.to_scuba_redacted_blob_accessed(ctx, key, operation);

                        if metadata.log_only {
                            Ok(&self.blobstore)
                        } else {
                            Err(
                                ErrorKind::Censored(key.to_string(), metadata.task.to_string())
                                    .into(),
                            )
                        }
                    })
            }
            None => Ok(&self.blobstore),
        }
    }

    pub fn to_scuba_redacted_blob_accessed(&self, ctx: &CoreContext, key: &str, operation: &str) {
        let sampling_rate = tunables()
            .redacted_logging_sampling_rate()
            .unwrap_or_default()
            .try_into()
            .ok()
            .and_then(NonZeroU64::new);

        let mut scuba_builder = self.config.scuba_builder.clone();

        if let Some(sampling_rate) = sampling_rate {
            scuba_builder.sampled(sampling_rate);
        } else {
            scuba_builder.unsampled();
        }

        scuba_builder
            .add("operation", operation)
            .add("key", key.to_string())
            .add("session_uuid", ctx.metadata().session_id().to_string());

        if let Some(unix_username) = ctx.metadata().unix_name() {
            scuba_builder.add("unix_username", unix_username);
        }

        scuba_builder.log();
    }
}

#[async_trait]
impl<T: Blobstore> Blobstore for RedactedBlobstoreInner<T> {
    async fn get<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: &'a str,
    ) -> Result<Option<BlobstoreGetData>> {
        let blobstore = self.access_blobstore(ctx, key, config::GET_OPERATION)?;
        blobstore.get(ctx, key).await
    }

    async fn put<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: String,
        value: BlobstoreBytes,
    ) -> Result<()> {
        let blobstore = self.access_blobstore(ctx, &key, config::PUT_OPERATION)?;
        blobstore.put(ctx, key, value).await
    }

    async fn is_present<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: &'a str,
    ) -> Result<BlobstoreIsPresent> {
        self.blobstore.is_present(ctx, key).await
    }

    async fn copy<'a>(
        &'a self,
        ctx: &'a CoreContext,
        old_key: &'a str,
        new_key: String,
    ) -> Result<()> {
        let blobstore = self.access_blobstore(ctx, old_key, config::GET_OPERATION)?;
        self.access_blobstore(ctx, &new_key, config::PUT_OPERATION)?;
        blobstore.copy(ctx, old_key, new_key).await
    }
}

#[async_trait]
impl<B: Blobstore> Blobstore for RedactedBlobstore<B> {
    async fn get<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: &'a str,
    ) -> Result<Option<BlobstoreGetData>> {
        self.inner.get(ctx, key).await
    }
    async fn put<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: String,
        value: BlobstoreBytes,
    ) -> Result<()> {
        self.inner.put(ctx, key, value).await
    }
    async fn is_present<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: &'a str,
    ) -> Result<BlobstoreIsPresent> {
        self.inner.is_present(ctx, key).await
    }
}

pub fn has_redaction_root_cause(e: &Error) -> bool {
    match e.root_cause().downcast_ref::<ErrorKind>() {
        Some(ErrorKind::Censored(_, _)) => true,
        None => false,
    }
}

#[cfg(test)]
mod test {
    use std::sync::Arc;

    use assert_matches::assert_matches;
    use borrowed::borrowed;
    use context::CoreContext;
    use fbinit::FacebookInit;
    use maplit::hashmap;
    use memblob::Memblob;
    use prefixblob::PrefixBlobstore;

    use super::*;

    #[fbinit::test]
    async fn test_redacted_key(fb: FacebookInit) {
        let unredacted_key = "foo";
        let redacted_key = "bar";
        let redacted_task = "bar task";

        let ctx = CoreContext::test_mock(fb);
        borrowed!(ctx);

        let inner = Memblob::default();
        let redacted_pairs = RedactedBlobs::FromSql(Arc::new(hashmap! {
            redacted_key.to_owned() => RedactedMetadata {
                task: redacted_task.to_owned(),
                log_only: false,
            },
        }));

        let blob = RedactedBlobstore::new(
            PrefixBlobstore::new(inner, "prefix"),
            RedactedBlobstoreConfig::new(
                Some(Arc::new(redacted_pairs)),
                MononokeScubaSampleBuilder::with_discard(),
            ),
        );

        //Test put with redacted key
        let res = blob
            .put(
                ctx,
                redacted_key.to_owned(),
                BlobstoreBytes::from_bytes("test bar"),
            )
            .await;

        assert_matches!(
            res.expect_err("the key should be redacted").downcast::<ErrorKind>(),
            Ok(ErrorKind::Censored(_, ref task)) if task == redacted_task
        );

        //Test key added to the blob
        let res = blob
            .put(
                ctx,
                unredacted_key.to_owned(),
                BlobstoreBytes::from_bytes("test foo"),
            )
            .await;
        assert!(res.is_ok(), "the key should be added successfully");

        // Test accessing a key which is redacted
        let res = blob.get(ctx, redacted_key).await;

        assert_matches!(
            res.expect_err("the key should be redacted").downcast::<ErrorKind>(),
            Ok(ErrorKind::Censored(_, ref task)) if *task == redacted_task
        );

        // Test accessing a key which exists and is accesible
        let res = blob.get(ctx, unredacted_key).await;
        assert!(res.is_ok(), "the key should be found and available");
    }

    #[fbinit::test]
    async fn test_log_only_redacted_key(fb: FacebookInit) -> Result<()> {
        let redacted_log_only_key = "bar";
        let redacted_task = "bar task";

        let ctx = CoreContext::test_mock(fb);
        borrowed!(ctx);

        let inner = Memblob::default();
        let redacted_pairs = RedactedBlobs::FromSql(Arc::new(hashmap! {
            redacted_log_only_key.to_owned() => RedactedMetadata {
                task: redacted_task.to_owned(),
                log_only: true,
            },
        }));

        let blob = RedactedBlobstore::new(
            PrefixBlobstore::new(inner, "prefix"),
            RedactedBlobstoreConfig::new(
                Some(Arc::new(redacted_pairs)),
                MononokeScubaSampleBuilder::with_discard(),
            ),
        );

        // Since this is a log-only mode it should succeed
        let val = BlobstoreBytes::from_bytes("test bar");
        blob.put(ctx, redacted_log_only_key.to_owned(), val.clone())
            .await?;

        let actual = blob.get(ctx, redacted_log_only_key).await?;
        assert_eq!(Some(val), actual.map(|val| val.into_bytes()));

        Ok(())
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::sync::Arc;

use anyhow::Result;
use async_trait::async_trait;
use blobstore::Blobstore;
use blobstore::BlobstoreGetData;
use blobstore::BlobstoreIsPresent;
use context::CoreContext;
use mononoke_types::BlobstoreBytes;
use prefixblob::PrefixBlobstore;

/// Blobstore where redaction sets are stored
#[facet::facet]
#[derive(Debug)]
pub struct RedactionConfigBlobstore(PrefixBlobstore<Arc<dyn Blobstore>>);

impl RedactionConfigBlobstore {
    pub fn new(blobstore: Arc<dyn Blobstore>) -> Self {
        Self(PrefixBlobstore::new(blobstore, "redactionconfig"))
    }
}

impl std::fmt::Display for RedactionConfigBlobstore {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(f, "RedactionConfigBlobstore<{}>", self.0)
    }
}

#[async_trait]
impl Blobstore for RedactionConfigBlobstore {
    async fn get<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: &'a str,
    ) -> Result<Option<BlobstoreGetData>> {
        self.0.get(ctx, key).await
    }
    async fn put<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: String,
        value: BlobstoreBytes,
    ) -> Result<()> {
        self.0.put(ctx, key, value).await
    }
    async fn is_present<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: &'a str,
    ) -> Result<BlobstoreIsPresent> {
        self.0.is_present(ctx, key).await
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::collections::HashMap;
use std::str::FromStr;
use std::sync::Arc;

use anyhow::Context;
use anyhow::Error;
use anyhow::Result;
use async_trait::async_trait;
use blobstore::Blobstore;
use blobstore::Loadable;
use cached_config::ConfigHandle;
use cached_config::ConfigStore;
use context::CoreContext;
use futures::stream;
use futures::stream::StreamExt;
use futures::stream::TryStreamExt;
use mononoke_types::typed_hash::RedactionKeyListId;
use mononoke_types::RedactionKeyList;
use mononoke_types::Timestamp;
use redaction_set::RedactionSets;
use reloader::Loader;
use reloader::Reloader;
use sql::Connection;
use sql_construct::SqlConstruct;
use sql_construct::SqlConstructFromMetadataDatabaseConfig;
use sql_ext::mononoke_queries;
use sql_ext::SqlConnections;

use crate::RedactionConfigBlobstore;

#[derive(Clone)]
pub struct SqlRedactedContentStore {
    read_connection: Connection,
    write_connection: Connection,
}

mononoke_queries! {

    write InsertRedactedBlobs(
        values: (content_key: String, task: String, add_timestamp: Timestamp, log_only: bool)
    ) {
        none,
        mysql(
            "INSERT INTO censored_contents(content_key, task, add_timestamp, log_only) VALUES {values}
            ON DUPLICATE KEY UPDATE task = VALUES(task), add_timestamp = VALUES(add_timestamp), log_ONLY = VALUES(log_only)
            "
        )
        sqlite(
            "REPLACE INTO censored_contents(content_key, task, add_timestamp, log_only) VALUES {values}"
        )
    }

    read GetAllRedactedBlobs() -> (String, String, Option<bool>) {
        "SELECT content_key, task, log_only
        FROM censored_contents"
    }

    write DeleteRedactedBlobs(>list content_keys: String) {
        none,
        "DELETE FROM censored_contents
         WHERE content_key IN {content_keys}"
    }
}

impl SqlConstruct for SqlRedactedContentStore {
    const LABEL: &'static str = "censored_contents";

    const CREATION_QUERY: &'static str = include_str!("../schemas/sqlite-redacted.sql");

    fn from_sql_connections(connections: SqlConnections) -> Self {
        Self {
            write_connection: connections.write_connection,
            read_connection: connections.read_connection,
        }
    }
}

impl SqlConstructFromMetadataDatabaseConfig for SqlRedactedContentStore {}

#[derive(Clone, Debug)]
pub struct RedactedMetadata {
    pub task: String,
    pub log_only: bool,
}

#[derive(Debug, Clone)]
pub enum RedactedBlobs {
    FromSql(Arc<HashMap<String, RedactedMetadata>>),
    FromConfigerator(Arc<ConfigeratorRedactedBlobs>),
}

impl RedactedBlobs {
    pub fn redacted(&self) -> Arc<HashMap<String, RedactedMetadata>> {
        match self {
            Self::FromSql(hm) => hm.clone(),
            Self::FromConfigerator(conf) => conf.get_map(),
        }
    }

    pub async fn from_configerator(
        store: &ConfigStore,
        config_path: &str,
        ctx: CoreContext,
        blobstore: Arc<RedactionConfigBlobstore>,
    ) -> Result<Self, Error> {
        let handle = store
            .get_config_handle(config_path.to_string())
            .with_context(|| format!("redaction sets not found at {}", config_path))?;
        Ok(Self::FromConfigerator(Arc::new(
            ConfigeratorRedactedBlobs::new(ctx, handle, blobstore).await?,
        )))
    }
}

#[derive(Debug)]
struct InnerConfig {
    #[allow(dead_code)]
    raw_config: Arc<RedactionSets>,
    map: Arc<HashMap<String, RedactedMetadata>>,
}

#[derive(Debug)]
pub struct ConfigeratorRedactedBlobs(Reloader<InnerConfig>);

struct InnerConfigLoader {
    ctx: CoreContext,
    handle: ConfigHandle<RedactionSets>,
    blobstore: Arc<RedactionConfigBlobstore>,
    last_config: Option<Arc<RedactionSets>>,
}
impl InnerConfigLoader {
    fn new(
        ctx: CoreContext,
        handle: ConfigHandle<RedactionSets>,
        blobstore: Arc<RedactionConfigBlobstore>,
    ) -> Self {
        Self {
            ctx,
            handle,
            blobstore,
            last_config: None,
        }
    }
}

#[async_trait]
impl Loader<InnerConfig> for InnerConfigLoader {
    async fn load(&mut self) -> Result<Option<InnerConfig>> {
        let new_config = self.handle.get();
        if match &self.last_config {
            Some(old) => !Arc::ptr_eq(old, &new_config) && *old != new_config,
            None => true,
        } {
            let res = Some(InnerConfig::new(new_config.clone(), &self.ctx, &self.blobstore).await?);
            self.last_config = Some(new_config);
            Ok(res)
        } else {
            Ok(None)
        }
    }
}

impl ConfigeratorRedactedBlobs {
    async fn new(
        ctx: CoreContext,
        handle: ConfigHandle<RedactionSets>,
        blobstore: Arc<RedactionConfigBlobstore>,
    ) -> Result<Self> {
        let loader = InnerConfigLoader::new(ctx.clone(), handle, blobstore);
        let reloader =
            Reloader::reload_periodically(ctx, || std::time::Duration::from_secs(60), loader)
                .await?;

        Ok(ConfigeratorRedactedBlobs(reloader))
    }

    fn get_map(&self) -> Arc<HashMap<String, RedactedMetadata>> {
        self.0.load().map.clone()
    }
}

impl InnerConfig {
    async fn new(
        config: Arc<RedactionSets>,
        ctx: &CoreContext,
        blobstore: &dyn Blobstore,
    ) -> Result<Self> {
        slog::debug!(ctx.logger(), "Reloading redacted config from configerator");
        let map: HashMap<String, RedactedMetadata> = stream::iter(
            config
                .all_redactions
                .iter()
                .map(|redaction| async move {
                    let keylist: RedactionKeyList = RedactionKeyListId::from_str(&redaction.id)
                        .with_context(|| format!("Invalid keylist id: {}", redaction.id))?
                        .load(ctx, &blobstore)
                        .await
                        .with_context(|| format!("Keylist with id {} not found", redaction.id))?;
                    let keys_with_metadata = keylist
                        .keys
                        .into_iter()
                        .map(|key| {
                            (
                                key,
                                RedactedMetadata {
                                    task: redaction.reason.clone(),
                                    log_only: !redaction.enforce,
                                },
                            )
                        })
                        .collect::<Vec<_>>();

                    Result::<_, Error>::Ok(keys_with_metadata)
                })
                // If we don't collect, it triggers a compile bug P421135476
                .collect::<Vec<_>>(),
        )
        .buffer_unordered(100)
        .try_collect::<Vec<Vec<_>>>()
        .await?
        .into_iter()
        .flatten()
        .collect();
        Ok(Self {
            raw_config: config,
            map: Arc::new(map),
        })
    }
}

impl SqlRedactedContentStore {
    pub async fn get_all_redacted_blobs(&self) -> Result<RedactedBlobs> {
        let redacted_blobs = GetAllRedactedBlobs::query(&self.read_connection).await?;
        Ok(RedactedBlobs::FromSql(Arc::new(
            redacted_blobs
                .into_iter()
                .map(|(key, task, log_only)| {
                    let redacted_metadata = RedactedMetadata {
                        task,
                        log_only: log_only.unwrap_or(false),
                    };
                    (key, redacted_metadata)
                })
                .collect(),
        )))
    }

    pub async fn insert_redacted_blobs(
        &self,
        content_keys: &[String],
        task: &String,
        add_timestamp: &Timestamp,
        log_only: bool,
    ) -> Result<()> {
        let log_only = &log_only;
        let redacted_inserts: Vec<_> = content_keys
            .iter()
            .map(move |key| (key, task, add_timestamp, log_only))
            .collect();

        InsertRedactedBlobs::query(&self.write_connection, &redacted_inserts[..])
            .await
            .map_err(Error::from)
            .map(|_| ())
    }

    pub async fn delete_redacted_blobs(&self, content_keys: &[String]) -> Result<()> {
        DeleteRedactedBlobs::query(&self.write_connection, content_keys)
            .await
            .map_err(Error::from)
            .map(|_| ())
    }
}

#[cfg(test)]
mod test {
    use super::*;

    #[fbinit::test]
    async fn test_redacted_store(_fb: fbinit::FacebookInit) {
        let key_a = "aaaaaaaaaaaaaaaaaaaa".to_string();
        let key_b = "bbbbbbbbbbbbbbbbbbbb".to_string();
        let key_c = "cccccccccccccccccccc".to_string();
        let key_d = "dddddddddddddddddddd".to_string();
        let task1 = "task1".to_string();
        let task2 = "task2".to_string();
        let redacted_keys1 = vec![key_a.clone(), key_b.clone()];
        let redacted_keys2 = vec![key_c.clone(), key_d.clone()];

        let store = SqlRedactedContentStore::with_sqlite_in_memory().unwrap();

        store
            .insert_redacted_blobs(&redacted_keys1, &task1, &Timestamp::now(), false)
            .await
            .expect("insert failed");
        store
            .insert_redacted_blobs(&redacted_keys2, &task2, &Timestamp::now(), true)
            .await
            .expect("insert failed");

        let all = store.get_all_redacted_blobs().await.expect("select failed");
        let res = all.redacted();
        assert_eq!(res.len(), 4);
        assert!(!res.get(&key_a).unwrap().log_only);
        assert!(!res.get(&key_b).unwrap().log_only);
        assert!(res.get(&key_c).unwrap().log_only);
        assert!(res.get(&key_d).unwrap().log_only);

        store
            .insert_redacted_blobs(&redacted_keys1, &task1, &Timestamp::now(), true)
            .await
            .expect("insert failed");
        let all = store.get_all_redacted_blobs().await.expect("select failed");
        let res = all.redacted();
        assert!(res.get(&key_a).unwrap().log_only);
        assert!(res.get(&key_b).unwrap().log_only);

        store
            .delete_redacted_blobs(&redacted_keys1)
            .await
            .expect("delete failed");
        let all = store.get_all_redacted_blobs().await.expect("select failed");
        let res = all.redacted();

        assert!(res.contains_key(&key_c));
        assert!(res.contains_key(&key_d));
        assert_eq!(res.len(), 2);
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::sync::Arc;

use anyhow::Result;
use async_trait::async_trait;
use blobstore::Blobstore;
use blobstore::BlobstoreGetData;
use blobstore::BlobstoreIsPresent;
use blobstore::BlobstorePutOps;
use blobstore::OverwriteStatus;
use blobstore::PutBehaviour;
use context::CoreContext;
use metaconfig_types::BlobstoreId;
use mononoke_types::BlobstoreBytes;

pub trait SamplingHandler: std::fmt::Debug + Send + Sync {
    fn sample_get(
        &self,
        ctx: &CoreContext,
        key: &str,
        value: Option<&BlobstoreGetData>,
    ) -> Result<()>;

    fn sample_put(&self, _ctx: &CoreContext, _key: &str, _value: &BlobstoreBytes) -> Result<()> {
        Ok(())
    }

    fn sample_is_present(
        &self,
        _ctx: &CoreContext,
        _key: &str,
        _value: &BlobstoreIsPresent,
    ) -> Result<()> {
        Ok(())
    }
}

/// A layer over an existing blobstore that allows sampling of blobs, e.g. for
/// corpus generation.
#[derive(Debug)]
pub struct SamplingBlobstore<T> {
    inner: T,
    handler: Arc<dyn SamplingHandler>,
}

impl<T: std::fmt::Display> std::fmt::Display for SamplingBlobstore<T> {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(f, "SamplingBlobstore<{}>", &self.inner)
    }
}

impl<T> SamplingBlobstore<T> {
    pub fn new(inner: T, handler: Arc<dyn SamplingHandler>) -> Self {
        Self { inner, handler }
    }
}

#[async_trait]
impl<T: Blobstore> Blobstore for SamplingBlobstore<T> {
    #[inline]
    async fn get<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: &'a str,
    ) -> Result<Option<BlobstoreGetData>> {
        let opt_blob = self.inner.get(ctx, key).await?;
        self.handler.sample_get(ctx, key, opt_blob.as_ref())?;
        Ok(opt_blob)
    }

    #[inline]
    async fn put<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: String,
        value: BlobstoreBytes,
    ) -> Result<()> {
        let sample_res = self.handler.sample_put(ctx, &key, &value);
        self.inner.put(ctx, key, value).await?;
        sample_res
    }

    #[inline]
    async fn is_present<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: &'a str,
    ) -> Result<BlobstoreIsPresent> {
        let result = self.inner.is_present(ctx, key).await?;
        self.handler.sample_is_present(ctx, key, &result)?;

        Ok(result)
    }
}

/// Used when you need the BlobstoreId (where there is one) in the sample
pub trait ComponentSamplingHandler: std::fmt::Debug + Send + Sync {
    fn sample_get(
        &self,
        ctx: &CoreContext,
        key: &str,
        value: Option<&BlobstoreGetData>,
        inner_id: Option<BlobstoreId>,
    ) -> Result<()>;

    fn sample_put(
        &self,
        _ctx: &CoreContext,
        _key: &str,
        _value: &BlobstoreBytes,
        _inner_id: Option<BlobstoreId>,
    ) -> Result<()> {
        Ok(())
    }

    fn sample_is_present(
        &self,
        _ctx: &CoreContext,
        _key: &str,
        _value: &BlobstoreIsPresent,
        _inner_id: Option<BlobstoreId>,
    ) -> Result<()> {
        Ok(())
    }
}

/// A lower level sampler that can provide BlobstoreId
#[derive(Debug)]
pub struct SamplingBlobstorePutOps<T> {
    inner: T,
    inner_id: Option<BlobstoreId>,
    handler: Arc<dyn ComponentSamplingHandler>,
}

impl<T: std::fmt::Display> std::fmt::Display for SamplingBlobstorePutOps<T> {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(
            f,
            "SamplingBlobstorePutOps<{}, {:?}>",
            &self.inner, self.inner_id
        )
    }
}

impl<T> SamplingBlobstorePutOps<T> {
    pub fn new(
        inner: T,
        inner_id: Option<BlobstoreId>,
        handler: Arc<dyn ComponentSamplingHandler>,
    ) -> Self {
        Self {
            inner,
            inner_id,
            handler,
        }
    }
}

#[async_trait]
impl<T: Blobstore + BlobstorePutOps> Blobstore for SamplingBlobstorePutOps<T> {
    #[inline]
    async fn get<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: &'a str,
    ) -> Result<Option<BlobstoreGetData>> {
        let opt_blob = self.inner.get(ctx, key).await?;
        self.handler
            .sample_get(ctx, key, opt_blob.as_ref(), self.inner_id)?;
        Ok(opt_blob)
    }

    #[inline]
    async fn put<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: String,
        value: BlobstoreBytes,
    ) -> Result<()> {
        let sample_res = self.handler.sample_put(ctx, &key, &value, self.inner_id);
        self.inner.put(ctx, key, value).await?;
        sample_res
    }

    #[inline]
    async fn is_present<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: &'a str,
    ) -> Result<BlobstoreIsPresent> {
        let result = self.inner.is_present(ctx, key).await?;
        self.handler
            .sample_is_present(ctx, key, &result, self.inner_id)?;

        Ok(result)
    }
}

#[async_trait]
impl<T: BlobstorePutOps> BlobstorePutOps for SamplingBlobstorePutOps<T> {
    async fn put_explicit<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: String,
        value: BlobstoreBytes,
        put_behaviour: PutBehaviour,
    ) -> Result<OverwriteStatus> {
        self.handler.sample_put(ctx, &key, &value, self.inner_id)?;
        self.inner
            .put_explicit(ctx, key, value, put_behaviour)
            .await
    }

    async fn put_with_status<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: String,
        value: BlobstoreBytes,
    ) -> Result<OverwriteStatus> {
        self.handler.sample_put(ctx, &key, &value, self.inner_id)?;
        self.inner.put_with_status(ctx, key, value).await
    }
}

#[cfg(test)]
mod test {
    use std::sync::atomic::AtomicBool;
    use std::sync::atomic::Ordering;

    use borrowed::borrowed;
    use context::SamplingKey;
    use fbinit::FacebookInit;
    use memblob::Memblob;

    use super::*;

    #[derive(Debug)]
    struct TestSamplingHandler {
        sampled: AtomicBool,
        looking_for: SamplingKey,
    }
    impl TestSamplingHandler {
        fn check_sample(&self, ctx: &CoreContext) -> Result<()> {
            if let Some(sampling_key) = ctx.sampling_key() {
                if sampling_key == &self.looking_for {
                    self.sampled.store(true, Ordering::Relaxed);
                }
            }
            Ok(())
        }
    }

    impl SamplingHandler for TestSamplingHandler {
        fn sample_get(
            &self,
            ctx: &CoreContext,
            _key: &str,
            _value: Option<&BlobstoreGetData>,
        ) -> Result<()> {
            self.check_sample(ctx)
        }
        fn sample_put(&self, ctx: &CoreContext, _key: &str, _value: &BlobstoreBytes) -> Result<()> {
            self.check_sample(ctx)
        }
        fn sample_is_present(
            &self,
            ctx: &CoreContext,
            _key: &str,
            _value: &BlobstoreIsPresent,
        ) -> Result<()> {
            self.check_sample(ctx)
        }
    }

    #[fbinit::test]
    async fn test_sample_called(fb: FacebookInit) {
        let ctx = CoreContext::test_mock(fb);
        borrowed!(ctx);
        let base = Memblob::default();
        let sample_this = SamplingKey::new();
        let handler = Arc::new(TestSamplingHandler {
            sampled: AtomicBool::new(false),
            looking_for: sample_this,
        });
        let wrapper =
            SamplingBlobstore::new(base.clone(), handler.clone() as Arc<dyn SamplingHandler>);
        let key = "foobar";

        let r = wrapper
            .put(
                ctx,
                key.to_owned(),
                BlobstoreBytes::from_bytes("test foobar"),
            )
            .await;
        assert!(r.is_ok());
        let was_sampled = handler.sampled.load(Ordering::Relaxed);
        assert!(!was_sampled);
        let ctx = ctx.clone_and_sample(sample_this);
        borrowed!(ctx);
        let base_present = base
            .is_present(ctx, key)
            .await
            .unwrap()
            .assume_not_found_if_unsure();
        assert!(base_present);
        let was_sampled = handler.sampled.load(Ordering::Relaxed);
        assert!(!was_sampled);
        let wrapper_present = wrapper
            .is_present(ctx, key)
            .await
            .unwrap()
            .assume_not_found_if_unsure();
        assert!(wrapper_present);
        let was_sampled = handler.sampled.load(Ordering::Relaxed);
        assert!(was_sampled);
    }

    impl ComponentSamplingHandler for TestSamplingHandler {
        fn sample_get(
            &self,
            ctx: &CoreContext,
            _key: &str,
            _value: Option<&BlobstoreGetData>,
            _inner_id: Option<BlobstoreId>,
        ) -> Result<()> {
            self.check_sample(ctx)
        }
        fn sample_put(
            &self,
            ctx: &CoreContext,
            _key: &str,
            _value: &BlobstoreBytes,
            _inner_id: Option<BlobstoreId>,
        ) -> Result<()> {
            self.check_sample(ctx)
        }
        fn sample_is_present(
            &self,
            ctx: &CoreContext,
            _key: &str,
            _value: &BlobstoreIsPresent,
            _inner_id: Option<BlobstoreId>,
        ) -> Result<()> {
            self.check_sample(ctx)
        }
    }

    #[fbinit::test]
    async fn test_component_sample_called(fb: FacebookInit) {
        let ctx = CoreContext::test_mock(fb);
        borrowed!(ctx);
        let base = Memblob::default();
        let sample_this = SamplingKey::new();
        let handler = Arc::new(TestSamplingHandler {
            sampled: AtomicBool::new(false),
            looking_for: sample_this,
        });
        let wrapper = SamplingBlobstorePutOps::new(
            base.clone(),
            None,
            handler.clone() as Arc<dyn ComponentSamplingHandler>,
        );
        let key = "foobar";

        let r = wrapper
            .put(
                ctx,
                key.to_owned(),
                BlobstoreBytes::from_bytes("test foobar"),
            )
            .await;
        assert!(r.is_ok());
        let was_sampled = handler.sampled.load(Ordering::Relaxed);
        assert!(!was_sampled);
        let ctx = ctx.clone_and_sample(sample_this);
        borrowed!(ctx);
        let base_present = base
            .is_present(ctx, key)
            .await
            .unwrap()
            .assume_not_found_if_unsure();
        assert!(base_present);
        let was_sampled = handler.sampled.load(Ordering::Relaxed);
        assert!(!was_sampled);
        let wrapper_present = wrapper
            .is_present(ctx, key)
            .await
            .unwrap()
            .assume_not_found_if_unsure();
        assert!(wrapper_present);
        let was_sampled = handler.sampled.load(Ordering::Relaxed);
        assert!(was_sampled);
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::cmp::max;
use std::cmp::min;
use std::num::NonZeroUsize;
use std::time::Duration;
use std::time::Instant;

use futures::stream::StreamExt;
use rand::thread_rng;
use rand::Rng;
use stats::prelude::*;
use tokio::sync::watch;

// This can be tweaked later.
pub(crate) const MAX_LAG: Duration = Duration::from_secs(5);

define_stats! {
    prefix = "mononoke.sqlblob.lag_delay";
    total_delay_ms: dynamic_timeseries("{}.total_delay_ms", (entity: String); Rate, Sum),
    raw_lag_ms: dynamic_timeseries("{}.raw_lag_ms", (entity: String); Rate, Sum),
}

#[derive(Clone)]
pub struct BlobDelay {
    lag_receivers: Vec<watch::Receiver<Duration>>,
    entity: Option<String>,
}

// Adds a small amount of random delay to desynchronise when waiting
async fn jitter_delay(raw_lag: Duration) {
    // Delay should not be greater than 1 second
    let delay = min(Duration::from_millis(1000), raw_lag);
    // Delay should also not be less than or equal to 0 seconds since that
    // will make the range empty and lead to panic
    let delay = max(Duration::from_millis(50), delay);
    let delay = thread_rng().gen_range(Duration::from_millis(0)..delay);
    tokio::time::sleep(delay).await;
}

impl BlobDelay {
    pub fn dummy(shard_count: NonZeroUsize) -> Self {
        let lag_receivers = vec![
            {
                let (_, ch) = watch::channel(Duration::new(0, 0));
                ch
            };
            shard_count.into()
        ];
        Self {
            lag_receivers,
            entity: None,
        }
    }

    #[cfg(fbcode_build)]
    pub fn from_channels(lag_receivers: Vec<watch::Receiver<Duration>>, name: String) -> Self {
        let entity = Some(name);
        Self {
            lag_receivers,
            entity,
        }
    }

    pub async fn delay(&self, shard_id: usize) {
        let mut lag_receiver =
            tokio_stream::wrappers::WatchStream::new(self.lag_receivers[shard_id].clone());
        let start_time = Instant::now();

        while let Some(raw_lag) = lag_receiver.next().await {
            if raw_lag < MAX_LAG {
                if start_time.elapsed() > Duration::from_secs(1) {
                    // No jittering for short delays, but jitter us about a bit if we've seen
                    // lag and waited for it to die down, so that next request is random
                    jitter_delay(raw_lag).await;
                }
                break;
            }
            if let Some(entity) = &self.entity {
                let raw_lag_ms = raw_lag.as_millis().try_into();
                if let Ok(raw_lag_ms) = raw_lag_ms {
                    STATS::raw_lag_ms.add_value(raw_lag_ms, (entity.clone(),))
                }
            }
        }
        if let Some(entity) = &self.entity {
            let total_delay_ms = start_time.elapsed().as_millis().try_into();
            if let Ok(total_delay_ms) = total_delay_ms {
                STATS::total_delay_ms.add_value(total_delay_ms, (entity.clone(),));
            }
        }
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

#![cfg_attr(not(fbcode_build), allow(unused_crate_dependencies))]

mod delay;
#[cfg(fbcode_build)]
mod facebook;
#[cfg(not(fbcode_build))]
mod myadmin_delay_dummy;
mod store;
#[cfg(test)]
mod tests;

use std::collections::HashMap;
use std::fmt;
use std::future::Future;
use std::num::NonZeroUsize;
use std::path::PathBuf;
use std::sync::Arc;
use std::time::Duration;
use std::time::SystemTime;

use anyhow::bail;
use anyhow::format_err;
use anyhow::Error;
use anyhow::Result;
use async_trait::async_trait;
use blobstore::Blobstore;
use blobstore::BlobstoreGetData;
use blobstore::BlobstoreIsPresent;
use blobstore::BlobstoreMetadata;
use blobstore::BlobstorePutOps;
use blobstore::BlobstoreUnlinkOps;
use blobstore::CountedBlobstore;
use blobstore::OverwriteStatus;
use blobstore::PutBehaviour;
use bytes::Bytes;
use bytes::BytesMut;
use cached_config::ConfigHandle;
use cached_config::ConfigStore;
use cached_config::ModificationTime;
use cached_config::TestSource;
use context::CoreContext;
use fbinit::FacebookInit;
use futures::stream::FuturesOrdered;
use futures::stream::FuturesUnordered;
use futures::stream::Stream;
use futures::stream::TryStreamExt;
use futures::TryFutureExt;
use mononoke_types::hash::Context as HashContext;
use mononoke_types::BlobstoreBytes;
use nonzero_ext::nonzero;
use sql::rusqlite::Connection as SqliteConnection;
use sql::Connection;
use sql_ext::facebook::create_mysql_connections_sharded;
use sql_ext::facebook::create_mysql_connections_unsharded;
use sql_ext::facebook::MysqlOptions;
use sql_ext::open_sqlite_in_memory;
use sql_ext::open_sqlite_path;
use sql_ext::SqlConnections;
use sql_ext::SqlShardedConnections;
use tokio::task::spawn_blocking;
use vec1::Vec1;
use xdb_gc_structs::XdbGc;

use crate::delay::BlobDelay;
#[cfg(fbcode_build)]
use crate::facebook::myadmin_delay;
#[cfg(not(fbcode_build))]
use crate::myadmin_delay_dummy as myadmin_delay;
use crate::store::ChunkGenerationState;
use crate::store::ChunkSqlStore;
use crate::store::ChunkingMethod;
use crate::store::DataSqlStore;

// Leaving some space for metadata
const MAX_KEY_SIZE: usize = 200;
// MySQL wants multiple chunks, each around 1 MiB, as a tradeoff between query latency and replication lag
const CHUNK_SIZE: usize = 1024 * 1024;
const SQLITE_SHARD_NUM: NonZeroUsize = nonzero!(2_usize);
const SINGLE_SHARD_NUM: NonZeroUsize = nonzero!(1_usize);
const GC_GENERATION_PATH: &str = "scm/mononoke/xdb_gc/default";

const SQLBLOB_LABEL: &str = "blobstore";

// Test setup data
const UPDATE_FREQUENCY: Duration = Duration::from_millis(1);
const INITIAL_VERSION: u64 = 0;

const COUNTED_ID: &str = "sqlblob";
pub type CountedSqlblob = CountedBlobstore<Sqlblob>;

pub struct Sqlblob {
    data_store: Arc<DataSqlStore>,
    chunk_store: Arc<ChunkSqlStore>,
    put_behaviour: PutBehaviour,
    allow_inline_put: bool,
    ctime_inline_grace: i64,
}

impl std::fmt::Display for Sqlblob {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "Sqlblob")
    }
}

fn get_gc_config_handle(config_store: &ConfigStore) -> Result<ConfigHandle<XdbGc>> {
    config_store.get_config_handle(GC_GENERATION_PATH.to_string())
}

const DEFAULT_ALLOW_INLINE_PUT: bool = true;

// One day
const DEFAULT_CTIME_INLINE_GRACE: i64 = 86400;

// base64 encoding for inline hash has an overhead
pub const MAX_INLINE_LEN: u64 = 255 * 3 / 4;

fn encode_small_value(raw: &[u8]) -> String {
    base64::encode_config(raw, base64::STANDARD_NO_PAD)
}

impl Sqlblob {
    pub async fn with_mysql(
        fb: FacebookInit,
        shardmap: String,
        shard_num: NonZeroUsize,
        mysql_options: MysqlOptions,
        readonly: bool,
        put_behaviour: PutBehaviour,
        config_store: &ConfigStore,
    ) -> Result<CountedSqlblob, Error> {
        let delay = if readonly {
            BlobDelay::dummy(shard_num)
        } else {
            myadmin_delay::sharded(fb, shardmap.clone(), shard_num)?
        };
        let config_handle = get_gc_config_handle(config_store)?;
        let shard_count = shard_num.clone().get();

        let SqlShardedConnections {
            read_connections,
            read_master_connections,
            write_connections,
        } = spawn_blocking({
            let shardmap = shardmap.clone();
            move || {
                create_mysql_connections_sharded(
                    fb,
                    mysql_options,
                    SQLBLOB_LABEL.into(),
                    shardmap,
                    0..shard_count,
                    readonly,
                )
            }
        })
        .await??;

        let write_connections = Arc::new(write_connections);
        let read_connections = Arc::new(read_connections);
        let read_master_connections = Arc::new(read_master_connections);
        Ok(Self::counted(
            Self {
                data_store: Arc::new(DataSqlStore::new(
                    shard_num,
                    write_connections.clone(),
                    read_connections.clone(),
                    read_master_connections.clone(),
                    delay.clone(),
                )),
                chunk_store: Arc::new(ChunkSqlStore::new(
                    shard_num,
                    write_connections,
                    read_connections,
                    read_master_connections,
                    delay,
                    config_handle,
                )),
                put_behaviour,
                allow_inline_put: DEFAULT_ALLOW_INLINE_PUT,
                ctime_inline_grace: DEFAULT_CTIME_INLINE_GRACE,
            },
            shardmap,
        ))
    }

    pub async fn with_mysql_unsharded(
        fb: FacebookInit,
        db_address: String,
        mysql_options: MysqlOptions,
        readonly: bool,
        put_behaviour: PutBehaviour,
        config_store: &ConfigStore,
    ) -> Result<CountedSqlblob, Error> {
        let delay = if readonly {
            BlobDelay::dummy(SINGLE_SHARD_NUM)
        } else {
            myadmin_delay::single(fb, db_address.clone())?
        };
        Self::with_connection_factory(
            delay,
            db_address.clone(),
            SINGLE_SHARD_NUM,
            put_behaviour,
            move |_shard_id| {
                let res = create_mysql_connections_unsharded(
                    fb,
                    mysql_options.clone(),
                    SQLBLOB_LABEL.into(),
                    db_address.clone(),
                    readonly,
                );
                async { res }
            },
            config_store,
            DEFAULT_ALLOW_INLINE_PUT,
            DEFAULT_CTIME_INLINE_GRACE,
        )
        .await
    }

    async fn with_connection_factory<CF, SF>(
        delay: BlobDelay,
        label: String,
        shard_num: NonZeroUsize,
        put_behaviour: PutBehaviour,
        connection_factory: CF,
        config_store: &ConfigStore,
        allow_inline_put: bool,
        ctime_inline_grace: i64,
    ) -> Result<CountedSqlblob, Error>
    where
        CF: Fn(usize) -> SF,
        SF: Future<Output = Result<SqlConnections, Error>> + Sized,
    {
        let shard_count = shard_num.get();

        let config_handle = get_gc_config_handle(config_store)?;

        let futs: FuturesOrdered<_> = (0..shard_count).map(connection_factory).collect();

        let shard_connections = futs.try_collect::<Vec<_>>().await?;
        let mut write_connections = Vec::with_capacity(shard_count);
        let mut read_connections = Vec::with_capacity(shard_count);
        let mut read_master_connections = Vec::with_capacity(shard_count);

        for connections in shard_connections {
            write_connections.push(connections.write_connection);
            read_connections.push(connections.read_connection);
            read_master_connections.push(connections.read_master_connection);
        }

        let write_connections: Arc<Vec1<Connection>> = Arc::new(write_connections.try_into()?);
        let read_connections: Arc<Vec1<Connection>> = Arc::new(read_connections.try_into()?);
        let read_master_connections: Arc<Vec1<Connection>> =
            Arc::new(read_master_connections.try_into()?);

        Ok(Self::counted(
            Self {
                data_store: Arc::new(DataSqlStore::new(
                    shard_num,
                    write_connections.clone(),
                    read_connections.clone(),
                    read_master_connections.clone(),
                    delay.clone(),
                )),
                chunk_store: Arc::new(ChunkSqlStore::new(
                    shard_num,
                    write_connections,
                    read_connections,
                    read_master_connections,
                    delay,
                    config_handle,
                )),
                put_behaviour,
                allow_inline_put,
                ctime_inline_grace,
            },
            label,
        ))
    }

    pub fn with_sqlite_in_memory(
        put_behaviour: PutBehaviour,
        config_store: &ConfigStore,
        allow_inline_put: bool,
        ctime_inline_grace: i64,
    ) -> Result<CountedSqlblob> {
        Self::with_sqlite(
            put_behaviour,
            |_| {
                let con = open_sqlite_in_memory()?;
                con.execute_batch(Self::CREATION_QUERY)?;
                Ok(con)
            },
            config_store,
            allow_inline_put,
            ctime_inline_grace,
        )
    }

    pub fn with_sqlite_path<P: Into<PathBuf>>(
        path: P,
        readonly_storage: bool,
        put_behaviour: PutBehaviour,
        config_store: &ConfigStore,
    ) -> Result<CountedSqlblob> {
        let pathbuf = path.into();
        Self::with_sqlite(
            put_behaviour,
            move |shard_id| {
                let con = open_sqlite_path(
                    pathbuf.join(format!("shard_{}.sqlite", shard_id)),
                    readonly_storage,
                )?;
                con.execute_batch(Self::CREATION_QUERY)?;
                Ok(con)
            },
            config_store,
            DEFAULT_ALLOW_INLINE_PUT,
            DEFAULT_CTIME_INLINE_GRACE,
        )
    }

    fn with_sqlite<F>(
        put_behaviour: PutBehaviour,
        mut constructor: F,
        config_store: &ConfigStore,
        allow_inline_put: bool,
        ctime_inline_grace: i64,
    ) -> Result<CountedSqlblob>
    where
        F: FnMut(usize) -> Result<SqliteConnection>,
    {
        let mut cons = Vec::with_capacity(SQLITE_SHARD_NUM.get());

        for i in 0..SQLITE_SHARD_NUM.get() {
            cons.push(Connection::with_sqlite(constructor(i)?));
        }

        let cons: Arc<Vec1<Connection>> = Arc::new(cons.try_into()?);

        // SQLite is predominately intended for tests, and has less concurrency
        // issues relating to GC, so cope with missing configerator
        let config_handle = get_gc_config_handle(config_store)
            .or_else(|_| get_gc_config_handle(&(get_test_config_store().1)))?;

        Ok(Self::counted(
            Self {
                data_store: Arc::new(DataSqlStore::new(
                    SQLITE_SHARD_NUM,
                    cons.clone(),
                    cons.clone(),
                    cons.clone(),
                    BlobDelay::dummy(SQLITE_SHARD_NUM),
                )),
                chunk_store: Arc::new(ChunkSqlStore::new(
                    SQLITE_SHARD_NUM,
                    cons.clone(),
                    cons.clone(),
                    cons,
                    BlobDelay::dummy(SQLITE_SHARD_NUM),
                    config_handle,
                )),
                put_behaviour,
                allow_inline_put,
                ctime_inline_grace,
            },
            "sqlite".into(),
        ))
    }

    const CREATION_QUERY: &'static str = include_str!("../schema/sqlite-sqlblob.sql");

    fn counted(self, label: String) -> CountedBlobstore<Self> {
        CountedBlobstore::new(format!("{}.{}", COUNTED_ID, label), self)
    }

    #[cfg(test)]
    pub(crate) fn get_data_store(&self) -> &DataSqlStore {
        &self.data_store
    }

    pub fn get_keys_from_shard(&self, shard_num: usize) -> impl Stream<Item = Result<String>> {
        self.data_store.get_keys_from_shard(shard_num)
    }

    /// Returns a HashMap from generation->(size, chunk_count)
    pub async fn get_chunk_sizes_by_generation(
        &self,
        shard_num: usize,
    ) -> Result<HashMap<Option<u64>, (u64, u64)>> {
        self.chunk_store
            .get_chunk_sizes_by_generation(shard_num)
            .await
    }

    pub async fn set_initial_generation(&self, shard_num: usize) -> Result<()> {
        self.chunk_store.set_initial_generation(shard_num).await
    }

    #[cfg(test)]
    pub async fn get_chunk_generations(&self, key: &str) -> Result<Vec<Option<u64>>> {
        let chunked = self.data_store.get(key).await?;
        if let Some(chunked) = chunked {
            let fetch_chunk_generations: FuturesOrdered<_> = (0..chunked.count)
                .map(|chunk_num| {
                    self.chunk_store
                        .get_generation(&chunked.id, chunk_num, chunked.chunking_method)
                })
                .collect();
            fetch_chunk_generations.try_collect().await
        } else {
            bail!("key does not exist");
        }
    }

    pub fn get_mark_generation(&self) -> u64 {
        self.chunk_store.get_mark_generation()
    }

    /// Mark the generation for a key
    /// If its value was small enough to inline, then also inline it if requested
    pub async fn set_generation(
        &self,
        key: &str,
        inline_small_values: bool,
        // Take the mark generation as param, so that marking for an entire run is consistent
        mark_generation: u64,
    ) -> Result<()> {
        let chunked = self.data_store.get(key).await?;
        if let Some(chunked) = chunked {
            let set_chunk_generations: FuturesUnordered<_> = (0..chunked.count)
                .map(|chunk_num| {
                    self.chunk_store
                        .set_generation(
                            &chunked.id,
                            chunk_num,
                            chunked.chunking_method,
                            mark_generation,
                        )
                        .map_ok(|value_len| {
                            if let Some(value_len) = value_len {
                                // Should not be chunked at all, request the key is inlined
                                // We do this after marking rather than short circuiting the above
                                // so that inlining following this call fails the key is still live
                                value_len <= MAX_INLINE_LEN
                            } else {
                                false
                            }
                        })
                })
                .collect();
            let can_inline: Vec<bool> = set_chunk_generations.try_collect().await?;
            if inline_small_values && can_inline.len() == 1 && can_inline[0] {
                // Value was small, so lets inline it
                let small_value = self.get_impl(key).await?;
                if let Some(small_value) = small_value {
                    // Double check length incase it changed since setting generation
                    let value_len: u64 = small_value.as_bytes().len().try_into()?;
                    if value_len <= MAX_INLINE_LEN {
                        if let Some(old_ctime) = small_value.as_meta().ctime() {
                            let ctime = {
                                match SystemTime::now().duration_since(SystemTime::UNIX_EPOCH) {
                                    Ok(offset) => offset.as_secs().try_into(),
                                    Err(negative) => {
                                        negative.duration().as_secs().try_into().map(|v: i64| -v)
                                    }
                                }
                            }?;
                            // Optimisitic update to convert to the inline form. This only updates for actually old ctimes
                            // (one day or more), so that we don't attempt to inline any new data from e.g. a packer write
                            if ctime - old_ctime >= self.ctime_inline_grace {
                                let small_value = encode_small_value(&small_value.into_raw_bytes());
                                self.data_store
                                    .update_optimistic(
                                        key,
                                        ctime,
                                        &small_value,
                                        0,
                                        ChunkingMethod::InlineBase64,
                                        old_ctime,
                                    )
                                    .await?;
                            }
                        }
                    }
                }
            }
            Ok(())
        } else {
            bail!("key does not exist");
        }
    }

    async fn get_impl<'a>(&'a self, key: &'a str) -> Result<Option<BlobstoreGetData>> {
        let chunked = self.data_store.get(key).await?;
        if let Some(chunked) = chunked {
            let blob = match chunked.chunking_method {
                ChunkingMethod::InlineBase64 => {
                    let decoded = base64::decode_config(&chunked.id, base64::STANDARD_NO_PAD)?;
                    Bytes::copy_from_slice(decoded.as_ref())
                }
                ChunkingMethod::ByContentHashBlake2 => {
                    let chunks = (0..chunked.count)
                        .map(|chunk_num| {
                            self.chunk_store
                                .get(&chunked.id, chunk_num, chunked.chunking_method)
                        })
                        .collect::<FuturesOrdered<_>>()
                        .try_collect::<Vec<_>>()
                        .await?;

                    let size = chunks.iter().map(|chunk| chunk.len()).sum();
                    let mut blob = BytesMut::with_capacity(size);
                    for chunk in chunks {
                        blob.extend_from_slice(&chunk);
                    }
                    blob.freeze()
                }
            };

            let meta = BlobstoreMetadata::new(Some(chunked.ctime), None);
            Ok(Some(BlobstoreGetData::new(
                meta,
                BlobstoreBytes::from_bytes(blob),
            )))
        } else {
            Ok(None)
        }
    }
}

impl fmt::Debug for Sqlblob {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("Sqlblob").finish()
    }
}

#[async_trait]
impl Blobstore for Sqlblob {
    async fn get<'a>(
        &'a self,
        _ctx: &'a CoreContext,
        key: &'a str,
    ) -> Result<Option<BlobstoreGetData>> {
        self.get_impl(key).await
    }

    async fn is_present<'a>(
        &'a self,
        _ctx: &'a CoreContext,
        key: &'a str,
    ) -> Result<BlobstoreIsPresent> {
        let present = self.data_store.is_present(key).await?;
        Ok(if present {
            BlobstoreIsPresent::Present
        } else {
            BlobstoreIsPresent::Absent
        })
    }

    async fn put<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: String,
        value: BlobstoreBytes,
    ) -> Result<()> {
        BlobstorePutOps::put_with_status(self, ctx, key, value).await?;
        Ok(())
    }

    async fn copy<'a>(
        &'a self,
        _ctx: &'a CoreContext,
        old_key: &'a str,
        new_key: String,
    ) -> Result<()> {
        let existing_data = self
            .data_store
            .get(old_key)
            .await?
            .ok_or_else(|| format_err!("Key {} does not exist in the blobstore", old_key))?;
        self.data_store
            .put(
                &new_key,
                existing_data.ctime,
                &existing_data.id,
                existing_data.count,
                existing_data.chunking_method,
            )
            .await
    }
}

#[async_trait]
impl BlobstorePutOps for Sqlblob {
    async fn put_explicit<'a>(
        &'a self,
        _ctx: &'a CoreContext,
        key: String,
        value: BlobstoreBytes,
        put_behaviour: PutBehaviour,
    ) -> Result<OverwriteStatus> {
        if key.as_bytes().len() > MAX_KEY_SIZE {
            return Err(format_err!(
                "Key {} exceeded max key size {}",
                key,
                MAX_KEY_SIZE
            ));
        }

        if put_behaviour == PutBehaviour::IfAbsent && self.data_store.is_present(&key).await? {
            // Can short circuit here as key already exists, and is keeping its chunks live
            return Ok(OverwriteStatus::Prevented);
        }

        let value_len: u64 = value.len().try_into()?;

        let chunking_method = if self.allow_inline_put && value_len <= MAX_INLINE_LEN {
            ChunkingMethod::InlineBase64
        } else {
            ChunkingMethod::ByContentHashBlake2
        };

        let put_fut = async {
            let ctime = {
                match SystemTime::now().duration_since(SystemTime::UNIX_EPOCH) {
                    Ok(offset) => offset.as_secs().try_into(),
                    Err(negative) => negative.duration().as_secs().try_into().map(|v: i64| -v),
                }
            }?;
            let (chunk_key, chunk_count, chunk_gen_insert_shard_id) = match chunking_method {
                ChunkingMethod::ByContentHashBlake2 => {
                    let chunk_key = {
                        let mut hash_context = HashContext::new(b"sqlblob");
                        hash_context.update(value.as_bytes());
                        hash_context.finish().to_hex().to_string()
                    };
                    let chunks = value.as_bytes().chunks(CHUNK_SIZE);
                    let chunk_count = chunks.len().try_into()?;
                    let mut updated_gen = false;
                    let mut chunk_gen_insert_shard_id = None;
                    for (chunk_num, chunk_value) in chunks.enumerate() {
                        let chunk_gen_state = self
                            .chunk_store
                            .put(
                                chunk_key.as_str(),
                                chunk_num.try_into()?,
                                chunking_method,
                                chunk_value,
                                value_len,
                            )
                            .await?;
                        match chunk_gen_state {
                            Some(ChunkGenerationState::Updated) => {
                                updated_gen = true;
                                // Seen an update, no need to insert later
                                chunk_gen_insert_shard_id = None;
                            }
                            Some(ChunkGenerationState::NeedsInsertToShard(shard_id)) => {
                                if !updated_gen {
                                    chunk_gen_insert_shard_id = Some(shard_id);
                                }
                            }
                            None => {}
                        };
                    }
                    (chunk_key, chunk_count, chunk_gen_insert_shard_id)
                }
                ChunkingMethod::InlineBase64 => {
                    (encode_small_value(value.as_bytes().as_ref()), 0, None)
                }
            };

            self.data_store
                .put(
                    &key,
                    ctime,
                    chunk_key.as_str(),
                    chunk_count,
                    chunking_method,
                )
                .await?;

            // Called after data_store.put to maintain invariant that chunk and data put complete
            // successfully before a generation is inserted (aka no dangling generations)
            if let Some(shard_id) = chunk_gen_insert_shard_id {
                self.chunk_store
                    .put_chunk_generation(&chunk_key, shard_id, value_len)
                    .await?
            }

            Ok(OverwriteStatus::NotChecked)
        };

        match put_behaviour {
            PutBehaviour::Overwrite => put_fut.await,
            PutBehaviour::IfAbsent | PutBehaviour::OverwriteAndLog => {
                match self.data_store.get(&key).await? {
                    None => {
                        put_fut.await?;
                        Ok(OverwriteStatus::New)
                    }
                    Some(chunked) => {
                        if put_behaviour.should_overwrite() {
                            put_fut.await?;
                            Ok(OverwriteStatus::Overwrote)
                        } else {
                            let chunk_count = chunked.count;
                            for chunk_num in 0..chunk_count {
                                self.chunk_store
                                    .update_generation(
                                        &chunked.id,
                                        chunk_num,
                                        chunked.chunking_method,
                                        value_len,
                                    )
                                    .await?;
                            }
                            Ok(OverwriteStatus::Prevented)
                        }
                    }
                }
            }
        }
    }

    async fn put_with_status<'a>(
        &'a self,
        ctx: &'a CoreContext,
        key: String,
        value: BlobstoreBytes,
    ) -> Result<OverwriteStatus> {
        self.put_explicit(ctx, key, value, self.put_behaviour).await
    }
}

#[async_trait]
impl BlobstoreUnlinkOps for Sqlblob {
    async fn unlink<'a>(&'a self, _ctx: &'a CoreContext, key: &'a str) -> Result<()> {
        if !self.data_store.is_present(key).await? {
            bail!(
                "Sqlblob::unlink: key {} does not exist in the blobstore",
                key
            )
        };
        self.data_store.unlink(key).await
    }
}

pub fn set_test_generations(
    source: &TestSource,
    put_generation: i64,
    mark_generation: i64,
    delete_generation: i64,
    mod_time: u64,
) -> u64 {
    source.insert_config(
        GC_GENERATION_PATH,
        &serde_json::to_string(&XdbGc {
            put_generation,
            mark_generation,
            delete_generation,
        })
        .expect("Invalid input config somehow"),
        ModificationTime::UnixTimestamp(mod_time),
    );
    source.insert_to_refresh(GC_GENERATION_PATH.to_string());
    mark_generation as u64
}

pub fn get_test_config_store() -> (Arc<TestSource>, ConfigStore) {
    let test_source = Arc::new(TestSource::new());
    set_test_generations(test_source.as_ref(), 2, 1, 0, INITIAL_VERSION);
    (
        test_source.clone(),
        ConfigStore::new(test_source, UPDATE_FREQUENCY, None),
    )
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::num::NonZeroUsize;

use anyhow::Result;
use fbinit::FacebookInit;

use crate::delay::BlobDelay;

pub fn sharded(
    _fb: FacebookInit,
    _shardmap: String,
    shard_count: NonZeroUsize,
) -> Result<BlobDelay> {
    Ok(BlobDelay::dummy(shard_count))
}

pub fn single(_fb: FacebookInit, _shard: String) -> Result<BlobDelay> {
    Ok(BlobDelay::dummy(
        NonZeroUsize::new(1).expect("1 should never be zero"),
    ))
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::collections::HashMap;
use std::hash::Hasher;
use std::num::NonZeroUsize;
use std::sync::Arc;

use anyhow::bail;
use anyhow::format_err;
use anyhow::Error;
use bytes::BytesMut;
use cached_config::ConfigHandle;
use futures::future::TryFutureExt;
use futures::stream;
use futures::stream::Stream;
use sql::Connection;
use sql_ext::mononoke_queries;
use twox_hash::XxHash32;
use vec1::Vec1;
use xdb_gc_structs::XdbGc;

use crate::delay::BlobDelay;

mod types {
    use sql::mysql;
    use sql::mysql_async::prelude::ConvIr;
    use sql::mysql_async::prelude::FromValue;
    use sql::mysql_async::FromValueError;
    use sql::mysql_async::Value;

    type FromValueResult<T> = Result<T, FromValueError>;

    #[derive(Clone, Copy, Debug, PartialEq, mysql::OptTryFromRowField)]
    pub enum ChunkingMethod {
        ByContentHashBlake2,
        InlineBase64,
    }

    impl From<ChunkingMethod> for Value {
        fn from(dtype: ChunkingMethod) -> Self {
            match dtype {
                // When you add here, please add the reverse transform
                // to impl ConvIr<ChunkingMethod> below
                ChunkingMethod::ByContentHashBlake2 => Value::UInt(1),
                ChunkingMethod::InlineBase64 => Value::UInt(2),
            }
        }
    }

    impl ConvIr<ChunkingMethod> for ChunkingMethod {
        fn new(v: Value) -> FromValueResult<Self> {
            match v {
                // Note that every value repeats 3 times - integer, unsigned, string - because MySQL can convert to
                // any of those for a response. We normally see UInt, but we want this to be safe against
                // surprises
                Value::Int(1) => Ok(ChunkingMethod::ByContentHashBlake2),
                Value::UInt(1) => Ok(ChunkingMethod::ByContentHashBlake2),
                Value::Bytes(ref b) if b == b"1" => Ok(ChunkingMethod::ByContentHashBlake2),
                Value::Int(2) => Ok(ChunkingMethod::InlineBase64),
                Value::UInt(2) => Ok(ChunkingMethod::InlineBase64),
                Value::Bytes(ref b) if b == b"2" => Ok(ChunkingMethod::InlineBase64),
                // If you need to add to this error path, ensure that the type you are adding cannot be converted to an integer
                // by MySQL
                v @ Value::NULL
                | v @ Value::Bytes(..)
                | v @ Value::Float(..)
                | v @ Value::Double(..)
                | v @ Value::Date(..)
                | v @ Value::Time(..)
                | v @ Value::Int(..)
                | v @ Value::UInt(..) => Err(FromValueError(v)),
            }
        }

        fn commit(self) -> ChunkingMethod {
            self
        }

        fn rollback(self) -> Value {
            self.into()
        }
    }

    impl FromValue for ChunkingMethod {
        type Intermediate = ChunkingMethod;
    }
}

pub use self::types::ChunkingMethod;

mononoke_queries! {
    write InsertData(values: (id: &str, ctime: i64, chunk_id: &str, chunk_count: u32, chunking_method: ChunkingMethod)) {
        insert_or_ignore,
        "{insert_or_ignore} INTO data (
            id
            , creation_time
            , chunk_id
            , chunk_count
            , chunking_method
        ) VALUES {values}"
    }

    write DeleteData(id: &str) {
        none,
        "DELETE FROM data WHERE id = {id}"
    }

    write UpdateData(id: &str, ctime: i64, chunk_id: &str, chunk_count: u32, chunking_method: ChunkingMethod) {
        none,
        "UPDATE data SET
            creation_time = {ctime}
            , chunk_id = {chunk_id}
            , chunk_count = {chunk_count}
            , chunking_method = {chunking_method}
        WHERE id = {id}"
    }

    write UpdateDataOptimistic(id: &str, ctime: i64, chunk_id: &str, chunk_count: u32, chunking_method: ChunkingMethod, old_ctime: i64) {
        none,
        "UPDATE data SET
            creation_time = {ctime}
            , chunk_id = {chunk_id}
            , chunk_count = {chunk_count}
            , chunking_method = {chunking_method}
        WHERE id = {id} AND creation_time = {old_ctime}"
    }

    write InsertChunk(values: (id: &str, chunk_num: u32, value: &[u8])) {
        insert_or_ignore,
        "{insert_or_ignore} INTO chunk (
            id
            , chunk_num
            , value
        ) VALUES {values}"
    }

    write UpdateGeneration(id: &str, generation: u64, value_len: u64) {
        none,
        "UPDATE chunk_generation
            SET last_seen_generation = {generation}, value_len = {value_len}
            WHERE id = {id} AND last_seen_generation < {generation}"
    }

    read SelectData(id: &str) -> (i64, Vec<u8>, u32, ChunkingMethod) {
        "SELECT creation_time, chunk_id, chunk_count, chunking_method
         FROM data
         WHERE id = {id}"
    }

    read SelectIsDataPresent(id: &str) -> (i32) {
        "SELECT 1
         FROM data
         WHERE id = {id}"
    }

    read SelectChunk(id: &str, chunk_num: u32) -> (Vec<u8>) {
        "SELECT value
         FROM chunk
         WHERE id = {id}
           AND chunk_num = {chunk_num}"
    }

    read SelectChunkLen(id: &str) -> (u64) {
        "SELECT CAST(SUM(LENGTH(value)) AS UNSIGNED)
         FROM chunk
         WHERE id = {id}"
    }

    read GetChunkGeneration(id: &str) -> (u64, u64) {
        "SELECT last_seen_generation, value_len
        FROM chunk_generation
        WHERE id = {id}"
    }

    write InsertGeneration(values: (id: &str, generation: u64, value_len: u64)) {
        insert_or_ignore,
        "{insert_or_ignore} INTO chunk_generation VALUES {values}"
    }

    read GetNeedsInitialGeneration(limit: u64) -> (Vec<u8>, Option<u64>) {
        "SELECT chunk.id, chunk_generation.value_len
        FROM chunk LEFT JOIN chunk_generation ON chunk.id = chunk_generation.id
        WHERE chunk_generation.last_seen_generation IS NULL
        LIMIT {limit}"
    }

    read GetAllKeys() -> (Vec<u8>) {
        "SELECT id FROM data"
    }

    read GetGenerationSizes() -> (Option<u64>, u64, u64) {
        "SELECT chunk_generation.last_seen_generation, CAST(SUM(chunk_generation.value_len) AS UNSIGNED), CAST(COUNT(1) AS UNSIGNED)
        FROM chunk_generation
        GROUP BY chunk_generation.last_seen_generation"
    }
}

pub struct Chunked {
    pub id: String,
    pub count: u32,
    pub ctime: i64,
    pub chunking_method: ChunkingMethod,
}

#[derive(Clone)]
pub(crate) struct DataSqlStore {
    shard_count: NonZeroUsize,
    write_connection: Arc<Vec1<Connection>>,
    read_connection: Arc<Vec1<Connection>>,
    read_master_connection: Arc<Vec1<Connection>>,
    delay: BlobDelay,
}

impl DataSqlStore {
    pub(crate) fn new(
        shard_count: NonZeroUsize,
        write_connection: Arc<Vec1<Connection>>,
        read_connection: Arc<Vec1<Connection>>,
        read_master_connection: Arc<Vec1<Connection>>,
        delay: BlobDelay,
    ) -> Self {
        Self {
            shard_count,
            write_connection,
            read_connection,
            read_master_connection,
            delay,
        }
    }

    pub(crate) async fn get(&self, key: &str) -> Result<Option<Chunked>, Error> {
        let shard_id = self.shard(key);

        let rows = {
            let rows = SelectData::query(&self.read_connection[shard_id], &key).await?;
            if rows.is_empty() {
                SelectData::query(&self.read_master_connection[shard_id], &key).await?
            } else {
                rows
            }
        };

        Ok(rows
            .into_iter()
            .next()
            .map(|(ctime, chunk_id, chunk_count, chunking_method)| Chunked {
                id: String::from_utf8_lossy(&chunk_id).to_string(),
                count: chunk_count,
                ctime,
                chunking_method,
            }))
    }

    pub(crate) async fn put(
        &self,
        key: &str,
        ctime: i64,
        chunk_id: &str,
        chunk_count: u32,
        chunking_method: ChunkingMethod,
    ) -> Result<(), Error> {
        let shard_id = self.shard(key);

        self.delay.delay(shard_id).await;

        let res = InsertData::query(
            &self.write_connection[shard_id],
            &[(&key, &ctime, &chunk_id, &chunk_count, &chunking_method)],
        )
        .await?;
        if res.affected_rows() == 0 {
            UpdateData::query(
                &self.write_connection[shard_id],
                &key,
                &ctime,
                &chunk_id,
                &chunk_count,
                &chunking_method,
            )
            .await?;
        }
        Ok(())
    }

    // Update optimistically using ctime as the optimistic lock check
    // Used from gc marking to inline small blobs where ctime hasn't changed
    pub(crate) async fn update_optimistic(
        &self,
        key: &str,
        ctime: i64,
        chunk_id: &str,
        chunk_count: u32,
        chunking_method: ChunkingMethod,
        old_ctime: i64,
    ) -> Result<(), Error> {
        let shard_id = self.shard(key);
        self.delay.delay(shard_id).await;

        UpdateDataOptimistic::query(
            &self.write_connection[shard_id],
            &key,
            &ctime,
            &chunk_id,
            &chunk_count,
            &chunking_method,
            &old_ctime,
        )
        .await?;

        Ok(())
    }

    pub(crate) async fn unlink(&self, key: &str) -> Result<(), Error> {
        let shard_id = self.shard(key);

        self.delay.delay(shard_id).await;

        // Deleting from data table does not remove the chunks as they are content addressed.  GC checks for orphaned chunks and removes them.
        let res = DeleteData::query(&self.write_connection[shard_id], &key).await?;
        if res.affected_rows() != 1 {
            bail!(
                "Unexpected row_count {} from sqlblob unlink for {}",
                res.affected_rows(),
                key
            );
        }
        Ok(())
    }

    pub(crate) async fn is_present(&self, key: &str) -> Result<bool, Error> {
        let shard_id = self.shard(key);

        let rows = {
            let rows = SelectIsDataPresent::query(&self.read_connection[shard_id], &key).await?;
            if rows.is_empty() {
                SelectIsDataPresent::query(&self.read_master_connection[shard_id], &key).await?
            } else {
                rows
            }
        };
        Ok(!rows.is_empty())
    }

    pub(crate) fn get_keys_from_shard(
        &self,
        shard_num: usize,
    ) -> impl Stream<Item = Result<String, Error>> {
        let conn = self.read_master_connection[shard_num].clone();
        async move {
            let keys = GetAllKeys::query(&conn).await?;
            Ok(stream::iter(
                keys.into_iter()
                    .map(|(id,)| Ok(String::from_utf8_lossy(&id).to_string())),
            ))
        }
        .try_flatten_stream()
    }

    fn shard(&self, key: &str) -> usize {
        let mut hasher = XxHash32::with_seed(0);
        hasher.write(key.as_bytes());
        (hasher.finish() % self.shard_count.get() as u64) as usize
    }
}
pub(crate) enum ChunkGenerationState {
    NeedsInsertToShard(usize),
    Updated,
}

#[derive(Clone)]
pub(crate) struct ChunkSqlStore {
    shard_count: NonZeroUsize,
    write_connection: Arc<Vec1<Connection>>,
    read_connection: Arc<Vec1<Connection>>,
    read_master_connection: Arc<Vec1<Connection>>,
    delay: BlobDelay,
    gc_generations: ConfigHandle<XdbGc>,
}

impl ChunkSqlStore {
    pub(crate) fn new(
        shard_count: NonZeroUsize,
        write_connection: Arc<Vec1<Connection>>,
        read_connection: Arc<Vec1<Connection>>,
        read_master_connection: Arc<Vec1<Connection>>,
        delay: BlobDelay,
        gc_generations: ConfigHandle<XdbGc>,
    ) -> Self {
        Self {
            shard_count,
            write_connection,
            read_connection,
            read_master_connection,
            delay,
            gc_generations,
        }
    }

    pub(crate) async fn get(
        &self,
        id: &str,
        chunk_num: u32,
        chunking_method: ChunkingMethod,
    ) -> Result<BytesMut, Error> {
        if let Some(shard_id) = self.shard(id, chunk_num, chunking_method) {
            let rows = {
                let rows =
                    SelectChunk::query(&self.read_connection[shard_id], &id, &chunk_num).await?;
                if rows.is_empty() {
                    SelectChunk::query(&self.read_master_connection[shard_id], &id, &chunk_num)
                        .await?
                } else {
                    rows
                }
            };
            rows.into_iter()
                .next()
                .map(|(value,)| (&*value).into())
                .ok_or_else(|| {
                    format_err!("Missing chunk with id {} shard {}", chunk_num, shard_id)
                })
        } else {
            bail!(
                "ChunkSqlStore::get() unexpectedly called for inline chunking_method {:?}",
                chunking_method
            )
        }
    }

    /// Returns the shard and number of chunk_generation rows updated
    pub(crate) async fn put(
        &self,
        key: &str,
        chunk_num: u32,
        chunking_method: ChunkingMethod,
        value: &[u8],
        full_value_len: u64,
    ) -> Result<Option<ChunkGenerationState>, Error> {
        if let Some(shard_id) = self.shard(key, chunk_num, chunking_method) {
            self.delay.delay(shard_id).await;
            let generation = self.gc_generations.get().put_generation as u64;
            let conn = &self.write_connection[shard_id];
            // Update generation incase it already exists
            let updated = UpdateGeneration::query(conn, &key, &generation, &full_value_len).await?;
            InsertChunk::query(conn, &[(&key, &chunk_num, &value)]).await?;
            if updated.affected_rows() > 0 {
                Ok(Some(ChunkGenerationState::Updated))
            } else {
                Ok(Some(ChunkGenerationState::NeedsInsertToShard(shard_id)))
            }
        } else {
            Ok(None)
        }
    }

    pub(crate) fn get_mark_generation(&self) -> u64 {
        self.gc_generations.get().mark_generation as u64
    }

    // Store an entry for value_len eagerly if it was missing on ChunkSqlStore::put()'s UpdateGeneration
    // Saves lazy computing it with associated MySQL read bandwidth from length(chunk.value) later.
    pub(crate) async fn put_chunk_generation(
        &self,
        key: &str,
        shard_id: usize,
        full_value_len: u64,
    ) -> Result<(), Error> {
        let generation = self.gc_generations.get().put_generation as u64;
        let conn = &self.write_connection[shard_id];
        InsertGeneration::query(conn, &[(&key, &generation, &full_value_len)]).await?;
        Ok(())
    }

    pub(crate) async fn update_generation(
        &self,
        key: &str,
        chunk_num: u32,
        chunking_method: ChunkingMethod,
        value_len: u64,
    ) -> Result<(), Error> {
        if let Some(shard_id) = self.shard(key, chunk_num, chunking_method) {
            self.delay.delay(shard_id).await;
            UpdateGeneration::query(
                &self.write_connection[shard_id],
                &key,
                &(self.gc_generations.get().put_generation as u64),
                &value_len,
            )
            .await?;
        }
        Ok(())
    }

    #[cfg(test)]
    pub(crate) async fn get_generation(
        &self,
        key: &str,
        chunk_num: u32,
        chunking_method: ChunkingMethod,
    ) -> Result<Option<u64>, Error> {
        if let Some(shard_id) = self.shard(key, chunk_num, chunking_method) {
            let rows = {
                let rows = GetChunkGeneration::query(&self.read_connection[shard_id], &key).await?;
                if rows.is_empty() {
                    GetChunkGeneration::query(&self.read_master_connection[shard_id], &key).await?
                } else {
                    rows
                }
            };
            Ok(rows.into_iter().next().map(|(v, _value_len)| v))
        } else {
            Ok(None)
        }
    }

    async fn get_len(&self, shard_id: usize, key: &str) -> Result<u64, Error> {
        let rows = {
            let rows = SelectChunkLen::query(&self.read_connection[shard_id], &key).await?;
            if rows.is_empty() {
                SelectChunkLen::query(&self.read_master_connection[shard_id], &key).await?
            } else {
                rows
            }
        };
        rows.into_iter()
            .next()
            .map(|(value_len,)| value_len)
            .ok_or_else(|| format_err!("Missing chunk with id {} shard {}", key, shard_id))
    }

    // Returns length of the chunk value if known
    pub(crate) async fn set_generation(
        &self,
        key: &str,
        chunk_num: u32,
        chunking_method: ChunkingMethod,
        // Take the mark generation as param, so that marking for an entire run is consistent
        mark_generation: u64,
    ) -> Result<Option<u64>, Error> {
        if let Some(shard_id) = self.shard(key, chunk_num, chunking_method) {
            // Take latest value for put generation
            let put_generation = self.gc_generations.get().put_generation as u64;

            // Short-circuit if we have a generation and that generation is >= mark_generation
            let found_generation = GetChunkGeneration::query(&self.read_connection[shard_id], &key)
                .await?
                .into_iter()
                .next();
            let (found_generation, value_len) =
                if let Some((found_generation, value_len)) = found_generation {
                    if found_generation >= mark_generation {
                        return Ok(Some(value_len));
                    }
                    (Some(found_generation), Some(value_len))
                } else {
                    let found_generation =
                        GetChunkGeneration::query(&self.read_master_connection[shard_id], &key)
                            .await?
                            .into_iter()
                            .next();

                    if let Some((found_generation, value_len)) = found_generation {
                        if found_generation >= mark_generation {
                            return Ok(Some(value_len));
                        }
                        (Some(found_generation), Some(value_len))
                    } else {
                        (None, None)
                    }
                };

            // Make sure we know how large the value is
            let value_len: u64 = if let Some(value_len) = value_len {
                value_len
            } else {
                // This chunk has never had value_len populated so get it from chunk.value
                self.get_len(shard_id, key).await?
            };

            // About to start writing so delay
            self.delay.delay(shard_id).await;

            if found_generation.is_none() {
                // First set the generation if unset, so that future writers will update it.
                InsertGeneration::query(
                    &self.write_connection[shard_id],
                    &[(&key, &put_generation, &value_len)],
                )
                .await?;
            }
            // Then update it in case it already existed
            UpdateGeneration::query(
                &self.write_connection[shard_id],
                &key,
                &mark_generation,
                &value_len,
            )
            .await?;

            return Ok(Some(value_len));
        }
        Ok(None)
    }

    // Returns a HashMap from generation->(size, chunk_id_count)
    // Its a chunk id count as some chunk ids have multiple chunks of CHUNK_SIZE
    // but chunk_generation doesn't record that (it doesn't need to)
    pub(crate) async fn get_chunk_sizes_by_generation(
        &self,
        shard_num: usize,
    ) -> Result<HashMap<Option<u64>, (u64, u64)>, Error> {
        GetGenerationSizes::query(&self.read_master_connection[shard_num])
            .await
            .map(|s| {
                s.into_iter()
                    .map(|(gen, size, count)| (gen, (size, count)))
                    .collect::<HashMap<_, (_, _)>>()
            })
    }

    pub(crate) async fn set_initial_generation(&self, shard_num: usize) -> Result<(), Error> {
        loop {
            self.delay.delay(shard_num).await;
            let conn = &self.write_connection[shard_num];
            let chunks_needing_gen = GetNeedsInitialGeneration::query(conn, &10000).await?;
            if chunks_needing_gen.is_empty() {
                return Ok(());
            }

            let generation = self.gc_generations.get().put_generation as u64;

            for (id, value_len) in chunks_needing_gen {
                let id = String::from_utf8_lossy(&id);
                let value_len = if let Some(value_len) = value_len {
                    value_len
                } else {
                    self.get_len(shard_num, &id).await?
                };

                InsertGeneration::query(conn, &[(&id.as_ref(), &generation, &value_len)]).await?;
            }
        }
    }

    // Returns None if the value is stored inline without needing chunk table lookup
    fn shard(&self, key: &str, chunk_id: u32, chunking_method: ChunkingMethod) -> Option<usize> {
        match chunking_method {
            ChunkingMethod::InlineBase64 => None,
            ChunkingMethod::ByContentHashBlake2 => {
                let mut hasher = XxHash32::with_seed(0);
                hasher.write(key.as_bytes());
                hasher.write_u32(chunk_id);
                Some((hasher.finish() % self.shard_count.get() as u64) as usize)
            }
        }
    }
}