Skip to content

Commit

Permalink
add connection string auth (#319)
Browse files Browse the repository at this point in the history 
* add connection string auth

* removed keywords

* manual execution false

* addressed review comments

* replaced with string utils blank

* fixed javadoc error
  • Loading branch information
@tanmaya-panda1
tanmaya-panda1 authored Oct 19, 2023
1 parent bfa2b6a commit 1398520
Show file tree
Hide file tree
Showing 3 changed files with 186 additions and 10 deletions.
147 changes: 137 additions & 10 deletions data/src/main/java/com/microsoft/azure/kusto/data/auth/ConnectionStringBuilder.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,9 @@

import java.security.PrivateKey;
import java.security.cert.X509Certificate;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.Callable;

public class ConnectionStringBuilder {
Expand All @@ -37,8 +39,42 @@ public class ConnectionStringBuilder {
private String applicationNameForTracing;
private static final String DEFAULT_DEVICE_AUTH_TENANT = "organizations";

private ConnectionStringBuilder(String clusterUrl) {
this.clusterUrl = clusterUrl;
private static final Map<String, ConnectionStringKeyword> stringKeywordMap = new HashMap<>();

static {
stringKeywordMap.put("Data Source".toLowerCase(), ConnectionStringKeyword.DataSource);
stringKeywordMap.put("Addr".toLowerCase(), ConnectionStringKeyword.DataSource);
stringKeywordMap.put("Address".toLowerCase(), ConnectionStringKeyword.DataSource);
stringKeywordMap.put("Network Address".toLowerCase(), ConnectionStringKeyword.DataSource);
stringKeywordMap.put("Server".toLowerCase(), ConnectionStringKeyword.DataSource);

// used for user prompt authentication
stringKeywordMap.put("User ID".toLowerCase(), ConnectionStringKeyword.UsernameHint);
stringKeywordMap.put("UID".toLowerCase(), ConnectionStringKeyword.UsernameHint);
stringKeywordMap.put("User".toLowerCase(), ConnectionStringKeyword.UsernameHint);

stringKeywordMap.put("Authority ID".toLowerCase(), ConnectionStringKeyword.AuthorityId);
stringKeywordMap.put("TenantId".toLowerCase(), ConnectionStringKeyword.AuthorityId);

stringKeywordMap.put("Application Client Id".toLowerCase(), ConnectionStringKeyword.ApplicationClientId);
stringKeywordMap.put("AppClientId".toLowerCase(), ConnectionStringKeyword.ApplicationClientId);

stringKeywordMap.put("Application Key".toLowerCase(), ConnectionStringKeyword.ApplicationKey);
stringKeywordMap.put("AppKey".toLowerCase(), ConnectionStringKeyword.ApplicationKey);

stringKeywordMap.put("User Token".toLowerCase(), ConnectionStringKeyword.UserToken);
stringKeywordMap.put("UserToken".toLowerCase(), ConnectionStringKeyword.UserToken);
stringKeywordMap.put("UsrToken".toLowerCase(), ConnectionStringKeyword.UserToken);

stringKeywordMap.put("Application Name for Tracing".toLowerCase(), ConnectionStringKeyword.ApplicationNameForTracing);

stringKeywordMap.put("User Name for Tracing".toLowerCase(), ConnectionStringKeyword.UserNameForTracing);

stringKeywordMap.put("Client Version for Tracing".toLowerCase(), ConnectionStringKeyword.ClientVersionForTracing);
}

private ConnectionStringBuilder() {
this.clusterUrl = null;
this.usernameHint = null;
this.applicationClientId = null;
this.applicationKey = null;
Expand All @@ -57,6 +93,46 @@ private ConnectionStringBuilder(String clusterUrl) {
this.applicationNameForTracing = null;
}

private void assignValue(String rawKey, String value) {
rawKey = rawKey.trim().toLowerCase();
ConnectionStringKeyword parsedKey = stringKeywordMap.get(rawKey);
if (parsedKey == null) {
throw new IllegalArgumentException("Error: unsupported key " + rawKey + " in connection string");
}
switch (parsedKey) {
case DataSource:
this.clusterUrl = value;
break;
case ApplicationClientId:
this.applicationClientId = value;
break;
case ApplicationKey:
this.applicationKey = value;
break;
case AuthorityId:
this.aadAuthorityId = value;
break;
case ApplicationNameForTracing:
this.applicationNameForTracing = value;
break;
case UserNameForTracing:
this.userNameForTracing = value;
break;
case ClientVersionForTracing:
this.appendedClientVersionForTracing = value;
break;
case UsernameHint:
this.usernameHint = value;
this.useUserPromptAuth = true;
break;
case UserToken:
this.accessToken = value;
break;
default:
throw new IllegalArgumentException("Error: unsupported key " + rawKey + " in connection string");
}
}

public ConnectionStringBuilder(ConnectionStringBuilder other) {
this.clusterUrl = other.clusterUrl;
this.usernameHint = other.usernameHint;
Expand All @@ -77,6 +153,35 @@ public ConnectionStringBuilder(ConnectionStringBuilder other) {
this.applicationNameForTracing = other.applicationNameForTracing;
}

/**
* Creates a ConnectionStringBuilder from a connection string. For more information please look at: https://docs.microsoft.com/azure/data-explorer/kusto/api/connection-strings/kusto
*
* @param connectionString The connection string should be of the format: https://clusterName.location.kusto.windows.net;AAD User ID="[email protected]";Password=P@ssWord
* @throws IllegalArgumentException If the connection string is invalid.
*/
public ConnectionStringBuilder(String connectionString) {
if (StringUtils.isBlank(connectionString)) {
throw new IllegalArgumentException("connectionString cannot be null or empty");
}

String[] connStrArr = connectionString.split(";");
if (!connStrArr[0].contains("=")) {
connStrArr[0] = "Data Source=" + connStrArr[0];
}

for (String kvp : connStrArr) {
kvp = kvp.trim();
if (StringUtils.isEmpty(kvp)) {
continue;
}
String[] kvpArr = kvp.split("=");
String val = kvpArr[1].trim();
if (!StringUtils.isEmpty(val)) {
assignValue(kvpArr[0], val);
}
}
}

public String getClusterUrl() {
return clusterUrl;
}
Expand Down Expand Up @@ -215,7 +320,8 @@ public static ConnectionStringBuilder createWithAadApplicationCredentials(String
throw new IllegalArgumentException("applicationKey cannot be null or empty");
}

ConnectionStringBuilder csb = new ConnectionStringBuilder(clusterUrl);
ConnectionStringBuilder csb = new ConnectionStringBuilder();
csb.clusterUrl = clusterUrl;
csb.applicationClientId = applicationClientId;
csb.applicationKey = applicationKey;
csb.aadAuthorityId = authorityId;
Expand All @@ -235,7 +341,8 @@ public static ConnectionStringBuilder createWithUserPrompt(String clusterUrl, St
throw new IllegalArgumentException("clusterUrl cannot be null or empty");
}

ConnectionStringBuilder csb = new ConnectionStringBuilder(clusterUrl);
ConnectionStringBuilder csb = new ConnectionStringBuilder();
csb.clusterUrl = clusterUrl;
csb.aadAuthorityId = authorityId;
csb.usernameHint = usernameHint;
csb.useUserPromptAuth = true;
Expand All @@ -251,7 +358,8 @@ public static ConnectionStringBuilder createWithDeviceCode(String clusterUrl, St
throw new IllegalArgumentException("clusterUrl cannot be null or empty");
}

ConnectionStringBuilder csb = new ConnectionStringBuilder(clusterUrl);
ConnectionStringBuilder csb = new ConnectionStringBuilder();
csb.clusterUrl = clusterUrl;
csb.aadAuthorityId = authorityId;
csb.useDeviceCodeAuth = true;
return csb;
Expand Down Expand Up @@ -282,7 +390,8 @@ public static ConnectionStringBuilder createWithAadApplicationCertificate(String
throw new IllegalArgumentException("privateKey cannot be null");
}

ConnectionStringBuilder csb = new ConnectionStringBuilder(clusterUrl);
ConnectionStringBuilder csb = new ConnectionStringBuilder();
csb.clusterUrl = clusterUrl;
csb.applicationClientId = applicationClientId;
csb.x509Certificate = x509Certificate;
csb.privateKey = privateKey;
Expand Down Expand Up @@ -315,7 +424,8 @@ public static ConnectionStringBuilder createWithAadApplicationCertificateSubject
throw new IllegalArgumentException("privateKey cannot be null");
}

ConnectionStringBuilder csb = new ConnectionStringBuilder(clusterUrl);
ConnectionStringBuilder csb = new ConnectionStringBuilder();
csb.clusterUrl = clusterUrl;
csb.applicationClientId = applicationClientId;
csb.x509CertificateChain = x509CertificateChain;
csb.privateKey = privateKey;
Expand All @@ -331,7 +441,8 @@ public static ConnectionStringBuilder createWithAadAccessTokenAuthentication(Str
throw new IllegalArgumentException("token cannot be null or empty");
}

ConnectionStringBuilder csb = new ConnectionStringBuilder(clusterUrl);
ConnectionStringBuilder csb = new ConnectionStringBuilder();
csb.clusterUrl = clusterUrl;
csb.accessToken = token;
return csb;
}
Expand All @@ -345,7 +456,8 @@ public static ConnectionStringBuilder createWithAadTokenProviderAuthentication(S
throw new IllegalArgumentException("tokenProviderCallback cannot be null");
}

ConnectionStringBuilder csb = new ConnectionStringBuilder(clusterUrl);
ConnectionStringBuilder csb = new ConnectionStringBuilder();
csb.clusterUrl = clusterUrl;
csb.tokenProvider = tokenProviderCallable;
return csb;
}
Expand All @@ -359,7 +471,8 @@ public static ConnectionStringBuilder createWithAadManagedIdentity(String cluste
throw new IllegalArgumentException("clusterUrl cannot be null or empty");
}

ConnectionStringBuilder csb = new ConnectionStringBuilder(clusterUrl);
ConnectionStringBuilder csb = new ConnectionStringBuilder();
csb.clusterUrl = clusterUrl;
csb.managedIdentityClientId = managedIdentityClientId;
csb.useManagedIdentityAuth = true;
return csb;
Expand All @@ -383,4 +496,18 @@ public void setConnectorDetails(String name, String version, @Nullable String ap
applicationNameForTracing = clientDetails.getApplicationForTracing();
userNameForTracing = clientDetails.getUserNameForTracing();
}

enum ConnectionStringKeyword {
DataSource,
UsernameHint,
UseDeviceCodeAuth,
TokenProvider,
ApplicationClientId,
ApplicationKey,
UserToken,
AuthorityId,
ApplicationNameForTracing,
UserNameForTracing,
ClientVersionForTracing
}
}
34 changes: 34 additions & 0 deletions data/src/test/java/com/microsoft/azure/kusto/data/ConnectionStringBuilderTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -115,6 +115,40 @@ void createWithAadAccessTokenAuthentication() {
.createWithAadAccessTokenAuthentication("resource.uri", "token"));
}

@Test
@DisplayName("validate createWithConnectionStringAuthentication which throws IllegalArgumentException exception when missing or invalid parameters")
void createWithConnectionStringAuthentication() {
// An empty connection string should throw an exception
Assertions.assertThrows(IllegalArgumentException.class,
() -> new ConnectionStringBuilder(""));
// A connection string with only spaces should throw an exception
Assertions.assertThrows(IllegalArgumentException.class,
() -> new ConnectionStringBuilder(" "));
// InvalidKey is not a valid connection string parameter
Assertions.assertThrows(IllegalArgumentException.class,
() -> new ConnectionStringBuilder("Data Source=mycluster.kusto.windows.net;AppClientId=myclientid;AppKey=myappkey;InvalidKey=invalidKey"));
}

@Test
@DisplayName("validate createWithConnectionStringAuthentication order in connection string is not important")
void testUnexpectedOrderWithConnectionString() {
String connectionString = "AppClientId=myclientid;Data Source=mycluster.kusto.windows.net;AppKey=myappkey;";
ConnectionStringBuilder builder = new ConnectionStringBuilder(connectionString);

// Assert that the fields have been set correctly
Assertions.assertEquals("mycluster.kusto.windows.net", builder.getClusterUrl());
}

@Test
@DisplayName("validate createWithConnectionStringAuthentication with valid connection string")
void createValidWithConnectionStringAuthentication() {
String connectionString = "Data Source=mycluster.kusto.windows.net;AppClientId=myclientid;AppKey=myappkey;";
ConnectionStringBuilder builder = new ConnectionStringBuilder(connectionString);

// Assert that the fields have been set correctly
Assertions.assertEquals("mycluster.kusto.windows.net", builder.getClusterUrl());
}

@Test
@DisplayName("validate createWithAadTokenProviderAuthentication throws IllegalArgumentException exception when missing or invalid parameters")
void createWithAadTokenProviderAuthentication() {
Expand Down
15 changes: 15 additions & 0 deletions ingest/src/test/java/com/microsoft/azure/kusto/ingest/E2ETest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -459,6 +459,14 @@ void testCreateWithUserPrompt() {
assertTrue(canAuthenticate(engineCsb));
}

@Test
void testCreateWithConnectionStringAndUserPrompt() {
Assumptions.assumeTrue(IsManualExecution());
ConnectionStringBuilder engineCsb = new ConnectionStringBuilder(
"Data Source=" + System.getenv("ENGINE_CONNECTION_STRING") + ";User ID=" + System.getenv("USERNAME_HINT"));
assertTrue(canAuthenticate(engineCsb));
}

@Test
void testCreateWithDeviceAuthentication() {
Assumptions.assumeTrue(IsManualExecution());
Expand All @@ -484,6 +492,13 @@ void testCreateWithAadApplicationCredentials() {
assertTrue(canAuthenticate(engineCsb));
}

@Test
void testCreateWithConnectionStringAndAadApplicationCredentials() {
ConnectionStringBuilder engineCsb = new ConnectionStringBuilder(
"Data Source=" + System.getenv("ENGINE_CONNECTION_STRING") + ";AppClientId=" + appId + ";AppKey=" + appKey + ";Authority ID=" + tenantId);
assertTrue(canAuthenticate(engineCsb));
}

@Test
void testCreateWithAadAccessTokenAuthentication() {
Assumptions.assumeTrue(StringUtils.isNotBlank(System.getenv("TOKEN")));
Expand Down

0 comments on commit 1398520

Please sign in to comment.