Skip to content

Commit 8589b16

Browse files
decsterdecster
authored
[Enhancement] Add groovy script engine to FE (StarRocks#21403)
Signed-off-by: Binglin Chang <[email protected]>
1 parent 27c8427 commit 8589b16

File tree

10 files changed

+355
-5
lines changed

10 files changed

+355
-5
lines changed

‎fe/fe-core/pom.xml

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -764,6 +764,12 @@ under the License.
764764
<scope>test</scope>
765765
</dependency>
766766

767+
<!-- https://mvnrepository.com/artifact/org.apache.groovy/groovy-groovysh -->
768+
<dependency>
769+
<groupId>org.apache.groovy</groupId>
770+
<artifactId>groovy-groovysh</artifactId>
771+
</dependency>
772+
767773
<dependency>
768774
<groupId>io.trino</groupId>
769775
<artifactId>trino-parser</artifactId>

‎fe/fe-core/src/main/java/com/starrocks/common/Config.java

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2173,10 +2173,12 @@ public class Config extends ConfigBase {
21732173
@ConfField(mutable = true)
21742174
public static long max_per_node_grep_log_limit = 500000;
21752175

2176-
21772176
/**
21782177
* only use compute node, test for multi-warehouse
21792178
*/
21802179
@ConfField(mutable = true)
21812180
public static boolean only_use_compute_node = false;
2181+
2182+
@ConfField
2183+
public static boolean enable_execute_script_on_frontend = true;
21822184
}

‎fe/fe-core/src/main/java/com/starrocks/monitor/jvm/JvmInfo.java

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,8 +24,11 @@
2424
import java.lang.management.ManagementPermission;
2525
import java.lang.management.MemoryMXBean;
2626
import java.lang.management.MemoryPoolMXBean;
27+
import java.lang.management.MemoryUsage;
2728
import java.lang.management.PlatformManagedObject;
2829
import java.lang.management.RuntimeMXBean;
30+
import java.lang.management.ThreadInfo;
31+
import java.lang.management.ThreadMXBean;
2932
import java.lang.reflect.Method;
3033
import java.util.Collections;
3134
import java.util.List;
@@ -466,4 +469,38 @@ public String toString() {
466469
return sb.toString();
467470
}
468471
}
472+
473+
public static String getThreadDump() {
474+
return getThreadDumpWithPattern(null);
475+
}
476+
477+
public static String getThreadDumpWithPattern(String pattern) {
478+
StringBuilder sb = new StringBuilder();
479+
ThreadMXBean threadMXBean = ManagementFactory.getThreadMXBean();
480+
ThreadInfo[] threadInfos = threadMXBean.dumpAllThreads(true, true);
481+
for (ThreadInfo threadInfo : threadInfos) {
482+
if (threadInfo != null) {
483+
String dump = threadInfo.toString();
484+
if (pattern == null || dump.contains(pattern)) {
485+
sb.append(dump);
486+
sb.append('\n');
487+
}
488+
}
489+
}
490+
return sb.toString();
491+
}
492+
493+
public static String getMemoryStats() {
494+
List<MemoryPoolMXBean> memoryPools = ManagementFactory.getMemoryPoolMXBeans();
495+
StringBuilder sb = new StringBuilder();
496+
for (MemoryPoolMXBean pool : memoryPools) {
497+
sb.append("Pool: " + pool.getName() + "\n");
498+
MemoryUsage usage = pool.getUsage();
499+
sb.append(" - Initial: " + usage.getInit() + "\n");
500+
sb.append(" - Used: " + usage.getUsed() + "\n");
501+
sb.append(" - Committed: " + usage.getCommitted() + "\n");
502+
sb.append(" - Maximum: " + usage.getMax() + "\n");
503+
}
504+
return sb.toString();
505+
}
469506
}

‎fe/fe-core/src/main/java/com/starrocks/qe/ExecuteScriptExecutor.java

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414

1515
package com.starrocks.qe;
1616

