Paul.masurel/add kv migration (#5920)

* add migration for in-metastore kv store

* store cluster identity

* test get-identity

---------

Co-authored-by: trinity Pointard <trinity.pointard@datadoghq.com>

Author: fulmicoton

quickwit/Cargo.lock

@@ -37,6 +37,7 @@ tokio-stream = { workspace = true }
 tower = { workspace = true }
 tracing = { workspace = true }
 ulid = { workspace = true, features = ["serde"] }
+uuid = { workspace = true }
 utoipa = { workspace = true }
 
 quickwit-common = { workspace = true }

quickwit/quickwit-metastore/Cargo.toml

@@ -0,0 +1 @@
+DROP TABLE kv;

quickwit/quickwit-metastore/migrations/postgresql/24_add-arbitrary-kv.down.sql

@@ -0,0 +1,4 @@
+CREATE TABLE IF NOT EXISTS kv (
+    key VARCHAR(50) PRIMARY KEY,
+    value TEXT NOT NULL
+);

quickwit/quickwit-metastore/migrations/postgresql/24_add-arbitrary-kv.up.sql

@@ -22,17 +22,17 @@ use quickwit_proto::metastore::{
     CreateIndexResponse, CreateIndexTemplateRequest, DeleteIndexRequest,
     DeleteIndexTemplatesRequest, DeleteQuery, DeleteShardsRequest, DeleteShardsResponse,
     DeleteSourceRequest, DeleteSplitsRequest, DeleteTask, EmptyResponse,
-    FindIndexTemplateMatchesRequest, FindIndexTemplateMatchesResponse, GetIndexTemplateRequest,
-    GetIndexTemplateResponse, IndexMetadataRequest, IndexMetadataResponse, IndexesMetadataRequest,
-    IndexesMetadataResponse, LastDeleteOpstampRequest, LastDeleteOpstampResponse,
-    ListDeleteTasksRequest, ListDeleteTasksResponse, ListIndexTemplatesRequest,
-    ListIndexTemplatesResponse, ListIndexesMetadataRequest, ListIndexesMetadataResponse,
-    ListShardsRequest, ListShardsResponse, ListSplitsRequest, ListSplitsResponse,
-    ListStaleSplitsRequest, MarkSplitsForDeletionRequest, MetastoreResult, MetastoreService,
-    MetastoreServiceClient, MetastoreServiceStream, OpenShardsRequest, OpenShardsResponse,
-    PruneShardsRequest, PublishSplitsRequest, ResetSourceCheckpointRequest, StageSplitsRequest,
-    ToggleSourceRequest, UpdateIndexRequest, UpdateSourceRequest, UpdateSplitsDeleteOpstampRequest,
-    UpdateSplitsDeleteOpstampResponse,
+    FindIndexTemplateMatchesRequest, FindIndexTemplateMatchesResponse, GetClusterIdentityRequest,
+    GetClusterIdentityResponse, GetIndexTemplateRequest, GetIndexTemplateResponse,
+    IndexMetadataRequest, IndexMetadataResponse, IndexesMetadataRequest, IndexesMetadataResponse,
+    LastDeleteOpstampRequest, LastDeleteOpstampResponse, ListDeleteTasksRequest,
+    ListDeleteTasksResponse, ListIndexTemplatesRequest, ListIndexTemplatesResponse,
+    ListIndexesMetadataRequest, ListIndexesMetadataResponse, ListShardsRequest, ListShardsResponse,
+    ListSplitsRequest, ListSplitsResponse, ListStaleSplitsRequest, MarkSplitsForDeletionRequest,
+    MetastoreResult, MetastoreService, MetastoreServiceClient, MetastoreServiceStream,
+    OpenShardsRequest, OpenShardsResponse, PruneShardsRequest, PublishSplitsRequest,
+    ResetSourceCheckpointRequest, StageSplitsRequest, ToggleSourceRequest, UpdateIndexRequest,
+    UpdateSourceRequest, UpdateSplitsDeleteOpstampRequest, UpdateSplitsDeleteOpstampResponse,
 };
 
 /// A [`MetastoreService`] implementation that proxies some requests to the control plane so it can
@@ -274,4 +274,11 @@ impl MetastoreService for ControlPlaneMetastore {
     ) -> MetastoreResult<EmptyResponse> {
         self.metastore.delete_index_templates(request).await
     }
+
+    async fn get_cluster_identity(
+        &self,
+        request: GetClusterIdentityRequest,
+    ) -> MetastoreResult<GetClusterIdentityResponse> {
+        self.metastore.get_cluster_identity(request).await
+    }
 }

quickwit/quickwit-metastore/src/metastore/control_plane_metastore.rs

