ylrz_saas_pharmacy/java/scripts/ApplyLobsterWorkflowLobsterColumnsMigration.java

import java.sql.*;
import java.util.ArrayList;
import java.util.List;

/**
 * Add missing company_workflow_lobster execution/template columns on tenant DB(s).
 * Usage: java -cp mysql-connector-j.jar ApplyLobsterWorkflowLobsterColumnsMigration [tenantId...]
 * Default tenantId: 33
 */
public class ApplyLobsterWorkflowLobsterColumnsMigration {

    private static final String MASTER_HOST = getenv("DB_HOST", "cq-cdb-8fjmemkb.sql.tencentcdb.com");
    private static final int MASTER_PORT = Integer.parseInt(getenv("DB_PORT", "27220"));
    private static final String MASTER_USER = getenv("DB_USER", "root");
    private static final String MASTER_PASSWORD = getenv("DB_PASSWORD", "Ylrz_1q2w3e4r5t6y");
    private static final String MASTER_DB = getenv("DB_NAME", "ylrz_saas");

    public static void main(String[] args) throws Exception {
        if (args.length == 1 && ("master".equalsIgnoreCase(args[0]) || "all".equalsIgnoreCase(args[0]))) {
            patchMaster();
            if ("master".equalsIgnoreCase(args[0])) {
                return;
            }
            args = new String[] { "all-tenants" };
        }

        List<Long> tenantIds = new ArrayList<>();
        if (args.length == 0) {
            tenantIds.add(33L);
        } else if (args.length == 1 && "all-tenants".equalsIgnoreCase(args[0])) {
            try (Connection master = open(MASTER_HOST, MASTER_PORT, MASTER_DB, MASTER_USER, MASTER_PASSWORD)) {
                try (Statement st = master.createStatement();
                     ResultSet rs = st.executeQuery("SELECT id FROM tenant_info WHERE status = 1 ORDER BY id")) {
                    while (rs.next()) {
                        tenantIds.add(rs.getLong("id"));
                    }
                }
            }
        } else {
            for (String arg : args) {
                tenantIds.add(Long.parseLong(arg.trim()));
            }
        }

        try (Connection master = open(MASTER_HOST, MASTER_PORT, MASTER_DB, MASTER_USER, MASTER_PASSWORD)) {
            for (Long tenantId : tenantIds) {
                TenantDb tenant = loadTenant(master, tenantId);
                if (tenant == null) {
                    System.err.println("[SKIP] tenant " + tenantId + " not found");
                    continue;
                }
                System.out.println("[RUN] tenant " + tenantId + " -> " + tenant.host + ":" + tenant.port + "/" + tenant.database);
                try (Connection tenantConn = open(tenant.host, tenant.port, tenant.database, tenant.user, tenant.password)) {
                    applyLobsterExecutionConfigTable(tenantConn);
                    applyWorkflowLobsterColumns(tenantConn);
                    applyTagBindingColumns(tenantConn);
                }
                System.out.println("[OK] tenant " + tenantId);
            }
        }
    }

    private static void patchMaster() throws SQLException {
        System.out.println("[RUN] master -> " + MASTER_DB);
        try (Connection conn = open(MASTER_HOST, MASTER_PORT, MASTER_DB, MASTER_USER, MASTER_PASSWORD)) {
            if (tableExists(conn, "company_workflow_lobster")) {
                applyWorkflowLobsterColumns(conn);
            } else {
                System.out.println("  skip: company_workflow_lobster not on master");
            }
            applyTagBindingColumns(conn);
        }
        System.out.println("[OK] master");
    }

    private static void applyLobsterExecutionConfigTable(Connection conn) throws SQLException {
        if (tableExists(conn, "lobster_execution_config")) {
            System.out.println("  table exists: lobster_execution_config");
            return;
        }
        try (Statement st = conn.createStatement()) {
            st.execute(
                    "CREATE TABLE IF NOT EXISTS `lobster_execution_config` ("
                            + "`id` bigint NOT NULL AUTO_INCREMENT,"
                            + "`company_id` bigint NOT NULL COMMENT 'tenant company id',"
                            + "`default_execution_mode` varchar(16) NOT NULL DEFAULT 'HYBRID',"
                            + "`enable_content_personalization` tinyint NOT NULL DEFAULT 1,"
                            + "`enable_flow_personalization` tinyint NOT NULL DEFAULT 1,"
                            + "`strict_fixed_workflow` tinyint NOT NULL DEFAULT 0,"
                            + "`auto_start_instance_on_tag` tinyint NOT NULL DEFAULT 0,"
                            + "`create_by` varchar(64) DEFAULT NULL,"
                            + "`create_time` datetime DEFAULT CURRENT_TIMESTAMP,"
                            + "`update_by` varchar(64) DEFAULT NULL,"
                            + "`update_time` datetime DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,"
                            + "PRIMARY KEY (`id`),"
                            + "UNIQUE KEY `uk_company_id` (`company_id`)"
                            + ") ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COMMENT='lobster execution policy'"
            );
            System.out.println("  created table: lobster_execution_config");
        }
    }

