go-db/factory/doris_factory.go

package factory

import (
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strings"
	"time"

	"git.x2erp.com/qdy/go-base/config"
)

// DorisFactory Doris HTTP客户端工厂
type DorisFactory struct {
	config     config.IConfig
	httpClient *http.Client
}

// NewDorisFactory 创建Doris HTTP客户端工厂
func NewDorisFactory(httpFactory *HTTPFactory) (*DorisFactory, error) {
	cfg := config.GetConfig()

	if err := config.GetInitError(); err != nil {
		return nil, fmt.Errorf("failed to load config: %v", err)
	}

	if !cfg.IsDorisConfigured() {
		return nil, fmt.Errorf("doris configuration is incomplete")
	}

	return &DorisFactory{
		config:     cfg,
		httpClient: httpFactory.CreateHTTPClient(),
	}, nil
}

// CreateDorisClient 创建Doris HTTP客户端
func (f *DorisFactory) CreateDorisClient() *DorisClient {
	dorisConfig := f.config.GetDoris()

	return &DorisClient{
		httpClient: f.httpClient,
		FEHost:     dorisConfig.FEHost,
		FEPort:     dorisConfig.FEPort,
		Username:   dorisConfig.FEUsername,
		Password:   dorisConfig.FEPassword,
		Timeout:    time.Duration(dorisConfig.StreamLoadTimeout) * time.Second,
	}
}

// DorisClient Doris HTTP客户端
type DorisClient struct {
	httpClient *http.Client
	FEHost     string
	FEPort     int
	Username   string
	Password   string
	Timeout    time.Duration
}

// InsertCSV 插入CSV数据到Doris 只认数据次序
func (c *DorisClient) DorisInsertCSV(database, table, csvData string, skipHeader bool) error {
	url := fmt.Sprintf("http://%s:%d/api/%s/%s/_stream_load", c.FEHost, c.FEPort, database, table)

	req, err := http.NewRequest("PUT", url, strings.NewReader(csvData))
	if err != nil {
		return fmt.Errorf("创建请求失败: %v", err)
	}

	req.SetBasicAuth(c.Username, c.Password)
	req.Header.Set("Content-Type", "text/plain")
	req.Header.Set("format", "csv")
	req.Header.Set("column_separator", ",")

	if skipHeader {
		req.Header.Set("skip_header", "1")
	}

	resp, err := c.httpClient.Do(req)
	if err != nil {
		return fmt.Errorf("请求失败: %v", err)
	}
	defer resp.Body.Close()

	body, _ := io.ReadAll(resp.Body)
	if resp.StatusCode != 200 {
		return fmt.Errorf("插入失败: %s", string(body))
	}

	// 解析Stream Load响应
	var result struct {
		Status  string `json:"Status"`
		Message string `json:"Message"`
	}
	if err := json.Unmarshal(body, &result); err != nil {
		return fmt.Errorf("解析响应失败: %v", err)
	}

	if result.Status != "Success" {
		return fmt.Errorf("插入失败: %s", result.Message)
	}

	return nil
}