Fix yugabyte tablegroups in colocated database (sqlancer#795)

zhenglin-charlie-li · web-flow · commit f0c7eedd1532 · 2023-05-30T09:29:31.000+08:00
* fix: yugabyte error: cannot use tablegroups in a colocated database

* fix: delete useless import

* fix: column index starts from 1

* fix: change check colocated logic when creating table

* fix: add exception handle

* style: check statement will not return an error

* fix: remove useless null check

* fix: PMD violation
diff --git a/src/sqlancer/yugabyte/ysql/YSQLSchema.java b/src/sqlancer/yugabyte/ysql/YSQLSchema.java
@@ -160,6 +160,18 @@ protected static List<YSQLColumn> getTableColumns(SQLConnection con, String tabl
         return columns;
     }
 
+    public boolean getDatabaseIsColocated(SQLConnection con) {
+        try (Statement s = con.createStatement(); ResultSet rs = s.executeQuery("SELECT yb_is_database_colocated();")) {
+            rs.next();
+            String result = rs.getString(1);
+            // The query will result in a 'f' for a non-colocated database
+            return !"f".equals(result);
+
+        } catch (SQLException e) {
+            throw new AssertionError(e);
+        }
+    }
+
     public YSQLTables getRandomTableNonEmptyTables() {
         return new YSQLTables(Randomly.nonEmptySubset(getDatabaseTables()));
     }
diff --git a/src/sqlancer/yugabyte/ysql/gen/YSQLCommon.java b/src/sqlancer/yugabyte/ysql/gen/YSQLCommon.java
@@ -159,8 +159,10 @@ public static void generateWith(StringBuilder sb, YSQLGlobalState globalState, E
             }
         } else if (Randomly.getBoolean()) {
             errors.add("Cannot use TABLEGROUP with TEMP table");
-            sb.append(" TABLEGROUP tg").append(
-                    Randomly.getNotCachedInteger(1, (int) YSQLTableGroupGenerator.UNIQUE_TABLEGROUP_COUNTER.get()));
+            if (!globalState.getSchema().getDatabaseIsColocated(globalState.getConnection())) {
+                sb.append(" TABLEGROUP tg").append(
+                        Randomly.getNotCachedInteger(1, (int) YSQLTableGroupGenerator.UNIQUE_TABLEGROUP_COUNTER.get()));
+            }
         }
     }
 

Original file line number	Diff line number	Diff line change
`@@ -159,8 +159,10 @@ public static void generateWith(StringBuilder sb, YSQLGlobalState globalState, E`
`159`	`159`	`}`
`160`	`160`	`} else if (Randomly.getBoolean()) {`
`161`	`161`	`errors.add("Cannot use TABLEGROUP with TEMP table");`
`162`		`- sb.append(" TABLEGROUP tg").append(`
`163`		`- Randomly.getNotCachedInteger(1, (int) YSQLTableGroupGenerator.UNIQUE_TABLEGROUP_COUNTER.get()));`
	`162`	`+ if (!globalState.getSchema().getDatabaseIsColocated(globalState.getConnection())) {`
	`163`	`+ sb.append(" TABLEGROUP tg").append(`
	`164`	`+ Randomly.getNotCachedInteger(1, (int) YSQLTableGroupGenerator.UNIQUE_TABLEGROUP_COUNTER.get()));`
	`165`	`+ }`
`164`	`166`	`}`
`165`	`167`	`}`
`166`	`168`