temp refactor

Author: Mouli Mukherjee

Cargo.lock

@@ -17,6 +17,7 @@ use std::time::Duration;
 use futures::Future;
 use itertools::Itertools;
 use mz_adapter_types::connection::ConnectionId;
+use mz_storage_types::connections::aws::AwsPrincipalContext;
 use tokio::sync::mpsc::UnboundedSender;
 use tokio::sync::MutexGuard;
 use tracing::{info, trace};
@@ -89,7 +90,7 @@ use mz_storage_types::sources::Timeline;
 use mz_transform::dataflow::DataflowMetainfo;
 
 pub use crate::catalog::builtin_table_updates::BuiltinTableUpdate;
-pub use crate::catalog::config::{AwsPrincipalContext, ClusterReplicaSizeMap, Config, StateConfig};
+pub use crate::catalog::config::{ClusterReplicaSizeMap, Config, StateConfig};
 pub use crate::catalog::open::BuiltinMigrationMetadata;
 pub use crate::catalog::state::CatalogState;
 use crate::command::CatalogDump;
@@ -4220,6 +4221,10 @@ impl SessionCatalog for ConnCatalog<'_> {
         self.state.aws_privatelink_availability_zones.clone()
     }
 
+    fn aws_principal_context(&self) -> Option<AwsPrincipalContext> {
+        self.state.aws_principal_context.clone()
+    }
+
     fn system_vars(&self) -> &SystemVars {
         &self.state.system_configuration
     }

src/adapter/src/catalog.rs

@@ -46,7 +46,7 @@ use mz_sql::catalog::{CatalogCluster, CatalogDatabase, CatalogSchema, CatalogTyp
 use mz_sql::func::FuncImplCatalogDetails;
 use mz_sql::names::{CommentObjectId, ResolvedDatabaseSpecifier, SchemaId, SchemaSpecifier};
 use mz_sql_parser::ast::display::AstDisplay;
-use mz_storage_types::connections::aws::{AwsAuth, AwsConfig};
+use mz_storage_types::connections::aws::{AwsAuth, AwsConfig, AwsPrincipalContext};
 use mz_storage_types::connections::inline::ReferencedConnection;
 use mz_storage_types::connections::{KafkaConnection, StringOrSecret};
 use mz_storage_types::sinks::{KafkaSinkConnection, StorageSinkConnection};
@@ -55,8 +55,8 @@ use mz_storage_types::sources::{
 };
 
 use crate::catalog::{
-    AwsPrincipalContext, CatalogItem, CatalogState, ClusterVariant, Connection, DataSourceDesc,
-    Database, DefaultPrivilegeObject, Func, Index, MaterializedView, Sink, Type, View,
+    CatalogItem, CatalogState, ClusterVariant, Connection, DataSourceDesc, Database,
+    DefaultPrivilegeObject, Func, Index, MaterializedView, Sink, Type, View,
 };
 use crate::coord::ConnMeta;
 use crate::subscribe::ActiveSubscribe;

src/adapter/src/catalog/builtin_table_updates.rs

@@ -15,15 +15,14 @@ use std::time::Duration;
 use bytesize::ByteSize;
 use mz_build_info::BuildInfo;
 use mz_catalog;
-use mz_cloud_resources::AwsExternalIdPrefix;
 use mz_controller::clusters::ReplicaAllocation;
 use mz_orchestrator::MemoryLimit;
 use mz_ore::cast::CastFrom;
 use mz_ore::metrics::MetricsRegistry;
-use mz_repr::GlobalId;
 use mz_secrets::SecretsReader;
 use mz_sql::catalog::EnvironmentId;
 use mz_sql::session::vars::ConnectionCounter;
+use mz_storage_types::connections::aws::AwsPrincipalContext;
 use serde::{Deserialize, Serialize};
 
 use crate::config::SystemParameterSyncConfig;
@@ -223,61 +222,3 @@ impl Default for ClusterReplicaSizeMap {
         Self(inner)
     }
 }