    private static void applyWorkflowLobsterColumns(Connection conn) throws SQLException {
        addColumnIfMissing(conn, "company_workflow_lobster", "execution_mode",
                "ALTER TABLE company_workflow_lobster ADD COLUMN execution_mode varchar(16) DEFAULT NULL "
                        + "COMMENT 'PROMPT/INSTANCE/HYBRID' AFTER canvas_data");
        addColumnIfMissing(conn, "company_workflow_lobster", "enable_content_personalization",
                "ALTER TABLE company_workflow_lobster ADD COLUMN enable_content_personalization tinyint DEFAULT NULL "
                        + "COMMENT 'null=inherit tenant' AFTER execution_mode");
        addColumnIfMissing(conn, "company_workflow_lobster", "enable_flow_personalization",
                "ALTER TABLE company_workflow_lobster ADD COLUMN enable_flow_personalization tinyint DEFAULT NULL "
                        + "COMMENT 'null=inherit tenant' AFTER enable_content_personalization");
        addColumnIfMissing(conn, "company_workflow_lobster", "strict_fixed_workflow",
                "ALTER TABLE company_workflow_lobster ADD COLUMN strict_fixed_workflow tinyint DEFAULT NULL "
                        + "COMMENT 'null=inherit tenant' AFTER enable_flow_personalization");
        addColumnIfMissing(conn, "company_workflow_lobster", "auto_start_instance_on_tag",
                "ALTER TABLE company_workflow_lobster ADD COLUMN auto_start_instance_on_tag tinyint DEFAULT NULL "
                        + "COMMENT 'null=inherit tenant' AFTER strict_fixed_workflow");
        addColumnIfMissing(conn, "company_workflow_lobster", "template_kind",
                "ALTER TABLE company_workflow_lobster ADD COLUMN template_kind varchar(20) NOT NULL DEFAULT 'MASTER' "
                        + "COMMENT 'MASTER or TASK_COPY' AFTER auto_start_instance_on_tag");
        addColumnIfMissing(conn, "company_workflow_lobster", "source_template_id",
                "ALTER TABLE company_workflow_lobster ADD COLUMN source_template_id bigint DEFAULT NULL "
                        + "COMMENT 'master template id when TASK_COPY' AFTER template_kind");
        try (Statement st = conn.createStatement()) {
            st.execute("UPDATE company_workflow_lobster SET template_kind = 'MASTER' "
                    + "WHERE template_kind IS NULL OR template_kind = ''");
        }
    }

    private static void applyTagBindingColumns(Connection conn) throws SQLException {
        if (!tableExists(conn, "company_tag_template_binding")) {
            try (Statement st = conn.createStatement()) {
                st.execute(
                        "CREATE TABLE IF NOT EXISTS company_tag_template_binding ("
                                + "id bigint NOT NULL AUTO_INCREMENT,"
                                + "company_id bigint NOT NULL,"
                                + "tag_code varchar(128) NOT NULL,"
                                + "tag_name varchar(128) NOT NULL,"
                                + "template_id bigint NOT NULL,"
                                + "template_name varchar(256) DEFAULT NULL,"
                                + "priority int DEFAULT 0,"
                                + "match_condition text,"
                                + "exclude_tags varchar(2000) DEFAULT NULL,"
                                + "target_channels varchar(500) DEFAULT NULL,"
                                + "status tinyint DEFAULT 1,"
                                + "del_flag tinyint DEFAULT 0,"
                                + "create_by varchar(64) DEFAULT '',"
                                + "create_time datetime DEFAULT NULL,"
                                + "update_by varchar(64) DEFAULT '',"
                                + "update_time datetime DEFAULT NULL,"
                                + "PRIMARY KEY (id),"
                                + "KEY idx_cttb_company_del (company_id, del_flag)"
                                + ") ENGINE=InnoDB DEFAULT CHARSET=utf8mb4"
                );
                System.out.println("  created table: company_tag_template_binding");
            }
        }
        if (!tableExists(conn, "company_tag_template_binding")) {
            return;
        }
        addColumnIfMissing(conn, "company_tag_template_binding", "source_template_id",
                "ALTER TABLE company_tag_template_binding ADD COLUMN source_template_id bigint DEFAULT NULL "
                        + "COMMENT 'master template id' AFTER template_id");
        addColumnIfMissing(conn, "company_tag_template_binding", "task_template_id",
                "ALTER TABLE company_tag_template_binding ADD COLUMN task_template_id bigint DEFAULT NULL "
                        + "COMMENT 'task copy template id' AFTER source_template_id");
        addColumnIfMissing(conn, "company_tag_template_binding", "exclude_tags",
                "ALTER TABLE company_tag_template_binding ADD COLUMN exclude_tags varchar(2000) DEFAULT NULL "
                        + "COMMENT 'exclude tag codes' AFTER match_condition");
        addColumnIfMissing(conn, "company_tag_template_binding", "target_channels",
                "ALTER TABLE company_tag_template_binding ADD COLUMN target_channels varchar(500) DEFAULT NULL "
                        + "COMMENT 'target channels' AFTER exclude_tags");
        try (Statement st = conn.createStatement()) {
            st.execute("UPDATE company_tag_template_binding SET source_template_id = template_id "
                    + "WHERE source_template_id IS NULL AND template_id IS NOT NULL");
        }
    }

