waf-platform/EdgeAdmin/internal/web/actions/default/db/clickhouse.go

//go:build plus

package db

import (
	"crypto/tls"
	"encoding/json"
	"fmt"
	"net/http"
	"strings"
	"time"

	"github.com/TeaOSLab/EdgeAdmin/internal/web/actions/actionutils"
	"github.com/TeaOSLab/EdgeCommon/pkg/langs/codes"
	"github.com/TeaOSLab/EdgeCommon/pkg/rpc/pb"
	"github.com/TeaOSLab/EdgeCommon/pkg/systemconfigs"
	"github.com/iwind/TeaGo/actions"
)

const clickhouseConfigCode = "clickhouseConfig"

type ClickHouseAction struct {
	actionutils.ParentAction
}

func (this *ClickHouseAction) Init() {
	this.Nav("db", "db", "clickhouse")
}

func (this *ClickHouseAction) RunGet(params struct{}) {
	this.Data["mainTab"] = "clickhouse"
	resp, err := this.RPC().SysSettingRPC().ReadSysSetting(this.AdminContext(), &pb.ReadSysSettingRequest{Code: clickhouseConfigCode})
	if err != nil {
		this.ErrorPage(err)
		return
	}
	cfg := &systemconfigs.ClickHouseSetting{Port: 8443, Database: "default", Scheme: "https"}
	if len(resp.ValueJSON) > 0 {
		_ = json.Unmarshal(resp.ValueJSON, cfg)
	}
	if cfg.Port <= 0 {
		cfg.Port = 8443
	}
	if cfg.Database == "" {
		cfg.Database = "default"
	}
	if strings.TrimSpace(cfg.Scheme) == "" {
		cfg.Scheme = "https"
	}
	this.Data["config"] = map[string]interface{}{
		"host":     cfg.Host,
		"port":     cfg.Port,
		"user":     cfg.User,
		"password": cfg.Password,
		"database": cfg.Database,
		"scheme":   cfg.Scheme,
	}

	// 自动检测连接状态
	connStatus := "unconfigured" // unconfigured / connected / disconnected
	connError := ""
	if strings.TrimSpace(cfg.Host) != "" {
		connStatus, connError = this.probeClickHouse(cfg)
	}
	this.Data["connStatus"] = connStatus
	this.Data["connError"] = connError

	this.Show()
}

func (this *ClickHouseAction) RunPost(params struct {
	Host          string
	Port          int
	User          string
	Password      string
	Database      string
	Scheme        string

	Must *actions.Must
}) {
	defer this.CreateLogInfo(codes.DBNode_LogUpdateDBNode, 0)
	if params.Database == "" {
		params.Database = "default"
	}
	if params.Scheme != "http" {
		params.Scheme = "https"
	}
	if params.Port <= 0 {
		params.Port = 8443
	}
	password := params.Password
	if password == "" {
		resp, _ := this.RPC().SysSettingRPC().ReadSysSetting(this.AdminContext(), &pb.ReadSysSettingRequest{Code: clickhouseConfigCode})
		if len(resp.ValueJSON) > 0 {
			var old systemconfigs.ClickHouseSetting
			if json.Unmarshal(resp.ValueJSON, &old) == nil {
				password = old.Password
			}
		}
	}
	cfg := &systemconfigs.ClickHouseSetting{
		Host:          params.Host,
		Port:          params.Port,
		User:          params.User,
		Password:      password,
		Database:      params.Database,
		Scheme:        params.Scheme,
		TLSSkipVerify: true,
		TLSServerName: "",
	}
	valueJSON, err := json.Marshal(cfg)
	if err != nil {
		this.ErrorPage(err)
		return
	}
	_, err = this.RPC().SysSettingRPC().UpdateSysSetting(this.AdminContext(), &pb.UpdateSysSettingRequest{
		Code:      clickhouseConfigCode,
		ValueJSON: valueJSON,
	})
	if err != nil {
		this.ErrorPage(err)
		return
	}
	this.Success()
}

// probeClickHouse 快速检测 ClickHouse 连接状态（SELECT 1）
func (this *ClickHouseAction) probeClickHouse(cfg *systemconfigs.ClickHouseSetting) (status string, errMsg string) {
	scheme := strings.ToLower(strings.TrimSpace(cfg.Scheme))
	if scheme == "" {
		scheme = "https"
	}
	port := cfg.Port
	if port <= 0 {
		port = 8443
	}
	db := cfg.Database
	if db == "" {
		db = "default"
	}

	testURL := fmt.Sprintf("%s://%s:%d/?query=SELECT+1&database=%s", scheme, cfg.Host, port, db)

	transport := &http.Transport{}
	if scheme == "https" {
		transport.TLSClientConfig = &tls.Config{
			InsecureSkipVerify: true,
		}
	}
	client := &http.Client{
		Timeout:   3 * time.Second,
		Transport: transport,
	}

	req, err := http.NewRequest(http.MethodGet, testURL, nil)
	if err != nil {
		return "disconnected", err.Error()
	}
	if cfg.User != "" || cfg.Password != "" {
		req.SetBasicAuth(cfg.User, cfg.Password)
	}

	resp, err := client.Do(req)
	if err != nil {
		return "disconnected", err.Error()
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return "disconnected", fmt.Sprintf("HTTP %d", resp.StatusCode)
	}
	return "connected", ""
}