-
-/// Context used to generate an AWS Principal.
-///
-/// In the case of AWS PrivateLink connections, Materialize will connect to the
-/// VPC endpoint as the AWS Principal generated via this context.
-#[derive(Debug, Clone, Serialize)]
-pub struct AwsPrincipalContext {
-    pub aws_account_id: String,
-    pub aws_external_id_prefix: AwsExternalIdPrefix,
-}
-
-impl AwsPrincipalContext {
-    pub fn to_privatelink_principal_string(&self, aws_external_id_suffix: GlobalId) -> String {
-        format!(
-            "arn:aws:iam::{}:role/mz_{}_{}",
-            self.aws_account_id, self.aws_external_id_prefix, aws_external_id_suffix
-        )
-    }
-
-    pub fn to_aws_connection_principal_string(&self) -> String {
-        format!(
-            "arn:aws:iam::{}:role/MaterializeConnection",
-            self.aws_account_id
-        )
-    }
-
-    pub fn to_aws_connection_external_id(&self, aws_external_id_suffix: GlobalId) -> String {
-        format!(
-            "mz_{}_{}",
-            self.aws_external_id_prefix, aws_external_id_suffix
-        )
-    }
-
-    pub fn to_aws_example_trust_policy(
-        &self,
-        aws_external_id_suffix: GlobalId,
-    ) -> serde_json::Value {
-        serde_json::json!(
-            {
-            "Version": "2012-10-17",
-            "Statement": [
-              {
-                "Effect": "Allow",
-                "Principal": {
-                  "AWS": self.to_aws_connection_principal_string()
-                },
-                "Action": "sts:AssumeRole",
-                "Condition": {
-                  "StringEquals": {
-                    "sts:ExternalId": self.to_aws_connection_external_id(aws_external_id_suffix)
-                  }
-                }
-              }
-            ]
-          }
-        )
-    }
-}

src/adapter/src/catalog/config.rs

@@ -17,6 +17,7 @@ use std::time::{Duration, Instant};
 use anyhow::bail;
 use itertools::Itertools;
 use mz_adapter_types::connection::ConnectionId;
+use mz_storage_types::connections::aws::AwsPrincipalContext;
 use serde::Serialize;
 use tokio::sync::mpsc;
 use tracing::info;
@@ -72,7 +73,7 @@ use mz_storage_types::connections::inline::{
 };
 use mz_transform::Optimizer;
 
-use crate::catalog::{AwsPrincipalContext, BuiltinTableUpdate, ClusterReplicaSizeMap, ConnCatalog};
+use crate::catalog::{BuiltinTableUpdate, ClusterReplicaSizeMap, ConnCatalog};
 use crate::coord::ConnMeta;
 use crate::optimize::{self, Optimize};
 use crate::session::Session;

src/adapter/src/catalog/state.rs

@@ -112,6 +112,7 @@ use mz_sql::rbac::UnauthorizedError;
 use mz_sql::session::user::{RoleMetadata, User};
 use mz_sql::session::vars::{self, ConnectionCounter};
 use mz_storage_client::controller::{CollectionDescription, DataSource, DataSourceOther};
+use mz_storage_types::connections::aws::AwsPrincipalContext;
 use mz_storage_types::connections::inline::IntoInlineConnection;
 use mz_storage_types::connections::ConnectionContext;
 use mz_storage_types::sources::Timeline;
@@ -127,7 +128,7 @@ use tracing_opentelemetry::OpenTelemetrySpanExt;
 use uuid::Uuid;
 
 use crate::catalog::{
-    self, AwsPrincipalContext, BuiltinMigrationMetadata, BuiltinTableUpdate, Catalog, CatalogItem,
+    self, BuiltinMigrationMetadata, BuiltinTableUpdate, Catalog, CatalogItem,
     ClusterReplicaSizeMap, Connection, DataSourceDesc, Source,
 };
 use crate::client::{Client, Handle};

src/adapter/src/coord.rs

