SNOW-1003125: Fix flaky test DatabaseMetaDataIT.testGetSchemas

src/test/java/net/snowflake/client/TestUtil.java

@@ -22,6 +22,20 @@ public class TestUtil {
   private static final Pattern QUERY_ID_REGEX =
       Pattern.compile("[a-z0-9]{8}-[a-z0-9]{4}-[a-z0-9]{4}-[a-z0-9]{4}-[a-z0-9]{12}");
 
+  public static final String GENERATED_SCHEMA_PREFIX = "GENERATED_";
+  public static final String ESCAPED_GENERATED_SCHEMA_PREFIX =
+      GENERATED_SCHEMA_PREFIX.replaceAll("_", "\\\\_");
+  private static final String GITHUB_SCHEMA_PREFIX =
+      "GITHUB_"; // created by JDBC CI jobs before tests
+  private static final String GITHUB_JOB_SCHEMA_PREFIX =
+      "GH_JOB_"; // created by other drivers e.g. Python driver
+
+  public static boolean isSchemaGeneratedInTests(String schema) {
+    return schema.startsWith(TestUtil.GITHUB_SCHEMA_PREFIX)
+        || schema.startsWith(TestUtil.GITHUB_JOB_SCHEMA_PREFIX)
+        || schema.startsWith(TestUtil.GENERATED_SCHEMA_PREFIX);
+  }
+
   /**
    * Util function to assert a piece will throw exception and assert on the error code
    *
@@ -99,7 +113,8 @@ public static void withSchema(Statement statement, String schemaName, ThrowingRu
    */
   public static void withRandomSchema(Statement statement, ThrowingConsumer<String> action)
       throws Exception {
-    String customSchema = "TEST_SCHEMA_" + SnowflakeUtil.randomAlphaNumeric(5);
+    String customSchema =
+        GENERATED_SCHEMA_PREFIX + SnowflakeUtil.randomAlphaNumeric(5).toUpperCase();
     try {
       statement.execute("CREATE OR REPLACE SCHEMA " + customSchema);
       action.call(customSchema);

src/test/java/net/snowflake/client/jdbc/DatabaseMetaDataIT.java

@@ -7,6 +7,7 @@
 import static java.sql.ResultSetMetaData.columnNullableUnknown;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.greaterThanOrEqualTo;
+import static org.hamcrest.Matchers.hasItem;
 import static org.junit.Assert.*;
 
 import com.google.common.base.Strings;
@@ -139,41 +140,50 @@ public void testGetSchemas() throws Throwable {
     // CLIENT_METADATA_REQUEST_USE_CONNECTION_CTX = false
     try (Connection connection = getConnection()) {
       DatabaseMetaData metaData = connection.getMetaData();
+      String currentSchema = connection.getSchema();
       assertEquals("schema", metaData.getSchemaTerm());
-      ResultSet resultSet = metaData.getSchemas();
-      verifyResultSetMetaDataColumns(resultSet, DBMetadataResultSetMetadata.GET_SCHEMAS);
-
       Set<String> schemas = new HashSet<>();
-      while (resultSet.next()) {
-        schemas.add(resultSet.getString(1));
+      try (ResultSet resultSet = metaData.getSchemas()) {
+        verifyResultSetMetaDataColumns(resultSet, DBMetadataResultSetMetadata.GET_SCHEMAS);
+        while (resultSet.next()) {
+          String schema = resultSet.getString(1);
+          if (currentSchema.equals(schema) || !TestUtil.isSchemaGeneratedInTests(schema)) {
+            schemas.add(schema);
+          }
+        }
       }
-      resultSet.close();
-      assertThat(schemas.size(), greaterThanOrEqualTo(1)); // one or more schemas
+      assertThat(schemas.size(), greaterThanOrEqualTo(1));
 
       Set<String> schemasInDb = new HashSet<>();
-      resultSet = metaData.getSchemas(connection.getCatalog(), "%");
-      while (resultSet.next()) {
-        schemasInDb.add(resultSet.getString(1));
+      try (ResultSet resultSet = metaData.getSchemas(connection.getCatalog(), "%")) {
+        while (resultSet.next()) {
+          String schema = resultSet.getString(1);
+          if (currentSchema.equals(schema) || !TestUtil.isSchemaGeneratedInTests(schema)) {
+            schemasInDb.add(schema);
+          }
+        }
       }
-      assertThat(schemasInDb.size(), greaterThanOrEqualTo(1)); // one or more schemas
+      assertThat(schemasInDb.size(), greaterThanOrEqualTo(1));
       assertThat(schemas.size(), greaterThanOrEqualTo(schemasInDb.size()));
-      assertTrue(schemas.containsAll(schemasInDb));
-      assertTrue(schemas.contains(connection.getSchema()));
+      schemasInDb.forEach(schemaInDb -> assertThat(schemas, hasItem(schemaInDb)));
+      assertTrue(schemas.contains(currentSchema));
+      assertTrue(schemasInDb.contains(currentSchema));
     }
 
     // CLIENT_METADATA_REQUEST_USE_CONNECTION_CTX = true
-    try (Connection connection = getConnection()) {
-      Statement statement = connection.createStatement();
+    try (Connection connection = getConnection();
+        Statement statement = connection.createStatement()) {
       statement.execute("alter SESSION set CLIENT_METADATA_REQUEST_USE_CONNECTION_CTX=true");
 
       DatabaseMetaData metaData = connection.getMetaData();
       assertEquals("schema", metaData.getSchemaTerm());
-      ResultSet resultSet = metaData.getSchemas();
-      Set<String> schemas = new HashSet<>();
-      while (resultSet.next()) {
-        schemas.add(resultSet.getString(1));
+      try (ResultSet resultSet = metaData.getSchemas()) {
+        Set<String> schemas = new HashSet<>();
+        while (resultSet.next()) {
+          schemas.add(resultSet.getString(1));
+        }
+        assertThat(schemas.size(), equalTo(1));
       }
-      assertThat(schemas.size(), equalTo(1)); // more than 1 schema
     }
   }
 

src/test/java/net/snowflake/client/jdbc/DatabaseMetaDataLatestIT.java

@@ -171,11 +171,20 @@ public void testDoubleQuotedDatabaseAndSchema() throws Exception {
       // To query the schema and table, we can use a normal java escaped quote. Wildcards are also
       // escaped here
       String schemaRandomPart = SnowflakeUtil.randomAlphaNumeric(5);
-      String querySchema = "TEST\\_SCHEMA\\_\"WITH\\_QUOTES" + schemaRandomPart + "\"";
+      String querySchema =
+          TestUtil.ESCAPED_GENERATED_SCHEMA_PREFIX
+              + "TEST\\_SCHEMA\\_\"WITH\\_QUOTES"
+              + schemaRandomPart
+              + "\"";
       String queryTable = "TESTTABLE\\_\"WITH\\_QUOTES\"";
       // Create the schema and table. With SQL commands, double quotes must be escaped with another
       // quote
-      String schemaName = "\"TEST_SCHEMA_\"\"WITH_QUOTES" + schemaRandomPart + "\"\"\"";
+      String schemaName =
+          "\""
+              + TestUtil.GENERATED_SCHEMA_PREFIX
+              + "TEST_SCHEMA_\"\"WITH_QUOTES"
+              + schemaRandomPart
+              + "\"\"\"";
       TestUtil.withSchema(
           statement,
           schemaName,
@@ -367,24 +376,26 @@ public void testGetFunctionSqlInjectionProtection() throws Throwable {
    */
   @Test
   public void testGetProcedureColumnsWildcards() throws SQLException {
-    try (Connection con = getConnection()) {
-      Statement statement = con.createStatement();
+    try (Connection con = getConnection();
+        Statement statement = con.createStatement()) {
       String database = con.getCatalog();
-      String schema1 = "SCH1";
-      String schema2 = "SCH2";
+      String schemaPrefix =
+          TestUtil.GENERATED_SCHEMA_PREFIX + SnowflakeUtil.randomAlphaNumeric(5).toUpperCase();
+      String schema1 = schemaPrefix + "SCH1";
+      String schema2 = schemaPrefix + "SCH2";
       // Create 2 schemas, each with the same stored procedure declared in them
       statement.execute("create or replace schema " + schema1);
       statement.execute(TEST_PROC);
       statement.execute("create or replace schema " + schema2);
       statement.execute(TEST_PROC);
       DatabaseMetaData metaData = con.getMetaData();
-      ResultSet rs = metaData.getProcedureColumns(database, "SCH_", "TESTPROC", "PARAM1");
-      // Assert 4 rows returned for the param PARAM1 that's present in each of the 2 identical
-      // stored procs in different schemas. A result row is returned for each procedure, making the
-      // total rowcount 4
-      assertEquals(4, getSizeOfResultSet(rs));
-      rs.close();
-      statement.close();
+      try (ResultSet rs =
+          metaData.getProcedureColumns(database, schemaPrefix + "SCH_", "TESTPROC", "PARAM1")) {
+        // Assert 4 rows returned for the param PARAM1 that's present in each of the 2 identical
+        // stored procs in different schemas. A result row is returned for each procedure, making
+        // the total rowcount 4
+        assertEquals(4, getSizeOfResultSet(rs));
+      }
     }
   }
 
@@ -453,7 +464,7 @@ public void testGetStringValueFromColumnDef() throws SQLException {
 
   @Test
   public void testGetColumnsNullable() throws Throwable {
-    try (Connection connection = getConnection()) {
+    try (Connection connection = getConnection(); ) {
       String database = connection.getCatalog();
       String schema = connection.getSchema();
       final String targetTable = "T0";
@@ -484,8 +495,8 @@ public void testSessionDatabaseParameter() throws Throwable {
     String altdb = "ALTERNATEDB";
     String altschema1 = "ALTERNATESCHEMA1";
     String altschema2 = "ALTERNATESCHEMA2";
-    try (Connection connection = getConnection()) {
-      Statement statement = connection.createStatement();
+    try (Connection connection = getConnection();
+        Statement statement = connection.createStatement()) {
       String catalog = connection.getCatalog();
       String schema = connection.getSchema();
       statement.execute("create or replace database " + altdb);
@@ -905,7 +916,12 @@ public void testHandlingSpecialChars() throws Exception {
       statement.execute("INSERT INTO \"TEST\\1\\_1\" (\"C%1\",\"C\\1\\\\11\") VALUES (0,0)");
       // test getColumns with escaped special characters in schema and table name
       String specialSchemaSuffix = SnowflakeUtil.randomAlphaNumeric(5);
-      String specialSchema = "\"SPECIAL%_\\SCHEMA" + specialSchemaSuffix + "\"";
+      String specialSchema =
+          "\""
+              + TestUtil.GENERATED_SCHEMA_PREFIX
+              + "SPECIAL%_\\SCHEMA"
+              + specialSchemaSuffix
+              + "\"";
       TestUtil.withSchema(
           statement,
           specialSchema,
@@ -939,7 +955,13 @@ public void testHandlingSpecialChars() throws Exception {
 
             String escapedTable2 = "TEST" + escapeChar + "_1" + escapeChar + "_1";
             String escapedSchema =
-                "SPECIAL%" + escapeChar + "_" + escapeChar + "\\SCHEMA" + specialSchemaSuffix;
+                TestUtil.ESCAPED_GENERATED_SCHEMA_PREFIX
+                    + "SPECIAL%"
+                    + escapeChar
+                    + "_"
+                    + escapeChar
+                    + "\\SCHEMA"
+                    + specialSchemaSuffix;
             resultSet = metaData.getColumns(database, escapedSchema, escapedTable2, null);
             assertTrue(resultSet.next());
             assertEquals("RNUM", resultSet.getString("COLUMN_NAME"));

src/test/java/net/snowflake/client/jdbc/StatementLatestIT.java

@@ -229,7 +229,9 @@ public void testPreparedStatementLogging() throws SQLException {
 
   @Test // SNOW-647217
   public void testSchemaWith255CharactersDoesNotCauseException() throws SQLException {
-    String schemaName = "a" + SnowflakeUtil.randomAlphaNumeric(254);
+    String schemaName =
+        TestUtil.GENERATED_SCHEMA_PREFIX
+            + SnowflakeUtil.randomAlphaNumeric(255 - TestUtil.GENERATED_SCHEMA_PREFIX.length());
     try (Connection con = getConnection()) {
       try (Statement stmt = con.createStatement()) {
         stmt.execute("create schema " + schemaName);