17+
import com.starrocks.common.Config;
1718
import com.starrocks.common.UserException;
1819
import com.starrocks.proto.ExecuteCommandRequestPB;
1920
import com.starrocks.proto.ExecuteCommandResultPB;
@@ -22,6 +23,8 @@
2223
import com.starrocks.sql.ast.ExecuteScriptStmt;
2324
import com.starrocks.system.Backend;
2425
import com.starrocks.thrift.TNetworkAddress;
26+
import groovy.lang.Binding;
27+
import groovy.lang.GroovyShell;
2528
import org.apache.commons.lang.StringUtils;
2629
import org.apache.logging.log4j.LogManager;
2730
import org.apache.logging.log4j.Logger;
@@ -33,6 +36,32 @@ public class ExecuteScriptExecutor {
3336
private static final Logger LOG = LogManager.getLogger(ExecuteScriptExecutor.class);
3437

3538
public static void execute(ExecuteScriptStmt stmt, ConnectContext ctx) throws UserException {
39+
if (stmt.isFrontendScript()) {
40+
executeFrontendScript(stmt, ctx);
41+
} else {
42+
executeBackendScript(stmt, ctx);
43+
}
44+
}
45+
46+
private static void executeFrontendScript(ExecuteScriptStmt stmt, ConnectContext ctx) throws UserException {
47+
if (!Config.enable_execute_script_on_frontend) {
48+
throw new UserException("execute script on frontend is disabled");
49+
}
50+
try {
51+
StringBuilder sb = new StringBuilder();
52+
Binding binding = new Binding();
53+
binding.setVariable("LOG", LOG);
54+
binding.setVariable("out", sb);
55+
binding.setVariable("globalState", GlobalStateMgr.getCurrentState());
56+
GroovyShell shell = new GroovyShell(binding);
57+
shell.evaluate(stmt.getScript());
58+
ctx.getState().setOk(0, 0, sb.toString());
59+
} catch (Exception e) {
60+
throw new UserException("execute script failed: " + e.getMessage());
61+
}
62+
}
63+
64+
private static void executeBackendScript(ExecuteScriptStmt stmt, ConnectContext ctx) throws UserException {
3665
Backend be = GlobalStateMgr.getCurrentSystemInfo().getBackend(stmt.getBeId());
3766
if (be == null) {
3867
throw new UserException("node not found: " + stmt.getBeId());

‎fe/fe-core/src/main/java/com/starrocks/sql/ast/ExecuteScriptStmt.java

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,10 @@ public String getScript() {
4444
return script;
4545
}
4646

47+
public boolean isFrontendScript() {
48+
return beId == -1;
49+
}
50+
4751
public void setTimeoutSec(long timeoutSec) {
4852
this.timeoutSec = timeoutSec;
4953
}

‎fe/fe-core/src/main/java/com/starrocks/sql/parser/AstBuilder.java

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -511,7 +511,7 @@ public ParseNode visitCreateDbStatement(StarRocksParser.CreateDbStatementContext
511511

512512
String dbName = getIdentifierName(context.database);
513513

514-
Map<String, String> properties = new HashMap<>();;
514+
Map<String, String> properties = new HashMap<>();
515515
if (context.properties() != null) {
516516
List<Property> propertyList = visit(context.properties().property(), Property.class);
517517
for (Property property : propertyList) {
@@ -3166,7 +3166,10 @@ public ParseNode visitSetWarehouseStatement(StarRocksParser.SetWarehouseStatemen
31663166

31673167
@Override
31683168
public ParseNode visitExecuteScriptStatement(StarRocksParser.ExecuteScriptStatementContext context) {
3169-
long beId = Long.parseLong(context.INTEGER_VALUE().getText());
3169+
long beId = -1;
3170+
if (context.INTEGER_VALUE() != null) {
3171+
beId = Long.parseLong(context.INTEGER_VALUE().getText());
3172+
}
31703173
StringLiteral stringLiteral = (StringLiteral) visit(context.string());
31713174
String script = stringLiteral.getStringValue();
31723175
return new ExecuteScriptStmt(beId, script, createPos(context));

‎fe/fe-core/src/main/java/com/starrocks/sql/parser/StarRocks.g4

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1522,7 +1522,7 @@ setWarehouseStatement
15221522
;
15231523

15241524
executeScriptStatement
1525-
: ADMIN EXECUTE ON INTEGER_VALUE string
1525+
: ADMIN EXECUTE ON (FRONTEND | INTEGER_VALUE) string
15261526
;
15271527

15281528
unsupportedStatement

‎fe/pom.xml

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -287,7 +287,7 @@ under the License.
287287
<artifactId>starrocks-bdb-je</artifactId>
288288
<version>18.3.13</version>
289289
</dependency>
290-
290+
291291
<!-- https://mvnrepository.com/artifact/de.jflex/jflex -->
292292
<dependency>
293293
<groupId>de.jflex</groupId>
@@ -704,6 +704,13 @@ under the License.
704704
<version>2.9.0</version>
705705
</dependency>
706706

707+
<!-- https://mvnrepository.com/artifact/org.apache.groovy/groovy-groovysh -->
708+
<dependency>
709+
<groupId>org.apache.groovy</groupId>
710+
<artifactId>groovy-groovysh</artifactId>
711+
<version>4.0.9</version>
712+
</dependency>
713+
707714
</dependencies>
708715
</dependencyManagement>
709716

0 commit comments

Comments
 (0)