cherry-pick refactor: secret ref columns in protobuf (#17335)

Co-authored-by: August <[email protected]>
risingwavelabs · Jun 19, 2024 · bdf9aff · bdf9aff
1 parent afc6106
commit bdf9aff
Show file tree

Hide file tree

Showing 12 changed files with 112 additions and 17 deletions.
diff --git a/proto/catalog.proto b/proto/catalog.proto
@@ -6,6 +6,7 @@ import "common.proto";
 import "data.proto";
 import "expr.proto";
 import "plan_common.proto";
+import "secret.proto";
 
 option java_package = "com.risingwave.proto";
 option optimize_for = SPEED;
@@ -83,8 +84,8 @@ message StreamSourceInfo {
   // Options specified by user in the FORMAT ENCODE clause.
   map<string, string> format_encode_options = 14;
 
-  // Handle the source relies on any sceret. The key is the propertity name and the value is the secret id.
-  map<string, uint32> secret_ref = 16;
+  // Handle the source relies on any sceret. The key is the propertity name and the value is the secret id and type.
+  map<string, secret.SecretRef> secret_ref = 16;
 }
 
 message Source {
@@ -180,8 +181,8 @@ message Sink {
   // Whether it should use background ddl or block until backfill finishes.
   CreateType create_type = 24;
 
-  // Handle the sink relies on any sceret. The key is the propertity name and the value is the secret id.
-  map<string, uint32> secret_ref = 25;
+  // Handle the sink relies on any sceret. The key is the propertity name and the value is the secret id and type.
+  map<string, secret.SecretRef> secret_ref = 25;
 }
 
 message Subscription {

diff --git a/proto/secret.proto b/proto/secret.proto
@@ -18,3 +18,14 @@ message Secret {
     SecretHashicropValutBackend hashicorp_vault = 2;
   }
 }
+
+message SecretRef {
+  enum RefAsType {
+    UNSPECIFIED = 0;
+    TEXT = 1;
+    // AS FILE
+    FILE = 2;
+  }
+  uint32 secret_id = 1;
+  RefAsType ref_as = 2;
+}
diff --git a/src/connector/src/sink/catalog/desc.rs b/src/connector/src/sink/catalog/desc.rs
@@ -19,6 +19,7 @@ use risingwave_common::catalog::{
     ColumnCatalog, ConnectionId, CreateType, DatabaseId, SchemaId, TableId, UserId,
 };
 use risingwave_common::util::sort_util::ColumnOrder;
+use risingwave_pb::secret::PbSecretRef;
 use risingwave_pb::stream_plan::PbSinkDesc;
 
 use super::{SinkCatalog, SinkFormatDesc, SinkId, SinkType};
@@ -83,7 +84,7 @@ impl SinkDesc {
         owner: UserId,
         connection_id: Option<ConnectionId>,
         dependent_relations: Vec<TableId>,
-        secret_ref: BTreeMap<String, u32>,
+        secret_ref: BTreeMap<String, PbSecretRef>,
     ) -> SinkCatalog {
         SinkCatalog {
             id: self.id,

diff --git a/src/connector/src/sink/catalog/mod.rs b/src/connector/src/sink/catalog/mod.rs
@@ -27,6 +27,7 @@ use risingwave_common::util::sort_util::ColumnOrder;
 use risingwave_pb::catalog::{
     PbCreateType, PbSink, PbSinkFormatDesc, PbSinkType, PbStreamJobStatus,
 };
+use risingwave_pb::secret::PbSecretRef;
 
 use super::{
     SinkError, CONNECTOR_TYPE_KEY, SINK_TYPE_APPEND_ONLY, SINK_TYPE_DEBEZIUM, SINK_TYPE_OPTION,
@@ -339,7 +340,7 @@ pub struct SinkCatalog {
     pub create_type: CreateType,
 
     /// The secret reference for the sink, mapping from property name to secret id.
-    pub secret_ref: BTreeMap<String, u32>,
+    pub secret_ref: BTreeMap<String, PbSecretRef>,
 }
 
 impl SinkCatalog {

diff --git a/src/ctl/src/cmd_impl/meta/migration.rs b/src/ctl/src/cmd_impl/meta/migration.rs
@@ -555,12 +555,12 @@ pub async fn migrate(from: EtcdBackend, target: String, force_clean: bool) -> an
                     *id = *connection_rewrite.get(id).unwrap();
                 }
                 for secret_id in s.secret_ref.values_mut() {
-                    *secret_id = *secret_rewrite.get(secret_id).unwrap();
+                    secret_id.secret_id = *secret_rewrite.get(&secret_id.secret_id).unwrap();
                 }
                 object_dependencies.extend(s.secret_ref.values().map(|id| {
                     object_dependency::ActiveModel {
                         id: NotSet,
-                        oid: Set(*id as _),
+                        oid: Set(id.secret_id as _),
                         used_by: Set(s.id as _),
                     }
                 }));

diff --git a/src/frontend/src/utils/with_options.rs b/src/frontend/src/utils/with_options.rs
@@ -19,6 +19,7 @@ use risingwave_connector::source::kafka::private_link::{
     insert_privatelink_broker_rewrite_map, CONNECTION_NAME_KEY, PRIVATELINK_ENDPOINT_KEY,
 };
 use risingwave_connector::WithPropertiesExt;
+use risingwave_pb::secret::PbSecretRef;
 use risingwave_sqlparser::ast::{
     CreateConnectionStatement, CreateSinkStatement, CreateSourceStatement,
     CreateSubscriptionStatement, SqlOption, Statement, Value,
@@ -119,7 +120,7 @@ impl WithOptions {
 pub(crate) fn resolve_secret_in_with_options(
     _with_options: &mut WithOptions,
     _session: &SessionImpl,
-) -> RwResult<BTreeMap<String, u32>> {
+) -> RwResult<BTreeMap<String, PbSecretRef>> {
     // todo: implement the function and take `resolve_privatelink_in_with_option` as reference
 
     Ok(BTreeMap::new())

diff --git a/src/meta/model_v2/migration/src/m20240525_090457_secret.rs b/src/meta/model_v2/migration/src/m20240525_090457_secret.rs
@@ -37,12 +37,22 @@ impl MigrationTrait for Migration {
             )
             .await?;
 
-        // Add a new column to the table
+        // Add a new column to the `sink` table
         manager
             .alter_table(
                 MigrationTable::alter()
                     .table(Sink::Table)
-                    .add_column(ColumnDef::new(Sink::SecretRef).json_binary())
+                    .add_column(ColumnDef::new(Sink::SecretRef).binary())
+                    .to_owned(),
+            )
+            .await?;
+
+        // Add a new column to the `source` table
+        manager
+            .alter_table(
+                MigrationTable::alter()
+                    .table(Source::Table)
+                    .add_column(ColumnDef::new(Source::SecretRef).binary())
                     .to_owned(),
             )
             .await?;
@@ -60,6 +70,14 @@ impl MigrationTrait for Migration {
                     .to_owned(),
             )
             .await?;
+        manager
+            .alter_table(
+                MigrationTable::alter()
+                    .table(Source::Table)
+                    .drop_column(Source::SecretRef)
+                    .to_owned(),
+            )
+            .await?;
         Ok(())
     }
 }
@@ -77,3 +95,9 @@ enum Sink {
     Table,
     SecretRef,
 }
+
+#[derive(DeriveIden)]
+enum Source {
+    Table,
+    SecretRef,
+}
diff --git a/src/meta/model_v2/src/lib.rs b/src/meta/model_v2/src/lib.rs
@@ -16,6 +16,7 @@ use std::collections::BTreeMap;
 
 use risingwave_pb::catalog::{PbCreateType, PbStreamJobStatus};
 use risingwave_pb::meta::table_fragments::PbState as PbStreamJobState;
+use risingwave_pb::secret::PbSecretRef;
 use risingwave_pb::stream_plan::PbStreamNode;
 use sea_orm::entity::prelude::*;
 use sea_orm::{DeriveActiveEnum, EnumIter, FromJsonQueryResult};
@@ -258,6 +259,55 @@ macro_rules! derive_array_from_blob {
     };
 }
 
+macro_rules! derive_btreemap_from_blob {
+    ($struct_name:ident, $key_type:ty, $value_type:ty, $field_type:ident) => {
+        #[derive(Clone, PartialEq, Eq, DeriveValueType, serde::Deserialize, serde::Serialize)]
+        pub struct $struct_name(#[sea_orm] Vec<u8>);
+
+        #[derive(Clone, PartialEq, ::prost::Message)]
+        pub struct $field_type {
+            #[prost(btree_map = "string, message")]
+            inner: BTreeMap<$key_type, $value_type>,
+        }
+        impl Eq for $field_type {}
+
+        impl $struct_name {
+            pub fn to_protobuf(&self) -> BTreeMap<$key_type, $value_type> {
+                let data: $field_type = prost::Message::decode(self.0.as_slice()).unwrap();
+                data.inner
+            }
+
+            fn from_protobuf(val: BTreeMap<$key_type, $value_type>) -> Self {
+                Self(prost::Message::encode_to_vec(&$field_type { inner: val }))
+            }
+        }
+
+        impl From<BTreeMap<$key_type, $value_type>> for $struct_name {
+            fn from(value: BTreeMap<$key_type, $value_type>) -> Self {
+                Self::from_protobuf(value)
+            }
+        }
+
+        impl std::fmt::Debug for $struct_name {
+            fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+                self.to_protobuf().fmt(f)
+            }
+        }
+
+        impl Default for $struct_name {
+            fn default() -> Self {
+                Self(vec![])
+            }
+        }
+
+        impl sea_orm::sea_query::Nullable for $struct_name {
+            fn null() -> Value {
+                Value::Bytes(None)
+            }
+        }
+    };
+}
+
 pub(crate) use {derive_array_from_blob, derive_from_blob};
 
 derive_from_json_struct!(I32Array, Vec<i32>);
@@ -286,7 +336,7 @@ impl From<BTreeMap<u32, Vec<u32>>> for ActorUpstreamActors {
     }
 }
 
-derive_from_json_struct!(SecretRef, BTreeMap<String, u32>);
+derive_btreemap_from_blob!(SecretRef, String, PbSecretRef, PbSecretRefMap);
 
 derive_from_blob!(StreamNode, PbStreamNode);
 derive_from_blob!(DataType, risingwave_pb::data::PbDataType);

diff --git a/src/meta/model_v2/src/sink.rs b/src/meta/model_v2/src/sink.rs
@@ -72,7 +72,7 @@ pub struct Model {
     pub sink_from_name: String,
     pub sink_format_desc: Option<SinkFormatDesc>,
     pub target_table: Option<TableId>,
-    // `secret_ref` stores a json string, mapping from property name to secret id.
+    // `secret_ref` stores the mapping info mapping from property name to secret id and type.
     pub secret_ref: Option<SecretRef>,
 }
 

diff --git a/src/meta/model_v2/src/source.rs b/src/meta/model_v2/src/source.rs
@@ -19,8 +19,8 @@ use sea_orm::ActiveValue::Set;
 use serde::{Deserialize, Serialize};
 
 use crate::{
-    ColumnCatalogArray, ConnectionId, I32Array, Property, SourceId, StreamSourceInfo, TableId,
-    WatermarkDescArray,
+    ColumnCatalogArray, ConnectionId, I32Array, Property, SecretRef, SourceId, StreamSourceInfo,
+    TableId, WatermarkDescArray,
 };
 
 #[derive(Clone, Debug, PartialEq, DeriveEntityModel, Eq, Serialize, Deserialize)]
@@ -39,6 +39,8 @@ pub struct Model {
     pub optional_associated_table_id: Option<TableId>,
     pub connection_id: Option<ConnectionId>,
     pub version: i64,
+    // `secret_ref` stores the mapping info mapping from property name to secret id and type.
+    pub secret_ref: Option<SecretRef>,
 }
 
 #[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
@@ -101,6 +103,7 @@ impl From<PbSource> for ActiveModel {
             optional_associated_table_id: Set(optional_associated_table_id),
             connection_id: Set(source.connection_id.map(|id| id as _)),
             version: Set(source.version as _),
+            secret_ref: Set(None),
         }
     }
 }
diff --git a/src/meta/src/controller/mod.rs b/src/meta/src/controller/mod.rs
@@ -201,9 +201,9 @@ impl From<ObjectModel<source::Model>> for PbSource {
 
 impl From<ObjectModel<sink::Model>> for PbSink {
     fn from(value: ObjectModel<sink::Model>) -> Self {
-        let mut secret_ref_map: BTreeMap<String, u32> = BTreeMap::new();
+        let mut secret_ref_map = BTreeMap::new();
         if let Some(secret_ref) = value.0.secret_ref {
-            secret_ref_map = secret_ref.into_inner();
+            secret_ref_map = secret_ref.to_protobuf();
         }
         Self {
             id: value.0.sink_id as _,

diff --git a/src/prost/build.rs b/src/prost/build.rs
@@ -66,6 +66,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
         ".plan_common.ExternalTableDesc",
         ".hummock.CompactTask",
         ".catalog.StreamSourceInfo",
+        ".catalog.SecretRef",
         ".catalog.Source",
         ".catalog.Sink",
         ".catalog.View",
@@ -111,6 +112,8 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
         // The requirement is from Source node -> SourceCatalog -> WatermarkDesc -> expr
         .type_attribute("catalog.WatermarkDesc", "#[derive(Eq, Hash)]")
         .type_attribute("catalog.StreamSourceInfo", "#[derive(Eq, Hash)]")
+        .type_attribute("secret.SecretRef", "#[derive(Eq, Hash)]")
+        .type_attribute("catalog.IndexColumnProperties", "#[derive(Eq, Hash)]")
         .type_attribute("expr.ExprNode", "#[derive(Eq, Hash)]")
         .type_attribute("data.DataType", "#[derive(Eq, Hash)]")
         .type_attribute("expr.ExprNode.rex_node", "#[derive(Eq, Hash)]")