Repository: linkedin/datahub-gma
Pull request #505: [feat] Smart schema evolution
Author: login: zhixuanjia
Created: 2025-03-17T23:10:49Z | Merged: 2025-03-19T20:53:42Z

## Summary
DDL operations on large table can be dangerous. For example, adding an index to a table contains a few millions of rows will likely take hour or so to complete and cause performance impact to database.

Instead of blindly applying the evolution scripts, this PR aims to classify the SQL script into low and high risk category.

- For low risk ones, flyway will automatically apply the script to maximize productivity.
- For high risk ones, we will work MySQL team to get it applied properly.

The low risk DDL are
- Create new table
- Add new column to existing table

The high risk DDL are
- Create index on table with more than 500K estimated row count.
- Drop table

## Other work needed
For high risk script, we will need to manually mark it as completed in version table after MySQL team applied the script.

## Testing Done
unit test

## Checklist

- [ ] The PR conforms to DataHub's [Contributing Guideline](https://github.com/linkedin/datahub/blob/master/docs/CONTRIBUTING.md) (particularly [Commit Message Format](https://github.com/linkedin/datahub/blob/master/docs/CONTRIBUTING.md#commit-message-format))
- [ ] Links to related issues (if applicable)
- [ ] Docs related to the changes have been added/updated (if applicable)