    private static void addColumnIfMissing(Connection conn, String table, String column, String ddl) throws SQLException {
        if (columnExists(conn, table, column)) {
            System.out.println("  column exists: " + table + "." + column);
            return;
        }
        try (Statement st = conn.createStatement()) {
            st.execute(ddl);
            System.out.println("  added column: " + table + "." + column);
        }
    }

    private static boolean tableExists(Connection conn, String table) throws SQLException {
        String sql = "SELECT COUNT(*) FROM information_schema.TABLES "
                + "WHERE TABLE_SCHEMA = DATABASE() AND TABLE_NAME = ?";
        try (PreparedStatement ps = conn.prepareStatement(sql)) {
            ps.setString(1, table);
            try (ResultSet rs = ps.executeQuery()) {
                return rs.next() && rs.getInt(1) > 0;
            }
        }
    }

    private static boolean columnExists(Connection conn, String table, String column) throws SQLException {
        String sql = "SELECT COUNT(*) FROM information_schema.COLUMNS "
                + "WHERE TABLE_SCHEMA = DATABASE() AND TABLE_NAME = ? AND COLUMN_NAME = ?";
        try (PreparedStatement ps = conn.prepareStatement(sql)) {
            ps.setString(1, table);
            ps.setString(2, column);
            try (ResultSet rs = ps.executeQuery()) {
                return rs.next() && rs.getInt(1) > 0;
            }
        }
    }

    private static TenantDb loadTenant(Connection master, long tenantId) throws SQLException {
        String sql = "SELECT db_url, db_account, db_pwd FROM tenant_info WHERE id = ? AND status = 1 LIMIT 1";
        try (PreparedStatement ps = master.prepareStatement(sql)) {
            ps.setLong(1, tenantId);
            try (ResultSet rs = ps.executeQuery()) {
                if (!rs.next()) {
                    return null;
                }
                return parseJdbcUrl(rs.getString("db_url"), rs.getString("db_account"), rs.getString("db_pwd"));
            }
        }
    }

    private static TenantDb parseJdbcUrl(String jdbcUrl, String user, String password) {
        if (jdbcUrl == null || !jdbcUrl.startsWith("jdbc:mysql://")) {
            throw new IllegalArgumentException("Unsupported db_url: " + jdbcUrl);
        }
        String body = jdbcUrl.substring("jdbc:mysql://".length());
        int slash = body.indexOf('/');
        if (slash < 0) {
            throw new IllegalArgumentException("Invalid db_url: " + jdbcUrl);
        }
        String hostPort = body.substring(0, slash);
        String rest = body.substring(slash + 1);
        int q = rest.indexOf('?');
        String database = q >= 0 ? rest.substring(0, q) : rest;
        String host;
        int port;
        int colon = hostPort.indexOf(':');
        if (colon >= 0) {
            host = hostPort.substring(0, colon);
            port = Integer.parseInt(hostPort.substring(colon + 1));
        } else {
            host = hostPort;
            port = 3306;
        }
        TenantDb t = new TenantDb();
        t.host = host;
        t.port = port;
        t.database = database;
        t.user = user != null ? user : MASTER_USER;
        t.password = password != null ? password : MASTER_PASSWORD;
        return t;
    }

    private static Connection open(String host, int port, String database, String user, String password)
            throws SQLException {
        String url = "jdbc:mysql://" + host + ":" + port + "/" + database
                + "?useUnicode=true&characterEncoding=utf8&serverTimezone=GMT%2B8&allowMultiQueries=true";
        return DriverManager.getConnection(url, user, password);
    }

    private static String getenv(String key, String defaultValue) {
        String v = System.getenv(key);
        return v != null && !v.isEmpty() ? v : defaultValue;
    }

    private static final class TenantDb {
        String host;
        int port;
        String database;
        String user;
        String password;
    }
}