Merge branch 'master' into low-allocations

Azure · Aug 21, 2024 · a34584d · a34584d
2 parents ccfb9d5 + ff919ae
commit a34584d
Show file tree

Hide file tree

Showing 3 changed files with 61 additions and 11 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -14,6 +14,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   - `WithAppCertificatePath` - Receives the path to the certificate file.
   - `WithAppCertificateBytes` - Receives the certificate bytes in-memory.  
   Both methods accept an optional password for the certificate.
+- `WithUserManagedIdentity` has been deprecated in favor of more specific functions:
+  - `WithUserAssignedIdentityClientId` - Receives the MSI client id
+  - `WithUserAssignedIdentityResourceId` - Receives the MSI resource id
 
 ### Fixed
 - Fixed Mapping Kind not working correctly with certain formats.

diff --git a/azkustodata/kcsb.go b/azkustodata/kcsb.go
@@ -28,14 +28,17 @@ type ConnectionStringBuilder struct {
 	MsiAuthentication              bool
 	WorkloadAuthentication         bool
 	FederationTokenFilePath        string
-	ManagedServiceIdentity         string
-	InteractiveLogin               bool
-	RedirectURL                    string
-	DefaultAuth                    bool
-	ClientOptions                  *azcore.ClientOptions
-	ApplicationForTracing          string
-	UserForTracing                 string
-	TokenCredential                azcore.TokenCredential
+	// Deprecated: Use ManagedServiceIdentityClientId or ManagedServiceIdentityResourceId instead
+	ManagedServiceIdentity           string
+	ManagedServiceIdentityClientId   string
+	ManagedServiceIdentityResourceId string
+	InteractiveLogin                 bool
+	RedirectURL                      string
+	DefaultAuth                      bool
+	ClientOptions                    *azcore.ClientOptions
+	ApplicationForTracing            string
+	UserForTracing                   string
+	TokenCredential                  azcore.TokenCredential
 }
 
 const (
@@ -164,6 +167,8 @@ func (kcsb *ConnectionStringBuilder) resetConnectionString() {
 	kcsb.MsiAuthentication = false
 	kcsb.WorkloadAuthentication = false
 	kcsb.ManagedServiceIdentity = ""
+	kcsb.ManagedServiceIdentityClientId = ""
+	kcsb.ManagedServiceIdentityResourceId = ""
 	kcsb.InteractiveLogin = false
 	kcsb.RedirectURL = ""
 	kcsb.ClientOptions = nil
@@ -254,13 +259,30 @@ func (kcsb *ConnectionStringBuilder) WithAzCli() *ConnectionStringBuilder {
 	return kcsb
 }
 
+// Deprecated: use WithUserManagedIdentityClientId or WithUserManagedIdentityResourceId instead
 // WithUserManagedIdentity Creates a Kusto Connection string builder that will authenticate with AAD application, using
 // an application token obtained from a Microsoft Service Identity endpoint using user assigned id.
 func (kcsb *ConnectionStringBuilder) WithUserManagedIdentity(clientID string) *ConnectionStringBuilder {
+	return kcsb.WithUserAssignedIdentityClientId(clientID)
+}
+
+// WithUserAssignedIdentityClientId Creates a Kusto Connection string builder that will authenticate with AAD application, using
+// an application token obtained from a Microsoft Service Identity endpoint using user assigned id.
+func (kcsb *ConnectionStringBuilder) WithUserAssignedIdentityClientId(clientID string) *ConnectionStringBuilder {
+	requireNonEmpty(dataSource, kcsb.DataSource)
+	kcsb.resetConnectionString()
+	kcsb.MsiAuthentication = true
+	kcsb.ManagedServiceIdentityClientId = clientID
+	return kcsb
+}
+
+// WithUserAssignedIdentityResourceId Creates a Kusto Connection string builder that will authenticate with AAD application, using
+// an application token obtained from a Microsoft Service Identity endpoint using an MSI's resourceID.
+func (kcsb *ConnectionStringBuilder) WithUserAssignedIdentityResourceId(resourceID string) *ConnectionStringBuilder {
 	requireNonEmpty(dataSource, kcsb.DataSource)
 	kcsb.resetConnectionString()
 	kcsb.MsiAuthentication = true
-	kcsb.ManagedServiceIdentity = clientID
+	kcsb.ManagedServiceIdentityResourceId = resourceID
 	return kcsb
 }
 
@@ -410,8 +432,17 @@ func (kcsb *ConnectionStringBuilder) newTokenProvider() (*TokenProvider, error)
 	case kcsb.MsiAuthentication:
 		init = func(ci *CloudInfo, cliOpts *azcore.ClientOptions, appClientId string) (azcore.TokenCredential, error) {
 			opts := &azidentity.ManagedIdentityCredentialOptions{ClientOptions: *cliOpts}
+			// legacy kcsb.ManagedServiceIdentity field takes precedence over
+			// new kcsb.ManagedServiceIdentityClientId field which takes precedence over
+			// new kcsb.ManagedServiceIdentityResourceId field
+			// if no client id is provided, the logic falls back to set up
+			// the system assigned identity
 			if !isEmpty(kcsb.ManagedServiceIdentity) {
 				opts.ID = azidentity.ClientID(kcsb.ManagedServiceIdentity)
+			} else if !isEmpty(kcsb.ManagedServiceIdentityClientId) {
+				opts.ID = azidentity.ClientID(kcsb.ManagedServiceIdentityClientId)
+			} else if !isEmpty(kcsb.ManagedServiceIdentityResourceId) {
+				opts.ID = azidentity.ResourceID(kcsb.ManagedServiceIdentityResourceId)
 			}
 
 			cred, err := azidentity.NewManagedIdentityCredential(opts)

diff --git a/azkustodata/kcsb_test.go b/azkustodata/kcsb_test.go
@@ -94,7 +94,7 @@ func TestWithAadUserPassAuthErr(t *testing.T) {
 
 }
 
-func TestWitAadUserToken(t *testing.T) {
+func TestWithAadUserToken(t *testing.T) {
 	want := ConnectionStringBuilder{
 		DataSource: "endpoint",
 		UserToken:  "token",
@@ -120,7 +120,7 @@ func TestWithWorkloadIdentity(t *testing.T) {
 	assert.EqualValues(t, want, *actual)
 }
 
-func TestWitAadUserTokenErr(t *testing.T) {
+func TestWithAadUserTokenErr(t *testing.T) {
 	defer func() {
 		if res := recover(); res == nil {
 			t.Errorf("Should have panic")
@@ -173,6 +173,22 @@ func TestGetTokenProviderHappy(t *testing.T) {
 				MsiAuthentication:      true,
 				ClientOptions:          &azcore.ClientOptions{},
 			},
+		}, {
+			name: "test_tokenprovider_managedui_clientId",
+			kcsb: ConnectionStringBuilder{
+				DataSource:                     "https://endpoint/test_tokenprovider_managedui_clientID",
+				ManagedServiceIdentityClientId: "00000000-0000-0000-0000-000000000000",
+				MsiAuthentication:              true,
+				ClientOptions:                  &azcore.ClientOptions{},
+			},
+		}, {
+			name: "test_tokenprovider_managedui_resourceId",
+			kcsb: ConnectionStringBuilder{
+				DataSource:                       "https://endpoint/test_tokenprovider_managedui_resourceID",
+				ManagedServiceIdentityResourceId: "/subscriptions/00000000-0000-0000-0000-000000000000/resourceGroups/testResourceGroup/providers/Microsoft.ManagedIdentity/userAssignedIdentities/testIdentity",
+				MsiAuthentication:                true,
+				ClientOptions:                    &azcore.ClientOptions{},
+			},
 		}, {
 			name: "test_tokenprovider_managedidauth2",
 			kcsb: ConnectionStringBuilder{