Add Azure Token Base Authentication (#18387)

* Add azure token based auth

* Add Azure Token Base Authentication

* Update azure-auth.md

* Update azure-auth.md

* feat: Add `azure-identity-extensions` library for passwordless database connection

---------

Co-authored-by: Ayush Shah <ayush@getcollate.io>
Co-authored-by: Akash-Jain <15995028+akash-jain-10@users.noreply.github.com>

openmetadata-docs/content/v1.5.x/deployment/azure-auth.md

@@ -0,0 +1,29 @@
+---
+title: How to enable Azure Auth
+slug: /deployment/azure-auth
+collate: false
+---
+
+# AZURE resources on Postgres/MySQL Auth
+https://learn.microsoft.com/en-us/azure/postgresql/flexible-server/concepts-extensions#how-to-use-postgresql-extensions
+# Requirements
+
+1. Azure Postgres or MySQL Cluster with auth enabled
+2. User on DB Cluster with authentication enabled
+
+# How to enable Azure Auth on postgresql
+
+Set the environment variables
+
+```Commandline
+  DB_PARAMS="azure=true&allowPublicKeyRetrieval=true&sslmode=require&serverTimezone=UTC" 
+  DB_USER_PASSWORD=none
+```
+
+Either through helm (if deployed in kubernetes) or as env vars.
+
+{% note %}
+
+The `DB_USER_PASSWORD` is still required and cannot be empty. Set it to a random/dummy string.
+
+{% /note %} 

openmetadata-docs/content/v1.5.x/menu.md

@@ -185,6 +185,8 @@ site_menu:
 
   - category: Deployment / How to enable AWS RDS IAM Auth
     url: /deployment/rds-iam-auth
+  - category: Deployment / How to enable Azure Database Auth
+    url: /deployment/azure-auth
   - category: Deployment / Server Configuration Reference
     url: /deployment/configuration
   - category: Deployment / Database Connection Pooling

openmetadata-docs/content/v1.6.x-SNAPSHOT/deployment/azure-auth.md

@@ -0,0 +1,29 @@
+---
+title: How to enable Azure Auth
+slug: /deployment/azure-auth
+collate: false
+---
+
+# AZURE resources on Postgres/MySQL Auth
+https://learn.microsoft.com/en-us/azure/postgresql/flexible-server/concepts-extensions#how-to-use-postgresql-extensions
+# Requirements
+
+1. Azure Postgres or MySQL Cluster with auth enabled
+2. User on DB Cluster with authentication enabled
+
+# How to enable Azure Auth on postgresql
+
+Set the environment variables
+
+```Commandline
+  DB_PARAMS="azure=true&allowPublicKeyRetrieval=true&sslmode=require&serverTimezone=UTC" 
+  DB_USER_PASSWORD=none
+```
+
+Either through helm (if deployed in kubernetes) or as env vars.
+
+{% note %}
+
+The `DB_USER_PASSWORD` is still required and cannot be empty. Set it to a random/dummy string.
+
+{% /note %} 

openmetadata-docs/content/v1.6.x-SNAPSHOT/menu.md

@@ -185,6 +185,8 @@ site_menu:
 
   - category: Deployment / How to enable AWS RDS IAM Auth
     url: /deployment/rds-iam-auth
+  - category: Deployment / How to enable Azure Auth
+    url: /deployment/azure-auth
   - category: Deployment / Server Configuration Reference
     url: /deployment/configuration
   - category: Deployment / Database Connection Pooling

openmetadata-service/pom.xml

@@ -17,8 +17,9 @@
     <sonar.tests>${project.basedir}/src/test/java</sonar.tests>
     <org.testcontainers.version>1.20.1</org.testcontainers.version>
     <awssdk.version>2.27.17</awssdk.version>
-    <azure-identity.version>1.13.2</azure-identity.version>
+    <azure-identity.version>1.14.0</azure-identity.version>
     <azure-kv.version>4.8.6</azure-kv.version>
+    <azure-identity-extensions.version>1.0.0</azure-identity-extensions.version>
     <expiring.map.version>0.5.11</expiring.map.version>
     <java.saml>2.9.0</java.saml>
     <xmlsec.version>2.3.4</xmlsec.version>
@@ -290,6 +291,11 @@
       <artifactId>azure-identity</artifactId>
       <version>${azure-identity.version}</version>
     </dependency>
+    <dependency>
+      <groupId>com.azure</groupId>
+      <artifactId>azure-identity-extensions</artifactId>
+      <version>${azure-identity-extensions.version}</version>
+  </dependency>
     <dependency>
       <groupId>io.dropwizard.modules</groupId>
       <artifactId>dropwizard-web</artifactId>
@@ -639,6 +645,10 @@
       <artifactId>jakarta.activation</artifactId>
       <version>2.0.1</version>
     </dependency>
+    <dependency>
+      <groupId>com.microsoft.azure</groupId>
+      <artifactId>msal4j</artifactId>
+    </dependency>
   </dependencies>
 
   <profiles>

openmetadata-service/src/main/java/org/openmetadata/service/OpenMetadataApplication.java

@@ -13,13 +13,13 @@
 
 package org.openmetadata.service;
 
+import static org.openmetadata.service.util.jdbi.JdbiUtils.createAndSetupJDBI;
+
 import io.dropwizard.Application;
 import io.dropwizard.configuration.EnvironmentVariableSubstitutor;
 import io.dropwizard.configuration.SubstitutingSourceProvider;
-import io.dropwizard.db.DataSourceFactory;
 import io.dropwizard.health.conf.HealthConfiguration;
 import io.dropwizard.health.core.HealthCheckBundle;
-import io.dropwizard.jdbi3.JdbiFactory;
 import io.dropwizard.jersey.errors.EarlyEofExceptionMapper;
 import io.dropwizard.jersey.errors.LoggingExceptionMapper;
 import io.dropwizard.jersey.jackson.JsonProcessingExceptionMapper;
@@ -59,7 +59,6 @@ import org.eclipse.jetty.websocket.server.WebSocketUpgradeFilter;
 import org.glassfish.jersey.media.multipart.MultiPartFeature;
 import org.glassfish.jersey.server.ServerProperties;
 import org.jdbi.v3.core.Jdbi;
-import org.jdbi.v3.core.statement.SqlStatements;
 import org.jdbi.v3.sqlobject.SqlObjects;
 import org.openmetadata.schema.api.security.AuthenticationConfiguration;
 import org.openmetadata.schema.api.security.AuthorizerConfiguration;
@@ -126,8 +125,6 @@ import org.openmetadata.service.socket.SocketAddressFilter;
 import org.openmetadata.service.socket.WebSocketManager;
 import org.openmetadata.service.util.MicrometerBundleSingleton;
 import org.openmetadata.service.util.incidentSeverityClassifier.IncidentSeverityClassifierInterface;
-import org.openmetadata.service.util.jdbi.DatabaseAuthenticationProviderFactory;
-import org.openmetadata.service.util.jdbi.OMSqlLogger;
 import org.pac4j.core.util.CommonHelper;
 import org.quartz.SchedulerException;
 
@@ -389,28 +386,6 @@ public class OpenMetadataApplication extends Application<OpenMetadataApplication
     }
   }
 