@@ -321,7 +321,6 @@ async fn run(args: Args) -> Result<(), anyhow::Error> {
         ConnectionContext::from_cli_args(
             &args.tracing.startup_log_filter,
             args.aws_external_id,
-            None, // TODO(mouli): fix passing to clusterd
             secrets_reader,
             None,
         ),

src/clusterd/src/bin/clusterd.rs

@@ -981,7 +981,6 @@ fn run(mut args: Args) -> Result<(), anyhow::Error> {
                 connection_context: ConnectionContext::from_cli_args(
                     &args.tracing.startup_log_filter,
                     args.aws_external_id_prefix,
-                    args.aws_account_id.clone(),
                     secrets_reader,
                     cloud_resource_reader,
                 ),

src/environmentd/src/bin/environmentd/main.rs

@@ -31,6 +31,7 @@ use mz_repr::explain::ExprHumanizer;
 use mz_repr::role_id::RoleId;
 use mz_repr::{ColumnName, GlobalId, RelationDesc};
 use mz_sql_parser::ast::{Expr, Ident, QualifiedReplica, UnresolvedItemName};
+use mz_storage_types::connections::aws::AwsPrincipalContext;
 use mz_storage_types::connections::inline::{ConnectionResolver, ReferencedConnection};
 use mz_storage_types::connections::Connection;
 use mz_storage_types::sources::SourceDesc;
@@ -275,6 +276,9 @@ pub trait SessionCatalog: fmt::Debug + ExprHumanizer + Send + Sync + ConnectionR
     /// Returns the set of supported AWS PrivateLink availability zone ids.
     fn aws_privatelink_availability_zones(&self) -> Option<BTreeSet<String>>;
 
+    /// Returns the AWS context for using AWS connections
+    fn aws_principal_context(&self) -> Option<AwsPrincipalContext>;
+
     /// Returns system vars
     fn system_vars(&self) -> &SystemVars;
 

src/sql/src/catalog.rs

@@ -184,9 +184,18 @@ impl ConnectionOptionExtracted {
 
                 let assume_role = match (self.assume_role_arn, self.assume_role_session_name) {
                     (Some(arn), session_name) => {
-                        Some(AwsAuth::AssumeRole(AwsAssumeRole { arn, session_name }))
+                        let aws_principal_context =
+                            scx.catalog.aws_principal_context().ok_or_else(|| {
+                                sql_err!("AWS principal context should have been provided")
+                            })?;
+                        Some(AwsAuth::AssumeRole(AwsAssumeRole {
+                            arn,
+                            session_name,
+                            mz_principal: aws_principal_context
+                                .to_aws_connection_principal_string(),
+                            external_id_prefix: aws_principal_context.aws_external_id_prefix.to_string(),
+                        }))
                     }
-
                     _ => None,
                 };
 

src/sql/src/plan/statement/ddl/connection.rs

@@ -45,6 +45,7 @@ prost = { version = "0.11.3", features = ["no-recursion-limit"] }
 rdkafka = { version = "0.29.0", features = ["cmake-build", "ssl-vendored", "libz-static", "zstd"] }
 scopeguard = "1.1.0"
 serde = { version = "1.0.152", features = ["derive"] }
+serde_json = "1.0.89"
 thiserror = "1.0.37"
 timely = { version = "0.12.0", default-features = false, features = ["bincode"] }
 tokio = { version = "1.24.2", features = ["fs", "rt", "sync", "test-util", "time"] }

src/storage-types/Cargo.toml

@@ -122,8 +122,6 @@ pub struct ConnectionContext {
     pub librdkafka_log_level: tracing::Level,
     /// A prefix for an external ID to use for all AWS AssumeRole operations.
     pub aws_external_id_prefix: Option<AwsExternalIdPrefix>,
-    /// The AWS principal for Materialize.
-    pub aws_principal: Option<String>,
     /// A secrets reader.
     pub secrets_reader: Arc<dyn SecretsReader>,
     /// A cloud resource reader, if supported in this configuration.
@@ -143,7 +141,6 @@ impl ConnectionContext {
     pub fn from_cli_args(
         startup_log_level: &CloneableEnvFilter,
         aws_external_id_prefix: Option<AwsExternalIdPrefix>,
-        aws_principal: Option<String>,
         secrets_reader: Arc<dyn SecretsReader>,
         cloud_resource_reader: Option<Arc<dyn CloudResourceReader>>,
     ) -> ConnectionContext {
@@ -153,7 +150,6 @@ impl ConnectionContext {
                 "librdkafka",
             ),
             aws_external_id_prefix,
-            aws_principal,
             secrets_reader,
             cloud_resource_reader,
             ssh_tunnel_manager: SshTunnelManager::default(),
@@ -165,7 +161,6 @@ impl ConnectionContext {
         ConnectionContext {
             librdkafka_log_level: tracing::Level::INFO,
             aws_external_id_prefix: None,
-            aws_principal: None,
             secrets_reader,
             cloud_resource_reader: None,
             ssh_tunnel_manager: SshTunnelManager::default(),

src/storage-types/src/connections.rs

@@ -36,4 +36,6 @@ message ProtoAwsCredentials {
 message ProtoAwsAssumeRole {
     string arn = 1;
     optional string session_name = 2;
+    string mz_principal = 3;
+    string external_id_prefix = 4;
 }