@@ -23,6 +23,7 @@ use quickwit_proto::types::{DocMappingUid, IndexId};
 use quickwit_storage::{OwnedBytes, Storage, StorageError, StorageErrorKind, StorageResult};
 use serde::{Deserialize, Serialize};
 use tracing::error;
+use uuid::Uuid;
 
 pub(super) const MANIFEST_FILE_NAME: &str = "manifest.json";
 
@@ -40,6 +41,7 @@ impl LegacyManifest {
         Manifest {
             indexes: self.indexes,
             templates: HashMap::new(),
+            identity: Uuid::nil(),
         }
     }
 }
@@ -64,6 +66,7 @@ pub(crate) struct Manifest {
     // The templates are serialized as a sorted `Vec<IndexTemplate>` so the btree map is
     // unnecessary here and we can pass the hash map as is to the `MetastoreState`
     pub templates: HashMap<IndexTemplateId, IndexTemplate>,
+    pub identity: Uuid,
 }
 
 #[derive(Clone, Debug, Serialize, Deserialize)]
@@ -103,6 +106,8 @@ impl From<VersionedManifest> for Manifest {
 struct ManifestV0_8 {
     indexes: BTreeMap<IndexId, IndexStatus>,
     templates: Vec<IndexTemplate>,
+    #[serde(default, skip_serializing_if = "Uuid::is_nil")]
+    identity: Uuid,
 }
 
 impl From<Manifest> for ManifestV0_8 {
@@ -115,6 +120,7 @@ impl From<Manifest> for ManifestV0_8 {
         ManifestV0_8 {
             indexes: manifest.indexes,
             templates,
+            identity: manifest.identity,
         }
     }
 }
@@ -127,7 +133,11 @@ impl From<ManifestV0_8> for Manifest {
             .into_iter()
             .map(|template| (template.template_id.clone(), template))
             .collect();
-        Manifest { indexes, templates }
+        Manifest {
+            indexes,
+            templates,
+            identity: manifest.identity,
+        }
     }
 }
 
@@ -144,7 +154,11 @@ impl quickwit_config::TestableForRegression for Manifest {
             "test-template-1".to_string(),
             IndexTemplate::sample_for_regression(),
         );
-        Manifest { indexes, templates }
+        Manifest {
+            indexes,
+            templates,
+            identity: Uuid::nil(),
+        }
     }
 
     fn assert_equality(&self, other: &Self) {
@@ -320,7 +334,11 @@ mod tests {
                 IndexTemplate::for_test("test-template-2", &["test-index-bar*"], 200),
             ),
         ]);
-        let manifest = Manifest { indexes, templates };
+        let manifest = Manifest {
+            indexes,
+            templates,
+            identity: Uuid::nil(),
+        };
         let manifest_json = serde_json::to_string_pretty(&manifest).unwrap();
         let manifest_deserialized: Manifest = serde_json::from_str(&manifest_json).unwrap();
         assert_eq!(manifest, manifest_deserialized);

quickwit/quickwit-metastore/src/metastore/file_backed/manifest.rs

@@ -44,24 +44,26 @@ use quickwit_proto::metastore::{
     CreateIndexResponse, CreateIndexTemplateRequest, DeleteIndexRequest,
     DeleteIndexTemplatesRequest, DeleteQuery, DeleteShardsRequest, DeleteShardsResponse,
     DeleteSourceRequest, DeleteSplitsRequest, DeleteTask, EmptyResponse, EntityKind,
-    FindIndexTemplateMatchesRequest, FindIndexTemplateMatchesResponse, GetIndexTemplateRequest,
-    GetIndexTemplateResponse, IndexMetadataFailure, IndexMetadataFailureReason,
-    IndexMetadataRequest, IndexMetadataResponse, IndexTemplateMatch, IndexesMetadataRequest,
-    IndexesMetadataResponse, LastDeleteOpstampRequest, LastDeleteOpstampResponse,
-    ListDeleteTasksRequest, ListDeleteTasksResponse, ListIndexTemplatesRequest,
-    ListIndexTemplatesResponse, ListIndexesMetadataRequest, ListIndexesMetadataResponse,
-    ListShardsRequest, ListShardsResponse, ListSplitsRequest, ListSplitsResponse,
-    ListStaleSplitsRequest, MarkSplitsForDeletionRequest, MetastoreError, MetastoreResult,
-    MetastoreService, MetastoreServiceStream, OpenShardSubrequest, OpenShardsRequest,
-    OpenShardsResponse, PruneShardsRequest, PublishSplitsRequest, ResetSourceCheckpointRequest,
-    StageSplitsRequest, ToggleSourceRequest, UpdateIndexRequest, UpdateSourceRequest,
-    UpdateSplitsDeleteOpstampRequest, UpdateSplitsDeleteOpstampResponse, serde_utils,
+    FindIndexTemplateMatchesRequest, FindIndexTemplateMatchesResponse, GetClusterIdentityRequest,
+    GetClusterIdentityResponse, GetIndexTemplateRequest, GetIndexTemplateResponse,
+    IndexMetadataFailure, IndexMetadataFailureReason, IndexMetadataRequest, IndexMetadataResponse,
+    IndexTemplateMatch, IndexesMetadataRequest, IndexesMetadataResponse, LastDeleteOpstampRequest,
+    LastDeleteOpstampResponse, ListDeleteTasksRequest, ListDeleteTasksResponse,
+    ListIndexTemplatesRequest, ListIndexTemplatesResponse, ListIndexesMetadataRequest,
+    ListIndexesMetadataResponse, ListShardsRequest, ListShardsResponse, ListSplitsRequest,
+    ListSplitsResponse, ListStaleSplitsRequest, MarkSplitsForDeletionRequest, MetastoreError,
+    MetastoreResult, MetastoreService, MetastoreServiceStream, OpenShardSubrequest,
+    OpenShardsRequest, OpenShardsResponse, PruneShardsRequest, PublishSplitsRequest,
+    ResetSourceCheckpointRequest, StageSplitsRequest, ToggleSourceRequest, UpdateIndexRequest,
+    UpdateSourceRequest, UpdateSplitsDeleteOpstampRequest, UpdateSplitsDeleteOpstampResponse,
+    serde_utils,
 };
 use quickwit_proto::types::{IndexId, IndexUid};
 use quickwit_storage::Storage;
 use time::OffsetDateTime;
 use tokio::sync::{Mutex, OwnedMutexGuard, RwLock};
 use ulid::Ulid;
