opsxdb/agents/jdbc/OpsxZkAgent.java

// OpsxDB ZooKeeper Agent —— 与 OpsxJdbcAgent 同协议的 ZK 专属桥（默认包，单文件）。
//
// 启动：java -cp opsx-agent-jdbc.jar:gson-2.12.1.jar OpsxZkAgent --jars <zookeeper.jar 目录>
//
// SQL 约定（工作台编辑器直写）：
//   GET <path>          读 znode：path/data/version/ctime/mtime/ephemeral/children
//   LS  <path>          列子节点（带 stat）
//   CREATE <path> [data]   持久节点（父须存在）；SET <path> <data>；DELETE <path>（须无子）
// 写动词（CREATE/SET/DELETE）由 Go 侧危险语句拦截 + 二次确认把守。

import com.google.gson.Gson;
import com.google.gson.JsonObject;
import com.google.gson.JsonParser;

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.io.OutputStreamWriter;
import java.io.PrintWriter;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;

import org.apache.zookeeper.CreateMode;
import org.apache.zookeeper.KeeperException;
import org.apache.zookeeper.WatchedEvent;
import org.apache.zookeeper.Watcher;
import org.apache.zookeeper.ZooDefs;
import org.apache.zookeeper.ZooKeeper;
import org.apache.zookeeper.data.Stat;

public final class OpsxZkAgent {
    static final String AGENT = "opsx-agent-jdbc";
    static final String VERSION = "0.1.0";
    static final int PROTOCOL = 1;

    private final Gson gson = new Gson();
    private final Map<String, ZooKeeper> handles = new HashMap<>();
    private final AtomicLong seq = new AtomicLong(0);

    public static void main(String[] args) throws Exception {
        new OpsxZkAgent().loop();
    }

    private void loop() throws Exception {
        BufferedReader in = new BufferedReader(new InputStreamReader(System.in, StandardCharsets.UTF_8));
        PrintWriter out = new PrintWriter(new OutputStreamWriter(System.out, StandardCharsets.UTF_8), true);
        out.println(gson.toJson(resp(0, true, helloResult(), null)));
        String line;
        while ((line = in.readLine()) != null) {
            line = line.trim();
            if (line.isEmpty()) continue;
            JsonObject req;
            try {
                req = JsonParser.parseString(line).getAsJsonObject();
            } catch (Exception e) {
                out.println(gson.toJson(resp(0, false, null, "bad json: " + e.getMessage())));
                continue;
            }
            long id = req.has("id") ? req.get("id").getAsLong() : 0;
            String op = req.has("op") ? req.get("op").getAsString() : "";
            JsonObject args = req.has("args") && req.get("args").isJsonObject() ? req.getAsJsonObject("args") : new JsonObject();
            try {
                out.println(gson.toJson(resp(id, true, dispatch(op, args), null)));
            } catch (KeeperException e) {
                out.println(gson.toJson(resp(id, false, null, "ZK " + e.code().name() + ": " + e.getMessage())));
            } catch (Exception e) {
                out.println(gson.toJson(resp(id, false, null, e.getMessage() == null ? String.valueOf(e) : e.getMessage())));
            }
            if ("shutdown".equals(op)) {
                out.flush();
                System.exit(0);
            }
        }
        System.exit(0);
    }

    private JsonObject helloResult() {
        JsonObject o = new JsonObject();
        o.addProperty("agent", AGENT);
        o.addProperty("version", VERSION);
        o.addProperty("protocol", PROTOCOL);
        return o;
    }

    private JsonObject dispatch(String op, JsonObject a) throws Exception {
        switch (op) {
            case "hello": return helloResult();
            case "open": return open(a);
            case "ping": {
                Stat st = zk(a).exists("/", false);
                if (st == null) throw new IllegalStateException("root 不存在");
                return ok();
            }
            case "query": return query(a);
            case "exec": return exec(a);
            case "metaDatabases": return children(zk(a), "/");
            case "metaTables": {
                String schema = str(a, "schema");
                String base = schema.isEmpty() ? "/" : normalize(schema);
                return children(zk(a), base);
            }
            case "close": {
                ZooKeeper z = handles.remove(str(a, "handle"));
                if (z != null) z.close();
                return ok();
            }
            default: throw new IllegalArgumentException("未知 op: " + op);
        }
    }

    private JsonObject open(JsonObject a) throws Exception {
        String connect = str(a, "url"); // host:port
        if (connect.isEmpty()) throw new IllegalArgumentException("缺少 connectString");
        CountDownLatch latch = new CountDownLatch(1);
        ZooKeeper zk = new ZooKeeper(connect, 10000, new Watcher() {
            public void process(WatchedEvent e) {
                if (e.getState() == Event.KeeperState.SyncConnected
                        || e.getState() == Event.KeeperState.ConnectedReadOnly) {
                    latch.countDown();
                }
            }
        });
        if (!latch.await(10, TimeUnit.SECONDS)) {
            zk.close();
            throw new IllegalStateException("ZK 连接超时: " + connect);
        }
        String handle = "z" + seq.incrementAndGet();
        handles.put(handle, zk);
        JsonObject o = new JsonObject();
        o.addProperty("handle", handle);
        return o;
    }