Diff:
diff --git a/build.gradle b/build.gradle
index 8efa37917..236e59743 100644
--- a/build.gradle
+++ b/build.gradle
@@ -50,6 +50,7 @@ project.ext.externalDependency = [
     'jacksonDataBind': 'com.fasterxml.jackson.core:jackson-databind:2.17.2',
     'javatuples': 'org.javatuples:javatuples:1.2',
     'jsonSimple': 'com.googlecode.json-simple:json-simple:1.1.1',
+    'jsql': 'com.github.jsqlparser:jsqlparser:4.6',
     'junitJupiterApi': "org.junit.jupiter:junit-jupiter-api:$junitJupiterVersion",
     'junitJupiterParams': "org.junit.jupiter:junit-jupiter-params:$junitJupiterVersion",
     'junitJupiterEngine': "org.junit.jupiter:junit-jupiter-engine:$junitJupiterVersion",
diff --git a/dao-impl/ebean-dao/build.gradle b/dao-impl/ebean-dao/build.gradle
index 633965677..e6643c871 100644
--- a/dao-impl/ebean-dao/build.gradle
+++ b/dao-impl/ebean-dao/build.gradle
@@ -16,6 +16,7 @@ dependencies {
   compile externalDependency.jsonSimple
   compile externalDependency.log4j
 
+  implementation 'com.github.jsqlparser:jsqlparser:4.6'
   compileOnly externalDependency.ebeanAgent
   compileOnly externalDependency.lombok
 
diff --git a/dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/EbeanLocalAccess.java b/dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/EbeanLocalAccess.java
index dd607cf9f..253e77ba9 100644
--- a/dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/EbeanLocalAccess.java
+++ b/dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/EbeanLocalAccess.java
@@ -81,7 +81,7 @@ public EbeanLocalAccess(EbeanServer server, ServerConfig serverConfig, @Nonnull
     _urnClass = urnClass;
     _urnPathExtractor = urnPathExtractor;
     _entityType = ModelUtils.getEntityTypeFromUrnClass(_urnClass);
-    _schemaEvolutionManager = createSchemaEvolutionManager(serverConfig);
+    _schemaEvolutionManager = createSchemaEvolutionManager(serverConfig, server);
     _nonDollarVirtualColumnsEnabled = nonDollarVirtualColumnsEnabled;
   }
 
@@ -597,7 +597,7 @@ private String toJsonString(@Nonnull URN urn) {
   }
 
   @Nonnull
-  private SchemaEvolutionManager createSchemaEvolutionManager(@Nonnull ServerConfig serverConfig) {
+  private SchemaEvolutionManager createSchemaEvolutionManager(@Nonnull ServerConfig serverConfig, @Nonnull EbeanServer server) {
     String identifier = serverConfig.getDataSourceConfig().getCustomProperties() != null
         ? serverConfig.getDataSourceConfig().getCustomProperties().getOrDefault(SERVICE_IDENTIFIER, null)
         : null;
@@ -607,7 +607,7 @@ private SchemaEvolutionManager createSchemaEvolutionManager(@Nonnull ServerConfi
         serverConfig.getDataSourceConfig().getUsername(),
         identifier);
 
-    return new FlywaySchemaEvolutionManager(config);
+    return new FlywaySchemaEvolutionManager(config, server);
   }
 
   /**
diff --git a/dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/FlywaySchemaEvolutionManager.java b/dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/FlywaySchemaEvolutionManager.java
index 1d00741bb..c66f1c628 100644
--- a/dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/FlywaySchemaEvolutionManager.java
+++ b/dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/FlywaySchemaEvolutionManager.java
@@ -1,24 +1,40 @@
 package com.linkedin.metadata.dao;
 
+import io.ebean.EbeanServer;
+import io.ebean.SqlRow;
 import java.io.IOException;
 import java.io.InputStream;
+import java.nio.file.Files;
+import java.nio.file.Paths;
+import java.util.List;
 import java.util.Properties;
+import lombok.extern.slf4j.Slf4j;
+import net.sf.jsqlparser.JSQLParserException;
 import org.flywaydb.core.Flyway;
+import org.flywaydb.core.api.MigrationInfo;
 
+import static com.linkedin.metadata.dao.utils.SQLSchemaUtils.*;
+import static com.linkedin.metadata.dao.utils.SQLStatementUtils.*;
 
+
+@Slf4j
 public class FlywaySchemaEvolutionManager implements SchemaEvolutionManager {
   private static final String EVOLUTION_SCRIPTS_LOCATION = "script_directory";
   private static final String VERSION_TABLE = "version_table";
   private static final String CONFIG_FILE_TEMPLATE = "%s.conf";
   private static final String CONFIG_FILE_TEMPLATE2 = "%s-%s.conf";
   private static final String DISABLE_CLEAN = "disable_clean";
+  private static final long HIGH_RISK_THRESHOLD = 2_000_000L;
   private final Flyway _flyway;
+  private final EbeanServer _server;
+  private final String _databaseName;
 
-  public FlywaySchemaEvolutionManager(Config config) {
-    String databaseName = getDatabaseName(config);
+  public FlywaySchemaEvolutionManager(Config config, EbeanServer server) {
+    _server = server;
+    _databaseName = getDatabaseName(config);
     String serviceIdentifier = config.getServiceIdentifier();
     String configFileName = serviceIdentifier == null
-        ? String.format(CONFIG_FILE_TEMPLATE, databaseName) : String.format(CONFIG_FILE_TEMPLATE2, serviceIdentifier, databaseName);
+        ? String.format(CONFIG_FILE_TEMPLATE, _databaseName) : String.format(CONFIG_FILE_TEMPLATE2, serviceIdentifier, _databaseName);
     Properties configProp = new Properties();
 
     try (InputStream configFile = getClass().getClassLoader().getResourceAsStream(configFileName)) {
@@ -46,7 +62,37 @@ public FlywaySchemaEvolutionManager(Config config) {
 
   @Override
   public void ensureSchemaUpToDate() {
-      _flyway.migrate();
+    //Retrieves the full set of infos about pending migrations, available locally, but not yet applied to the DB
+    MigrationInfo[] pendingMigrations = _flyway.info().pending();
+
+    for (MigrationInfo pendingMigration : pendingMigrations) {
+      try {
+        String sql = new String(Files.readAllBytes(Paths.get(pendingMigration.getPhysicalLocation())));
+        List<String> highRiskSQL = detectPotentialHighRiskSQL(sql);
+
+        for (String statement : highRiskSQL) {
+          String table = extractTableName(statement);
+          SqlRow sqlRow = _server.createSqlQuery(getEstimatedRowCount(_databaseName, table)).findOne();
+
+          if (sqlRow != null) {
+            long rowCount = sqlRow.getLong("table_rows");
+
+            if (rowCount > HIGH_RISK_THRESHOLD) {
+              log.error("Estimated row count {}. Apply SQL {} is high risk. Please work with MySQL team to get it applied", rowCount, statement);
+              return;
+            }
+          }
+        }
+
+      } catch (JSQLParserException e) {
+        log.error("Unable to parse the SQL script for potential high risk SQL. Please work with MySQL team to get it applied", e);
+        return;
+      } catch (IOException e) {
+        throw new RuntimeException(e);
+      }
+    }
+
+    _flyway.migrate();
   }
 
   @Override
diff --git a/dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/utils/SQLSchemaUtils.java b/dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/utils/SQLSchemaUtils.java
index f5bdf30b3..025f533bd 100644
--- a/dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/utils/SQLSchemaUtils.java
+++ b/dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/utils/SQLSchemaUtils.java
@@ -8,9 +8,20 @@
 import com.linkedin.metadata.dao.GlobalAssetRegistry;
 import com.linkedin.metadata.dao.exception.MissingAnnotationException;
 import com.linkedin.metadata.dao.exception.ModelValidationException;
+import java.util.ArrayList;
+import java.util.List;
 import java.util.Map;
 import javax.annotation.Nonnull;
 import lombok.extern.slf4j.Slf4j;
+import net.sf.jsqlparser.JSQLParserException;
+import net.sf.jsqlparser.parser.CCJSqlParserUtil;
+import net.sf.jsqlparser.statement.Statement;
+import net.sf.jsqlparser.statement.Statements;
+import net.sf.jsqlparser.statement.alter.Alter;
+import net.sf.jsqlparser.statement.alter.AlterExpression;
+import net.sf.jsqlparser.statement.create.index.CreateIndex;
+import net.sf.jsqlparser.statement.create.table.CreateTable;
+import net.sf.jsqlparser.statement.drop.Drop;
 
 
 /**
@@ -251,4 +262,77 @@ private static String getColumnNameFromAnnotation(@Nonnull final String aspectCa
           String.format("Aspect %s should be annotated with @gma.aspect.column.name.", aspectCanonicalName), e);
     }
   }
+
+  /**
+   * Given a sql script that has one or more SQL statements, detect all potentially high-risk SQL statements.
+   * @param sqlScript A string contains one or more SQL statements.
+   * @return A list of potentially high-risk SQL statements
+   * @throws JSQLParserException if the sql script cannot be properly parsed.
+   */
+  @Nonnull
+  public static List<String> detectPotentialHighRiskSQL(String sqlScript) throws JSQLParserException {
+
+    List<String> potentialHighRiskSQL = new ArrayList<>();
+    Statements sql = CCJSqlParserUtil.parseStatements(sqlScript);
+
+    for (Statement statement : sql.getStatements()) {
+      if (statement instanceof CreateTable) {
+        log.info("Create table is generally low risk. SQL {}", statement);
+      } else if (statement instanceof Alter) {
+        Alter alterStatement = (Alter) statement;
+        List<AlterExpression> expressions = alterStatement.getAlterExpressions();
+
+        // Check if the statement is trying to alter table ... add index ...
+        // Note that alter table ... add column ... is considered low-risk though.
+        for (AlterExpression expression : expressions) {
+          if (expression.getIndex() != null) {
+            log.info("Add index to table is potentially high risk. SQL {}", statement);
+            potentialHighRiskSQL.add(statement.toString());
+            break;
+          }
+        }
+
+      } else if (statement instanceof CreateIndex) {
+        log.info("Create index is potentially high risk. SQL {}", statement);
+        potentialHighRiskSQL.add(statement.toString());
+      } else {
+        // catch all bucket. Any SQL that is not create table, create index, alter table is considered high-risk.
+        log.info("SQL is potentially high-risk {}", statement);
+        potentialHighRiskSQL.add(statement.toString());
+      }
+    }
+
+    return potentialHighRiskSQL;
+  }
+
+  /**
+   * For a SQL statement, extract the table name used in the SQL.
+   * @param sql such as "alter table foo add index bar (abc);"
+   * @return The table name such as "foo"
+   * @throws JSQLParserException if the SQL statement is not one of the supported type.
+   */
+  public static String extractTableName(String sql) throws JSQLParserException {
+    Statement statement = CCJSqlParserUtil.parse(sql);
+
+    if (statement instanceof CreateTable) {
+
+      CreateTable createTable = (CreateTable) statement;
+      return createTable.getTable().getName();
+
+    } else if (statement instanceof Alter) {
+
+      Alter alterStatement = (Alter) statement;
+      return alterStatement.getTable().getName();
+
+    } else if (statement instanceof CreateIndex) {
+      CreateIndex createIndex = (CreateIndex) statement;
+      return createIndex.getTable().getName();
+
+    } else if (statement instanceof Drop) {
+      Drop dropStatement = (Drop) statement;
+      return  dropStatement.getName().getName();
+    }
+
+    throw new JSQLParserException("Unsupported statement type: " + statement.getClass().getName());
+  }
 }
diff --git a/dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/utils/SQLStatementUtils.java b/dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/utils/SQLStatementUtils.java
index 42e24daff..b2dabb406 100644
--- a/dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/utils/SQLStatementUtils.java
+++ b/dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/utils/SQLStatementUtils.java
@@ -141,6 +141,9 @@ public class SQLStatementUtils {
       String.format("SELECT urn, %%s, lastmodifiedon, lastmodifiedby, (SELECT COUNT(urn) FROM %%s) as _total_count "
           + "FROM %%s WHERE %s LIMIT %%d OFFSET %%d", SOFT_DELETED_CHECK);
 
+  private static final String GET_ESTIMATED_COUNT = "select table_rows from information_schema.tables where "
+      + "table_schema = '%s' and table_name = '%s'";
+
   public static final String SOURCE = "source";
   public static final String DESTINATION = "destination";
 
@@ -148,6 +151,13 @@ private SQLStatementUtils() {
     // Util class
   }
 
+  /**
+   * Get estimated row count for a table.
+   */
+  public static String getEstimatedRowCount(String db, String table) {
+    return String.format(GET_ESTIMATED_COUNT, db, table);
+  }
+
   /**
    * Create entity exist SQL statement.
    * @param urn entity urn
diff --git a/dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/FlywaySchemaEvolutionManagerTest.java b/dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/FlywaySchemaEvolutionManagerTest.java
index f7010c590..c36946f42 100644
--- a/dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/FlywaySchemaEvolutionManagerTest.java
+++ b/dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/FlywaySchemaEvolutionManagerTest.java
@@ -30,14 +30,14 @@ public void init() throws IOException {
         EmbeddedMariaInstance.SERVER_CONFIG_MAP.get(_server.getName()).getDataSourceConfig().getCustomProperties().get("SERVICE_IDENTIFIER")
     );
 
-    _schemaEvolutionManager = new FlywaySchemaEvolutionManager(config);
+    _schemaEvolutionManager = new FlywaySchemaEvolutionManager(config, _server);
   }
 
   @Test
   public void testSchemaUpToDate() {
     _schemaEvolutionManager.clean();
 
-    // make sure table did not exists
+    // make sure table did not exist.
     assertFalse(checkTableExists("metadata_entity_foo"));
     assertFalse(checkTableExists("metadata_entity_bar"));
 
diff --git a/dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/FlywaySchemaEvolutionManagerTestWithoutServiceIdentifier.java b/dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/FlywaySchemaEvolutionManagerTestWithoutServiceIdentifier.java
index 76fa119f1..91bc95c2b 100644
--- a/dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/FlywaySchemaEvolutionManagerTestWithoutServiceIdentifier.java
+++ b/dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/FlywaySchemaEvolutionManagerTestWithoutServiceIdentifier.java
@@ -35,7 +35,7 @@ public void init() throws IOException {
         null
     );
 
-    _schemaEvolutionManager = new FlywaySchemaEvolutionManager(config);
+    _schemaEvolutionManager = new FlywaySchemaEvolutionManager(config, _server);
   }
 
   @Test
diff --git a/dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/utils/SQLSchemaUtilsTest.java b/dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/utils/SQLSchemaUtilsTest.java
index 1d9f2e46a..c048b7290 100644
--- a/dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/utils/SQLSchemaUtilsTest.java
+++ b/dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/utils/SQLSchemaUtilsTest.java
@@ -5,6 +5,11 @@
 import com.linkedin.testing.BarAsset;
 import com.linkedin.testing.urn.BarUrn;
 import com.linkedin.testing.urn.FooUrn;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Set;
+import org.testng.annotations.DataProvider;
 import org.testng.annotations.Test;
 
 import static com.linkedin.metadata.dao.utils.SQLSchemaUtils.*;
@@ -13,6 +18,55 @@
 
 public class SQLSchemaUtilsTest {
 
+  @DataProvider
+  public static Object[][] sqlStatement() {
+
+    // each row has three element - ["sql script", "number of high risk sql", "table used in high risk sql"]
+    return new Object[][]{
+
+        // Create new table is low-risk.
+        {"CREATE TABLE IF NOT EXISTS metadata_entity_espressodatabase ("
+            + "    urn VARCHAR(500) NOT NULL,\n"
+            + "    lastmodifiedon TIMESTAMP NOT NULL,\n"
+            + "    lastmodifiedby VARCHAR(255) NOT NULL,\n"
+            + "    createdfor VARCHAR(255),\n"
+            + "    CONSTRAINT pk_metadata_entity_espressodatabase PRIMARY KEY (urn)\n"
+            + ");", 0, Collections.emptySet()},
+
+        // Add new column is low-risk.
+        {"ALTER TABLE metadata_entity_dataset ADD COLUMN a_datasetreferentialauditinfo JSON;", 0, Collections.emptySet()},
+
+        // Add index is high-risk.
+        {"ALTER TABLE metadata_relationship_downstreamof\n"
+            + "ADD INDEX idx_deleted_ts (deleted_ts),\n"
+            + "ADD INDEX idx_source_deleted_ts (source, deleted_ts);", 1,
+            Collections.singleton("metadata_relationship_downstreamof")},
+
+        // Create index is high-risk.
+        {"CREATE INDEX idx_source_deleted_ts ON metadata_relationship_ownedby (source, deleted_ts);", 1,
+            Collections.singleton("metadata_relationship_ownedby")},
+
+        // Mix low-risk and high-risk in one script. Add new virtual column is low-risk but adding index is high risk.
+        {"ALTER TABLE metadata_entity_datapolicy ADD i_datapolicyinfo$purposeUrn VARCHAR(128)"
+            + " GENERATED ALWAYS AS (REPLACE(JSON_EXTRACT(a_datapolicyinfo, '$.aspect.purposeUrn'), '\"', ''));\n"
+            + "CREATE INDEX i_datapolicyinfo$purposeUrn ON metadata_entity_datapolicy (urn, i_datapolicyinfo$purposeUrn);", 1,
+            Collections.singleton("metadata_entity_datapolicy")},
+
+        // Both creating new table and adding new columns is low risk.
+        {"CREATE TABLE IF NOT EXISTS metadata_entity_dataset (\n"
+            + "    urn VARCHAR(500) NOT NULL,\n"
+            + "    lastmodifiedon TIMESTAMP NOT NULL,\n"
+            + "    lastmodifiedby VARCHAR(255) NOT NULL,\n"
+            + "    createdfor VARCHAR(255),\n"
+            + "    CONSTRAINT pk_metadata_entity_dataset PRIMARY KEY (urn)\n"
+            + ");\n"
+            + "ALTER TABLE metadata_entity_dataset ADD COLUMN a_clusterdatasetsla JSON;\n"
+            + "ALTER TABLE metadata_entity_dataset ADD COLUMN a_cmonjiratickets JSON;\n"
+            + "ALTER TABLE metadata_entity_dataset ADD COLUMN a_complianceannotation JSON;\n"
+            + "ALTER TABLE metadata_entity_dataset ADD COLUMN a_complianceinfo JSON;", 0, Collections.emptySet()},
+    };
+  }
+
   @Test
   public void testGetGeneratedColumnName() {
     String generatedColumnName =
@@ -30,4 +84,17 @@ public void testGetAspectColumnName() {
     assertEquals("a_aspect_bar",
         SQLSchemaUtils.getAspectColumnName(BarUrn.ENTITY_TYPE, "com.linkedin.testing.AspectBar"));
   }
+
+  @Test(dataProvider = "sqlStatement")
+  public void testDetectHighRiskSQL(String sql, int numOfHighRiskSQL, Set<String> tables) throws Exception {
+    List<String> highRiskSQL = SQLSchemaUtils.detectPotentialHighRiskSQL(sql);
+    assertEquals(highRiskSQL.size(), numOfHighRiskSQL);
+
+    Set<String> highRiskSet = new HashSet<>();
+    for (String statement : highRiskSQL) {
+      highRiskSet.add(extractTableName(statement));
+    }
+
+    assertEquals(highRiskSet, tables);
+  }
 }
\ No newline at end of file


Reviews:
user: login: copilot-pull-request-reviewer[bot]
body: Copilot encountered an error and was unable to review this pull request. You can try again by re-requesting a review.
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-18T06:10:32Z
user: login: copilot-pull-request-reviewer[bot]
body: ## Pull Request Overview

This PR introduces a smart schema evolution mechanism that classifies DDL operations into low and high risk categories to prevent dangerous operations from being executed automatically. Key changes include:
- Adding utility methods (detectPotentialHighRiskSQL and extractTableName) to parse and classify SQL statements.
- Updating FlywaySchemaEvolutionManager to integrate high-risk checks before migration.
- Enhancing unit tests to validate the risk classification logic.

### Reviewed Changes

Copilot reviewed 9 out of 9 changed files in this pull request and generated 1 comment.

<details>
<summary>Show a summary per file</summary>

| File | Description |
| ---- | ----------- |
| dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/utils/SQLSchemaUtils.java | New SQL parsing methods for detecting high-risk DDL are added. |
| dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/utils/SQLSchemaUtilsTest.java | Added tests to validate the SQL risk detection functionality. |
| dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/FlywaySchemaEvolutionManager.java | Schema evolution manager updated to use the new risk detection logic. |
| dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/utils/SQLStatementUtils.java | Added helper for generating the estimated row count query. |
| dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/EbeanLocalAccess.java | Updated instantiation of the schema evolution manager. |
| Test files | Modified to use the updated FlywaySchemaEvolutionManager constructor. |
</details>





state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-18T06:11:51Z
user: login: yangyangv2
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-18T16:26:10Z
user: login: yangyangv2
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-18T16:29:05Z
user: login: yangyangv2
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-18T16:34:44Z
user: login: zhixuanjia
state: COMMENTED
author_association: NONE
submitted_at: 2025-03-18T20:18:31Z
user: login: zhixuanjia
state: COMMENTED
author_association: NONE
submitted_at: 2025-03-18T23:46:03Z
user: login: yangyangv2
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-19T19:03:11Z
user: login: yangyangv2
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-19T19:04:07Z
user: login: yangyangv2
body: LGTM!
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-03-19T19:05:09Z
user: login: rakhiagr
body: Thank you for answering these questions. LGTM!
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-03-19T19:54:10Z

Review comments:
diff_hunk: @@ -251,4 +261,79 @@ private static String getColumnNameFromAnnotation(@Nonnull final String aspectCa
           String.format("Aspect %s should be annotated with @gma.aspect.column.name.", aspectCanonicalName), e);
     }
   }
+
+  /**
+   * Given a sql script that has one or more SQL statements, detect all potentially high-risk SQL statements.
+   * @param sqlScript A string contains one or more SQL statements.
+   * @return A list of potentially high-risk SQL statements
+   * @throws JSQLParserException if the sql script cannot be properly parsed.
+   */
+  public static List<String> detectPotentialHighRiskSQL(String sqlScript) throws JSQLParserException {
+
+    // Split SQL into individual statements
+    String[] statements = sqlScript.split(";\\s*"); // Splits on semicolons.
+    List<String> potentialHighRiskSQL = new ArrayList<>();
+
+    for (String statement : statements) {
path: dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/utils/SQLSchemaUtils.java
user: login: Copilot
body: Consider trimming each statement and skipping empty strings after splitting the SQL script to avoid parse errors with CCJSqlParserUtil.parse().
```suggestion
    for (String statement : statements) {
      statement = statement.trim();
      if (statement.isEmpty()) {
        continue;
      }
```
created_at: 2025-03-18T06:11:51Z
author_association: CONTRIBUTOR
diff_hunk: @@ -251,4 +261,79 @@ private static String getColumnNameFromAnnotation(@Nonnull final String aspectCa
           String.format("Aspect %s should be annotated with @gma.aspect.column.name.", aspectCanonicalName), e);
     }
   }
+
+  /**
+   * Given a sql script that has one or more SQL statements, detect all potentially high-risk SQL statements.
+   * @param sqlScript A string contains one or more SQL statements.
+   * @return A list of potentially high-risk SQL statements
+   * @throws JSQLParserException if the sql script cannot be properly parsed.
+   */
+  public static List<String> detectPotentialHighRiskSQL(String sqlScript) throws JSQLParserException {
+
+    // Split SQL into individual statements
+    String[] statements = sqlScript.split(";\\s*"); // Splits on semicolons.
path: dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/utils/SQLSchemaUtils.java
user: login: yangyangv2
body: split by ";" is not safe for edge cases, consider using SQL parser such as calcite

```
import org.apache.calcite.sql.parser.SqlParser;

SqlParser parser = SqlParser.create(sql);
 List<String> queries = new ArrayList<>();
        
while (!parser.parseStmtList().getList().isEmpty()) {
      SqlNode stmt = parser.parseStmt();
      queries.add(stmt.toString());
}
```

Li has wide usage of calcite internally https://jarvis.corp.linkedin.com/codesearch/results?query=org.apache.calcite.sql.parser.SqlParser
created_at: 2025-03-18T16:26:10Z
author_association: CONTRIBUTOR
diff_hunk: @@ -141,13 +141,23 @@ public class SQLStatementUtils {
       String.format("SELECT urn, %%s, lastmodifiedon, lastmodifiedby, (SELECT COUNT(urn) FROM %%s) as _total_count "
           + "FROM %%s WHERE %s LIMIT %%d OFFSET %%d", SOFT_DELETED_CHECK);
 
+  private static final String GET_ESTIMATED_COUNT = "select table_rows from information_schema.tables where "
path: dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/utils/SQLStatementUtils.java
user: login: yangyangv2
body: Good use case for estimate rows!
created_at: 2025-03-18T16:29:04Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1,24 +1,40 @@
 package com.linkedin.metadata.dao;
 
+import io.ebean.EbeanServer;
+import io.ebean.SqlRow;
 import java.io.IOException;
 import java.io.InputStream;
+import java.nio.file.Files;
+import java.nio.file.Paths;
+import java.util.List;
 import java.util.Properties;
+import lombok.extern.slf4j.Slf4j;
+import net.sf.jsqlparser.JSQLParserException;
 import org.flywaydb.core.Flyway;
+import org.flywaydb.core.api.MigrationInfo;
 
+import static com.linkedin.metadata.dao.utils.SQLSchemaUtils.*;
+import static com.linkedin.metadata.dao.utils.SQLStatementUtils.*;
 
+
+@Slf4j
 public class FlywaySchemaEvolutionManager implements SchemaEvolutionManager {
   private static final String EVOLUTION_SCRIPTS_LOCATION = "script_directory";
   private static final String VERSION_TABLE = "version_table";
   private static final String CONFIG_FILE_TEMPLATE = "%s.conf";
   private static final String CONFIG_FILE_TEMPLATE2 = "%s-%s.conf";
   private static final String DISABLE_CLEAN = "disable_clean";
+  private static final long HIGH_RISK_THRESHOLD = 500_000L;
path: dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/FlywaySchemaEvolutionManager.java
user: login: yangyangv2
body: last week we tried adding indexes for 2 million rows and only takes several seconds (< 20s). 

is 500k the recommended number from MySQL team or industry recommendation? 
created_at: 2025-03-18T16:34:44Z
author_association: CONTRIBUTOR
diff_hunk: @@ -251,4 +261,79 @@ private static String getColumnNameFromAnnotation(@Nonnull final String aspectCa
           String.format("Aspect %s should be annotated with @gma.aspect.column.name.", aspectCanonicalName), e);
     }
   }
+
+  /**
+   * Given a sql script that has one or more SQL statements, detect all potentially high-risk SQL statements.
+   * @param sqlScript A string contains one or more SQL statements.
+   * @return A list of potentially high-risk SQL statements
+   * @throws JSQLParserException if the sql script cannot be properly parsed.
+   */
+  public static List<String> detectPotentialHighRiskSQL(String sqlScript) throws JSQLParserException {
+
+    // Split SQL into individual statements
+    String[] statements = sqlScript.split(";\\s*"); // Splits on semicolons.
path: dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/utils/SQLSchemaUtils.java
user: login: zhixuanjia
body: good point, JSqlParser is able parse sql script to multiple statements. Updated to
```
 Statements sql = CCJSqlParserUtil.parseStatements(sqlScript);
```

For calcite, I've tried it but it always throw Unexpected Token exception... give it up after a few hours of investigation with no luck. JSqlParser seems also widely used with 5k+ stars on github. I think it should be reliable solution too.
created_at: 2025-03-18T20:18:31Z
author_association: NONE
diff_hunk: @@ -1,24 +1,40 @@
 package com.linkedin.metadata.dao;
 
+import io.ebean.EbeanServer;
+import io.ebean.SqlRow;
 import java.io.IOException;
 import java.io.InputStream;
+import java.nio.file.Files;
+import java.nio.file.Paths;
+import java.util.List;
 import java.util.Properties;
+import lombok.extern.slf4j.Slf4j;
+import net.sf.jsqlparser.JSQLParserException;
 import org.flywaydb.core.Flyway;
+import org.flywaydb.core.api.MigrationInfo;
 
+import static com.linkedin.metadata.dao.utils.SQLSchemaUtils.*;
+import static com.linkedin.metadata.dao.utils.SQLStatementUtils.*;
 
+
+@Slf4j
 public class FlywaySchemaEvolutionManager implements SchemaEvolutionManager {
   private static final String EVOLUTION_SCRIPTS_LOCATION = "script_directory";
   private static final String VERSION_TABLE = "version_table";
   private static final String CONFIG_FILE_TEMPLATE = "%s.conf";
   private static final String CONFIG_FILE_TEMPLATE2 = "%s-%s.conf";
   private static final String DISABLE_CLEAN = "disable_clean";
+  private static final long HIGH_RISK_THRESHOLD = 500_000L;
path: dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/FlywaySchemaEvolutionManager.java
user: login: zhixuanjia
body: The execution time depends on several factors 
1. composite index vs single index
2. the data type we try to index
3. server load at the exact moment when we execute DDL.
...

I think this threshold is subject to tuning. If we find it unnecessarily classify too many DDL as high risk, we can change it a higher number. Based your feedback, I've tune up the threshold to 2M. Let's see how it goes?
created_at: 2025-03-18T23:46:03Z
author_association: NONE
diff_hunk: @@ -251,4 +261,79 @@ private static String getColumnNameFromAnnotation(@Nonnull final String aspectCa
           String.format("Aspect %s should be annotated with @gma.aspect.column.name.", aspectCanonicalName), e);
     }
   }
+
+  /**
+   * Given a sql script that has one or more SQL statements, detect all potentially high-risk SQL statements.
+   * @param sqlScript A string contains one or more SQL statements.
+   * @return A list of potentially high-risk SQL statements
+   * @throws JSQLParserException if the sql script cannot be properly parsed.
+   */
+  public static List<String> detectPotentialHighRiskSQL(String sqlScript) throws JSQLParserException {
+
+    // Split SQL into individual statements
+    String[] statements = sqlScript.split(";\\s*"); // Splits on semicolons.
path: dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/utils/SQLSchemaUtils.java
user: login: yangyangv2
body: sounds good, it's a surprise Calcite is this fragile.  
created_at: 2025-03-19T19:03:11Z
author_association: CONTRIBUTOR
diff_hunk: @@ -1,24 +1,40 @@
 package com.linkedin.metadata.dao;
 
+import io.ebean.EbeanServer;
+import io.ebean.SqlRow;
 import java.io.IOException;
 import java.io.InputStream;
+import java.nio.file.Files;
+import java.nio.file.Paths;
+import java.util.List;
 import java.util.Properties;
+import lombok.extern.slf4j.Slf4j;
+import net.sf.jsqlparser.JSQLParserException;
 import org.flywaydb.core.Flyway;
+import org.flywaydb.core.api.MigrationInfo;
 
+import static com.linkedin.metadata.dao.utils.SQLSchemaUtils.*;
+import static com.linkedin.metadata.dao.utils.SQLStatementUtils.*;
 
+
+@Slf4j
 public class FlywaySchemaEvolutionManager implements SchemaEvolutionManager {
   private static final String EVOLUTION_SCRIPTS_LOCATION = "script_directory";
   private static final String VERSION_TABLE = "version_table";
   private static final String CONFIG_FILE_TEMPLATE = "%s.conf";
   private static final String CONFIG_FILE_TEMPLATE2 = "%s-%s.conf";
   private static final String DISABLE_CLEAN = "disable_clean";
+  private static final long HIGH_RISK_THRESHOLD = 500_000L;
path: dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/FlywaySchemaEvolutionManager.java
user: login: yangyangv2
body: let's keep an eye on the real-life practice and adjust it accordingly. 
created_at: 2025-03-19T19:04:07Z
author_association: CONTRIBUTOR

PR comments:
user: login: rakhiagr
created_at: 2025-03-18T06:10:32Z
body: Few clarifying questions:

1. What happens if a script contains both low risk DDL and high risk DDL?
2. For high risk script, how are we planning to manually mark it as completed in version table after MySQL team applied the script? (We don't want write operation on db directly, we also need checksum value in flyway table).
author_association: CONTRIBUTOR
user: login: yangyangv2
created_at: 2025-03-18T16:37:38Z
body: > Few clarifying questions:
> 
> 1. What happens if a script contains both low risk DDL and high risk DDL?
> 2. For high risk script, how are we planning to manually mark it as completed in version table after MySQL team applied the script? (We don't want write operation on db directly, we also need checksum value in flyway table).

Similar question, how do we track which statements are executed which one are not. 
Will there be execution sequence dependencies in our DDLs? 
author_association: CONTRIBUTOR
user: login: zhixuanjia
created_at: 2025-03-18T20:16:03Z
body: > Few clarifying questions:
> 
> 1. What happens if a script contains both low risk DDL and high risk DDL?
> 2. For high risk script, how are we planning to manually mark it as completed in version table after MySQL team applied the script? (We don't want write operation on db directly, we also need checksum value in flyway table).

1. If a script has both low and hight risk DDL, we will not automatically apply it. We need to keep the integrity of a single script. From this perspective, it is best practice to keep high risk DDL in its own script.
2. This is discussed in the "Other work needed" section - For high risk script, we will need to manually mark it as completed in version table after MySQL team applied the script. We will calculate the checksum of the script by ourselves. 
author_association: NONE
user: login: zhixuanjia
created_at: 2025-03-18T23:48:41Z
body: > > Few clarifying questions:
> > 
> > 1. What happens if a script contains both low risk DDL and high risk DDL?
> > 2. For high risk script, how are we planning to manually mark it as completed in version table after MySQL team applied the script? (We don't want write operation on db directly, we also need checksum value in flyway table).
> 
> Similar question, how do we track which statements are executed which one are not. Will there be execution sequence dependencies in our DDLs?

So we will first fetch all the pending scripts to be applied to the target DB. If anyone of the scripts contains a high-risk SQL, we will not apply any of the pending scripts.
author_association: NONE
user: login: yangyangv2
created_at: 2025-03-19T19:04:49Z
body: > > > Few clarifying questions:
> > > 
> > > 1. What happens if a script contains both low risk DDL and high risk DDL?
> > > 2. For high risk script, how are we planning to manually mark it as completed in version table after MySQL team applied the script? (We don't want write operation on db directly, we also need checksum value in flyway table).
> > 
> > 
> > Similar question, how do we track which statements are executed which one are not. Will there be execution sequence dependencies in our DDLs?
> 
> So we will first fetch all the pending scripts to be applied to the target DB. If anyone of the scripts contains a high-risk SQL, we will not apply any of the pending scripts.

+1 to the execution integrity. 

LGTM, thanks for driving the onboarding productivity improvement!!
author_association: CONTRIBUTOR