+use uuid::Uuid;
 
 use self::file_backed_index::FileBackedIndex;
 pub use self::file_backed_metastore_factory::FileBackedMetastoreFactory;
@@ -1202,6 +1204,32 @@ impl MetastoreService for FileBackedMetastore {
         }
         Ok(EmptyResponse {})
     }
+
+    // Get cluster identity api
+
+    // this returns a constant uuid. on first call, it generate said uuid if it doesn't already
+    // exists
+    async fn get_cluster_identity(
+        &self,
+        _: GetClusterIdentityRequest,
+    ) -> MetastoreResult<GetClusterIdentityResponse> {
+        let mut state_wlock_guard = self.state.write().await;
+
+        if state_wlock_guard.identity.is_nil() {
+            state_wlock_guard.identity = Uuid::new_v4();
+
+            let manifest = state_wlock_guard.as_manifest();
+
+            if let Err(error) = save_manifest(&*self.storage, &manifest).await {
+                state_wlock_guard.identity = Uuid::nil();
+                return Err(error);
+            }
+        }
+
+        Ok(GetClusterIdentityResponse {
+            uuid: state_wlock_guard.identity.hyphenated().to_string(),
+        })
+    }
 }
 
 impl MetastoreServiceExt for FileBackedMetastore {}

quickwit/quickwit-metastore/src/metastore/file_backed/mod.rs

@@ -20,6 +20,7 @@ use quickwit_config::{IndexTemplate, IndexTemplateId};
 use quickwit_proto::metastore::MetastoreResult;
 use quickwit_proto::types::IndexId;
 use quickwit_storage::Storage;
+use uuid::Uuid;
 
 use super::LazyIndexStatus;
 use super::index_template_matcher::IndexTemplateMatcher;
@@ -31,6 +32,7 @@ pub(super) struct MetastoreState {
     pub indexes: HashMap<IndexId, LazyIndexStatus>,
     pub templates: HashMap<IndexTemplateId, IndexTemplate>,
     pub template_matcher: IndexTemplateMatcher,
+    pub identity: Uuid,
 }
 
 impl MetastoreState {
@@ -64,6 +66,7 @@ impl MetastoreState {
             indexes,
             templates: manifest.templates,
             template_matcher,
+            identity: manifest.identity,
         };
         Ok(state)
     }
@@ -82,6 +85,10 @@ impl MetastoreState {
             })
             .collect();
         let templates = self.templates.clone();
-        Manifest { indexes, templates }
+        Manifest {
+            indexes,
+            templates,
+            identity: self.identity,
+        }
     }
 }

quickwit/quickwit-metastore/src/metastore/file_backed/state.rs