    /** query 承载读动词：GET / LS */
    private JsonObject query(JsonObject a) throws Exception {
        ZooKeeper zk = zk(a);
        String sql = str(a, "sql").trim();
        String[] parts = sql.split("\\s+", 2);
        String verb = parts.length > 0 ? parts[0].toUpperCase() : "";
        String arg = parts.length > 1 ? parts[1].trim() : "/";
        if (verb.equals("GET")) {
            return statRow(zk, normalize(arg), false);
        }
        if (verb.equals("LS")) {
            return statRow(zk, normalize(arg), true);
        }
        throw new IllegalArgumentException("ZK query 仅支持 GET/LS，收到: " + verb);
    }

    /** exec 承载写动词：CREATE / SET / DELETE */
    private JsonObject exec(JsonObject a) throws Exception {
        ZooKeeper zk = zk(a);
        String sql = str(a, "sql").trim();
        String[] parts = sql.split("\\s+", 2);
        String verb = parts.length > 0 ? parts[0].toUpperCase() : "";
        String rest = parts.length > 1 ? parts[1] : "";
        int sp = rest.indexOf(' ');
        String path = normalize(sp < 0 ? rest : rest.substring(0, sp));
        String data = sp < 0 ? "" : rest.substring(sp + 1);
        switch (verb) {
            case "CREATE":
                zk.create(path, data.getBytes(StandardCharsets.UTF_8), ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
                break;
            case "SET":
                zk.setData(path, data.getBytes(StandardCharsets.UTF_8), -1);
                break;
            case "DELETE":
                zk.delete(path, -1);
                break;
            default:
                throw new IllegalArgumentException("ZK exec 仅支持 CREATE/SET/DELETE，收到: " + verb);
        }
        JsonObject o = new JsonObject();
        o.addProperty("affectedRows", 1);
        return o;
    }

    /** statRow：GET 单行 / LS 多行（每个子节点一行） */
    private JsonObject statRow(ZooKeeper zk, String path, boolean listChildren) throws Exception {
        List<JsonObject> rows = new ArrayList<>();
        if (listChildren) {
            for (String child : zk.getChildren(path, false)) {
                String childPath = "/".equals(path) ? "/" + child : path + "/" + child;
                rows.add(row(zk, childPath));
            }
        } else {
            rows.add(row(zk, path));
        }
        JsonObject o = new JsonObject();
        o.add("columns", gson.toJsonTree(new String[]{"path", "data", "version", "ctime", "mtime", "ephemeral", "children"}));
        o.add("rows", gson.toJsonTree(rows));
        o.addProperty("truncated", false);
        return o;
    }

    private JsonObject row(ZooKeeper zk, String path) throws Exception {
        Stat st = new Stat();
        byte[] data = zk.getData(path, false, st);
        JsonObject r = new JsonObject();
        r.addProperty("path", path);
        r.addProperty("data", data == null ? "" : new String(data, StandardCharsets.UTF_8));
        r.addProperty("version", st.getVersion());
        r.addProperty("ctime", st.getCtime());
        r.addProperty("mtime", st.getMtime());
        r.addProperty("ephemeral", st.getEphemeralOwner() != 0);
        r.addProperty("children", st.getNumChildren());
        return r;
    }

    private JsonObject children(ZooKeeper zk, String base) throws Exception {
        List<String> names = new ArrayList<>();
        if (zk.getChildren(base, false).isEmpty() && "/".equals(base)) {
            names.add("/");
        } else {
            names.addAll(zk.getChildren(base, false));
        }
        JsonObject o = new JsonObject();
        o.add("names", gson.toJsonTree(names));
        return o;
    }

    private ZooKeeper zk(JsonObject a) {
        ZooKeeper z = handles.get(str(a, "handle"));
        if (z == null) throw new IllegalStateException("句柄不存在或已关闭");
        return z;
    }

    private static String normalize(String path) {
        if (path == null || path.isEmpty()) return "/";
        return path.startsWith("/") ? path : "/" + path;
    }

    private static String str(JsonObject a, String k) {
        return a.has(k) && !a.get(k).isJsonNull() ? a.get(k).getAsString() : "";
    }

    private static JsonObject ok() { return new JsonObject(); }

    private static JsonObject resp(long id, boolean ok, JsonObject result, String error) {
        JsonObject o = new JsonObject();
        o.addProperty("id", id);
        o.addProperty("ok", ok);
        if (result != null) o.add("result", result);
        if (error != null) o.addProperty("error", error);
        return o;
    }
}