-  private Jdbi createAndSetupJDBI(Environment environment, DataSourceFactory dbFactory) {
-    // Check for db auth providers.
-    DatabaseAuthenticationProviderFactory.get(dbFactory.getUrl())
-        .ifPresent(
-            databaseAuthenticationProvider -> {
-              String token =
-                  databaseAuthenticationProvider.authenticate(
-                      dbFactory.getUrl(), dbFactory.getUser(), dbFactory.getPassword());
-              dbFactory.setPassword(token);
-            });
-
-    Jdbi jdbiInstance = new JdbiFactory().build(environment, dbFactory, "database");
-    jdbiInstance.setSqlLogger(new OMSqlLogger());
-    // Set the Database type for choosing correct queries from annotations
-    jdbiInstance
-        .getConfig(SqlObjects.class)
-        .setSqlLocator(new ConnectionAwareAnnotationSqlLocator(dbFactory.getDriverClass()));
-    jdbiInstance.getConfig(SqlStatements.class).setUnusedBindingAllowed(true);
-
-    return jdbiInstance;
-  }
-
   @SneakyThrows
   @Override
   public void initialize(Bootstrap<OpenMetadataApplicationConfig> bootstrap) {

openmetadata-service/src/main/java/org/openmetadata/service/util/AzureTokenProvider.java

@@ -0,0 +1,28 @@
+package org.openmetadata.service.util;
+
+import com.microsoft.aad.msal4j.*;
+import java.net.MalformedURLException;
+import java.util.Set;
+
+public class AzureTokenProvider {
+
+  private static final String CLIENT_ID = "your-client-id"; // From Azure AD App Registration
+  private static final String TENANT_ID = "your-tenant-id"; // Your Azure AD tenant ID
+  private static final String CLIENT_SECRET = "your-client-secret"; // Generated in App Registration
+  private static final String SCOPE =
+      "https://ossrdbms-aad.database.windows.net/.default"; // Scope for PostgreSQL
+
+  public static String getAccessToken() throws MalformedURLException {
+    ConfidentialClientApplication app =
+        ConfidentialClientApplication.builder(
+                CLIENT_ID, ClientCredentialFactory.createFromSecret(CLIENT_SECRET))
+            .authority("https://login.microsoftonline.com/" + TENANT_ID) // Azure AD authority
+            .build();
+
+    Set<String> scopes = Set.of(SCOPE);
+    ClientCredentialParameters parameters = ClientCredentialParameters.builder(scopes).build();
+    IAuthenticationResult result = app.acquireToken(parameters).join(); // Get the token
+
+    return result.accessToken(); // Return the access token
+  }
+}

openmetadata-service/src/main/java/org/openmetadata/service/util/OpenMetadataOperations.java

@@ -8,7 +8,6 @@ import static org.openmetadata.service.util.AsciiTable.printOpenMetadataText;
 
 import ch.qos.logback.classic.Level;
 import ch.qos.logback.classic.Logger;
-import com.codahale.metrics.NoopMetricRegistry;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.google.gson.Gson;
 import com.google.gson.JsonElement;
@@ -37,8 +36,6 @@ import lombok.extern.slf4j.Slf4j;
 import org.flywaydb.core.Flyway;
 import org.flywaydb.core.api.MigrationVersion;
 import org.jdbi.v3.core.Jdbi;
-import org.jdbi.v3.sqlobject.SqlObjectPlugin;
-import org.jdbi.v3.sqlobject.SqlObjects;
 import org.openmetadata.schema.EntityInterface;
 import org.openmetadata.schema.ServiceEntityInterface;
 import org.openmetadata.schema.entity.app.App;
@@ -63,7 +60,6 @@ import org.openmetadata.service.jdbi3.IngestionPipelineRepository;
 import org.openmetadata.service.jdbi3.ListFilter;
 import org.openmetadata.service.jdbi3.MigrationDAO;
 import org.openmetadata.service.jdbi3.SystemRepository;
-import org.openmetadata.service.jdbi3.locator.ConnectionAwareAnnotationSqlLocator;
 import org.openmetadata.service.jdbi3.locator.ConnectionType;
 import org.openmetadata.service.migration.api.MigrationWorkflow;
 import org.openmetadata.service.resources.CollectionRegistry;
@@ -73,6 +69,7 @@ import org.openmetadata.service.secrets.SecretsManager;
 import org.openmetadata.service.secrets.SecretsManagerFactory;
 import org.openmetadata.service.secrets.SecretsManagerUpdateService;
 import org.openmetadata.service.util.jdbi.DatabaseAuthenticationProviderFactory;
+import org.openmetadata.service.util.jdbi.JdbiUtils;
 import org.slf4j.LoggerFactory;
 import picocli.CommandLine;
 import picocli.CommandLine.Command;
@@ -211,6 +208,7 @@ public class OpenMetadataOperations implements Callable<Integer> {
       flyway.clean();
       LOG.info("Creating the OpenMetadata Schema.");
       flyway.migrate();
+      LOG.info("Running the Native Migrations.");
       validateAndRunSystemDataMigrations(true);
       LOG.info("OpenMetadata Database Schema is Updated.");
       LOG.info("create indexes.");
@@ -544,6 +542,9 @@ public class OpenMetadataOperations implements Callable<Integer> {
     String jdbcUrl = dataSourceFactory.getUrl();
     String user = dataSourceFactory.getUser();
     String password = dataSourceFactory.getPassword();
+    LOG.info("JDBC URL: {}", jdbcUrl);
+    LOG.info("User: {}", user);
+    LOG.info("Password: {}", password);
     assert user != null && password != null;
 
     String flywayRootPath = config.getMigrationConfiguration().getFlywayPath();
@@ -568,12 +569,8 @@ public class OpenMetadataOperations implements Callable<Integer> {
             .load();
     nativeSQLScriptRootPath = config.getMigrationConfiguration().getNativePath();
     extensionSQLScriptRootPath = config.getMigrationConfiguration().getExtensionPath();
-    jdbi = Jdbi.create(dataSourceFactory.build(new NoopMetricRegistry(), "open-metadata-ops"));
-    jdbi.installPlugin(new SqlObjectPlugin());
-    jdbi.getConfig(SqlObjects.class)
-        .setSqlLocator(
-            new ConnectionAwareAnnotationSqlLocator(
-                config.getDataSourceFactory().getDriverClass()));
+
+    jdbi = JdbiUtils.createAndSetupJDBI(dataSourceFactory);
 
     searchRepository = new SearchRepository(config.getElasticSearchConfiguration());
 

openmetadata-service/src/main/java/org/openmetadata/service/util/jdbi/AwsRdsDatabaseAuthenticationProvider.java

@@ -2,13 +2,8 @@ package org.openmetadata.service.util.jdbi;
 
 import java.net.MalformedURLException;
 import java.net.URI;
-import java.net.URL;
-import java.net.URLDecoder;
-import java.nio.charset.StandardCharsets;
-import java.util.LinkedHashMap;
 import java.util.Map;
 import java.util.Objects;
-import org.jetbrains.annotations.NotNull;
 import software.amazon.awssdk.auth.credentials.DefaultCredentialsProvider;
 import software.amazon.awssdk.regions.Region;
 import software.amazon.awssdk.services.rds.RdsUtilities;
@@ -27,9 +22,8 @@ public class AwsRdsDatabaseAuthenticationProvider implements DatabaseAuthenticat
 
   @Override
   public String authenticate(String jdbcUrl, String username, String password) {
-    // !!
     try {
-      // Prepare
+
       URI uri = URI.create(PROTOCOL + removeProtocolFrom(jdbcUrl));
       Map<String, String> queryParams = parseQueryParams(uri.toURL());
 
@@ -63,27 +57,4 @@ public class AwsRdsDatabaseAuthenticationProvider implements DatabaseAuthenticat
       throw new DatabaseAuthenticationProviderException(e);
     }
   }
-
-  @NotNull
-  private static String removeProtocolFrom(String jdbcUrl) {
-    return jdbcUrl.substring(jdbcUrl.indexOf("://") + 3);
-  }
-
-  private Map<String, String> parseQueryParams(URL url) {
-    // Prepare
-    Map<String, String> queryPairs = new LinkedHashMap<>();
-    String query = url.getQuery();
-    String[] pairs = query.split("&");
-
-    // Loop
-    for (String pair : pairs) {
-      int idx = pair.indexOf("=");
-      // Add
-      queryPairs.put(
-          URLDecoder.decode(pair.substring(0, idx), StandardCharsets.UTF_8),
-          URLDecoder.decode(pair.substring(idx + 1), StandardCharsets.UTF_8));
-    }
-    // Return
-    return queryPairs;
-  }
 }

openmetadata-service/src/main/java/org/openmetadata/service/util/jdbi/AzureDatabaseAuthenticationProvider.java

@@ -0,0 +1,36 @@
+package org.openmetadata.service.util.jdbi;
+
+import com.azure.core.credential.AccessToken;
+import com.azure.core.credential.TokenRequestContext;
+import com.azure.identity.DefaultAzureCredential;
+import com.azure.identity.DefaultAzureCredentialBuilder;
+
+public class AzureDatabaseAuthenticationProvider implements DatabaseAuthenticationProvider {
+  public static final String AZURE = "azure";
+
+  @Override
+  public String authenticate(String jdbcUrl, String username, String password) {
+    try {
+      return fetchAzureADToken();
+    } catch (Exception e) {
+      throw new DatabaseAuthenticationProviderException(e);
+    }
+  }
+
+  private String fetchAzureADToken() {
+    try {
+      DefaultAzureCredential defaultCredential = new DefaultAzureCredentialBuilder().build();
+      TokenRequestContext requestContext =
+          new TokenRequestContext().addScopes("https://ossrdbms-aad.database.windows.net/.default");
+      AccessToken token = defaultCredential.getToken(requestContext).block();
+
+      if (token != null) {
+        return token.getToken();
+      } else {
+        throw new DatabaseAuthenticationProviderException("Failed to fetch token");
+      }
+    } catch (Exception e) {
+      throw new DatabaseAuthenticationProviderException("Error fetching Azure AD token", e);
+    }
+  }
+}

openmetadata-service/src/main/java/org/openmetadata/service/util/jdbi/DatabaseAuthenticationProvider.java

@@ -1,5 +1,12 @@
 package org.openmetadata.service.util.jdbi;
 
+import java.net.URL;
+import java.net.URLDecoder;
+import java.nio.charset.StandardCharsets;
+import java.util.LinkedHashMap;
+import java.util.Map;
+import javax.validation.constraints.NotNull;
+
 /**
  * Database authentication provider is the main interface responsible for all implementation that requires additional
  * authentication steps required by the database in order to authorize a user to be able to operate on it.
@@ -15,4 +22,27 @@ public interface DatabaseAuthenticationProvider {
    * @return authorization token
    */
   String authenticate(String jdbcUrl, String username, String password);
+
+  @NotNull
+  default String removeProtocolFrom(String jdbcUrl) {
+    return jdbcUrl.substring(jdbcUrl.indexOf("://") + 3);
+  }
+
+  default Map<String, String> parseQueryParams(URL url) {
+    // Prepare
+    Map<String, String> queryPairs = new LinkedHashMap<>();
+    String query = url.getQuery();
+    String[] pairs = query.split("&");
+
+    // Loop
+    for (String pair : pairs) {
+      int idx = pair.indexOf("=");
+      // Add
+      queryPairs.put(
+          URLDecoder.decode(pair.substring(0, idx), StandardCharsets.UTF_8),
+          URLDecoder.decode(pair.substring(idx + 1), StandardCharsets.UTF_8));
+    }
+    // Return
+    return queryPairs;
+  }
 }

openmetadata-service/src/main/java/org/openmetadata/service/util/jdbi/DatabaseAuthenticationProviderFactory.java

@@ -17,8 +17,9 @@ public class DatabaseAuthenticationProviderFactory {
     // Check
     if (jdbcURL.contains(AwsRdsDatabaseAuthenticationProvider.AWS_REGION)
         && jdbcURL.contains(AwsRdsDatabaseAuthenticationProvider.ALLOW_PUBLIC_KEY_RETRIEVAL)) {
-      // Return AWS RDS Auth provider
       return Optional.of(new AwsRdsDatabaseAuthenticationProvider());
+    } else if (jdbcURL.contains(AzureDatabaseAuthenticationProvider.AZURE)) {
+      return Optional.of(new AzureDatabaseAuthenticationProvider());
     }
 
     // Return empty

openmetadata-service/src/main/java/org/openmetadata/service/util/jdbi/JdbiUtils.java

@@ -0,0 +1,55 @@
+package org.openmetadata.service.util.jdbi;
+
+import com.codahale.metrics.NoopMetricRegistry;
+import io.dropwizard.db.DataSourceFactory;
+import io.dropwizard.jdbi3.JdbiFactory;
+import io.dropwizard.setup.Environment;
+import org.jdbi.v3.core.Jdbi;
+import org.jdbi.v3.core.statement.SqlStatements;
+import org.jdbi.v3.sqlobject.SqlObjectPlugin;
+import org.jdbi.v3.sqlobject.SqlObjects;
+import org.openmetadata.service.jdbi3.locator.ConnectionAwareAnnotationSqlLocator;
+
+public class JdbiUtils {
+
+  public static Jdbi createAndSetupJDBI(Environment environment, DataSourceFactory dbFactory) {
+    DatabaseAuthenticationProviderFactory.get(dbFactory.getUrl())
+        .ifPresent(
+            databaseAuthenticationProvider -> {
+              String token =
+                  databaseAuthenticationProvider.authenticate(
+                      dbFactory.getUrl(), dbFactory.getUser(), dbFactory.getPassword());
+              dbFactory.setPassword(token);
+            });
+
+    Jdbi jdbiInstance = new JdbiFactory().build(environment, dbFactory, "database");
+    jdbiInstance.setSqlLogger(new OMSqlLogger());
+    // Set the Database type for choosing correct queries from annotations
+    jdbiInstance
+        .getConfig(SqlObjects.class)
+        .setSqlLocator(new ConnectionAwareAnnotationSqlLocator(dbFactory.getDriverClass()));
+    jdbiInstance.getConfig(SqlStatements.class).setUnusedBindingAllowed(true);
+
+    return jdbiInstance;
+  }
+
+  public static Jdbi createAndSetupJDBI(DataSourceFactory dbFactory) {
+    DatabaseAuthenticationProviderFactory.get(dbFactory.getUrl())
+        .ifPresent(
+            databaseAuthenticationProvider -> {
+              String token =
+                  databaseAuthenticationProvider.authenticate(
+                      dbFactory.getUrl(), dbFactory.getUser(), dbFactory.getPassword());
+              dbFactory.setPassword(token);
+            });
+
+    Jdbi jdbiInstance = Jdbi.create(dbFactory.build(new NoopMetricRegistry(), "open-metadata-ops"));
+    jdbiInstance.installPlugin(new SqlObjectPlugin());
+    jdbiInstance
+        .getConfig(SqlObjects.class)
+        .setSqlLocator(new ConnectionAwareAnnotationSqlLocator(dbFactory.getDriverClass()));
+    jdbiInstance.getConfig(SqlStatements.class).setUnusedBindingAllowed(true);
+
+    return jdbiInstance;
+  }
+}

pom.xml

@@ -551,6 +551,17 @@
         <version>${jsonpath.version}</version>
       </dependency>
 
+      <dependency>
+        <groupId>com.microsoft.azure</groupId>
+        <artifactId>msal4j</artifactId>
+        <version>1.17.2</version>
+      </dependency>
+      <dependency>
+        <groupId>com.azure</groupId>
+        <artifactId>azure-identity</artifactId>
+        <version>1.14.0</version>
+      </dependency>
+
       <!-- avoid security issue https://security.snyk.io/vuln/SNYK-JAVA-ORGYAML-2806360 -->
     </dependencies>
   </dependencyManagement>