@@ -31,15 +31,15 @@ use quickwit_proto::metastore::{
     CreateIndexResponse, CreateIndexTemplateRequest, DeleteIndexRequest,
     DeleteIndexTemplatesRequest, DeleteQuery, DeleteShardsRequest, DeleteShardsResponse,
     DeleteSourceRequest, DeleteSplitsRequest, DeleteTask, EmptyResponse, EntityKind,
-    FindIndexTemplateMatchesRequest, FindIndexTemplateMatchesResponse, GetIndexTemplateRequest,
-    GetIndexTemplateResponse, IndexMetadataFailure, IndexMetadataFailureReason,
-    IndexMetadataRequest, IndexMetadataResponse, IndexTemplateMatch, IndexesMetadataRequest,
-    IndexesMetadataResponse, LastDeleteOpstampRequest, LastDeleteOpstampResponse,
-    ListDeleteTasksRequest, ListDeleteTasksResponse, ListIndexTemplatesRequest,
-    ListIndexTemplatesResponse, ListIndexesMetadataRequest, ListIndexesMetadataResponse,
-    ListShardsRequest, ListShardsResponse, ListShardsSubresponse, ListSplitsRequest,
-    ListSplitsResponse, ListStaleSplitsRequest, MarkSplitsForDeletionRequest, MetastoreError,
-    MetastoreResult, MetastoreService, MetastoreServiceStream, OpenShardSubrequest,
+    FindIndexTemplateMatchesRequest, FindIndexTemplateMatchesResponse, GetClusterIdentityRequest,
+    GetClusterIdentityResponse, GetIndexTemplateRequest, GetIndexTemplateResponse,
+    IndexMetadataFailure, IndexMetadataFailureReason, IndexMetadataRequest, IndexMetadataResponse,
+    IndexTemplateMatch, IndexesMetadataRequest, IndexesMetadataResponse, LastDeleteOpstampRequest,
+    LastDeleteOpstampResponse, ListDeleteTasksRequest, ListDeleteTasksResponse,
+    ListIndexTemplatesRequest, ListIndexTemplatesResponse, ListIndexesMetadataRequest,
+    ListIndexesMetadataResponse, ListShardsRequest, ListShardsResponse, ListShardsSubresponse,
+    ListSplitsRequest, ListSplitsResponse, ListStaleSplitsRequest, MarkSplitsForDeletionRequest,
+    MetastoreError, MetastoreResult, MetastoreService, MetastoreServiceStream, OpenShardSubrequest,
     OpenShardSubresponse, OpenShardsRequest, OpenShardsResponse, PruneShardsRequest,
     PublishSplitsRequest, ResetSourceCheckpointRequest, StageSplitsRequest, ToggleSourceRequest,
     UpdateIndexRequest, UpdateSourceRequest, UpdateSplitsDeleteOpstampRequest,
@@ -51,6 +51,7 @@ use sea_query_binder::SqlxBinder;
 use sqlx::{Acquire, Executor, Postgres, Transaction};
 use time::OffsetDateTime;
 use tracing::{debug, info, instrument, warn};
+use uuid::Uuid;
 
 use super::error::convert_sqlx_err;
 use super::migrator::run_migrations;
@@ -1664,6 +1665,26 @@ impl MetastoreService for PostgresqlMetastore {
             .await?;
         Ok(EmptyResponse {})
     }
+
+    async fn get_cluster_identity(
+        &self,
+        _: GetClusterIdentityRequest,
+    ) -> MetastoreResult<GetClusterIdentityResponse> {
+        let (uuid,) = sqlx::query_as(
+            r"
+                WITH insert AS (
+                    INSERT INTO kv (key, value)
+                           VALUES ('cluster_identity', $1)
+                           ON CONFLICT (key) DO NOTHING
+                )
+                SELECT value FROM kv where key = 'cluster_identity';
+                ",
+        )
+        .bind(Uuid::new_v4().hyphenated().to_string())
+        .fetch_one(&self.connection_pool)
+        .await?;
+        Ok(GetClusterIdentityResponse { uuid })
+    }
 }
 
 async fn open_or_fetch_shard<'e>(

quickwit/quickwit-metastore/src/metastore/postgres/metastore.rs

@@ -0,0 +1,48 @@
+// Copyright 2021-Present Datadog, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// Index API tests
+//
+//  - create_index
+//  - index_exists
+//  - index_metadata
+//  - list_indexes
+//  - delete_index
+
+use quickwit_proto::metastore::{GetClusterIdentityRequest, MetastoreService};
+use uuid::Uuid;
+
+use super::DefaultForTest;
+use crate::MetastoreServiceExt;
+
+pub async fn test_metastore_get_identity<
+    MetastoreToTest: MetastoreService + MetastoreServiceExt + DefaultForTest,
+>() {
+    let metastore = MetastoreToTest::default_for_test().await;
+
+    let identity_1 = metastore
+        .get_cluster_identity(GetClusterIdentityRequest {})
+        .await
+        .unwrap()
+        .uuid;
+
+    let identity_2 = metastore
+        .get_cluster_identity(GetClusterIdentityRequest {})
+        .await
+        .unwrap()
+        .uuid;
+
+    assert_eq!(identity_1, identity_2);
+    assert_ne!(identity_1, Uuid::nil().hyphenated().to_string());
+}