udeServer/internal/audit/archive.go

package audit

import (
	"bufio"
	"context"
	"database/sql"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"
	"time"

	"udeServer/internal/model"

	"github.com/zeromicro/go-zero/core/stores/sqlx"
)

// LogArchive builds a temporary SQL dump file with one INSERT statement per
// log entry and returns the temp file together with the MinIO object key.
func LogArchive(logs []*model.TSystemLog, archiveName string) (*os.File, string, error) {
	tmpDir := os.TempDir()
	filename := fmt.Sprintf("%s_%d.sql", archiveName, time.Now().UnixNano())
	tmpPath := filepath.Join(tmpDir, filename)
	tmpFile, err := os.Create(tmpPath)
	objectKey := fmt.Sprintf("%s/%s", ObjectDirArchives, filename)
	if err != nil {
		return nil, "", fmt.Errorf("create temp file failed: %w", err)
	}
	defer func(tmpFile *os.File) {
		_ = tmpFile.Close()
	}(tmpFile)

	for _, entry := range logs {
		stmt := getLogStmt(entry)
		if _, err := tmpFile.WriteString(stmt); err != nil {
			return nil, "", fmt.Errorf("writing to file failed: %w", err)
		}
	}

	return tmpFile, objectKey, nil
}

// getLogStmt renders one INSERT statement for the entry. The table is left
// unqualified so restoring works against whatever schema the DSN points to.
func getLogStmt(entry *model.TSystemLog) string {
	return fmt.Sprintf(
		"INSERT INTO `t_system_log` ("+
			"`id`, `iden`, `trace_id`, `span_id`, `parent_span_id`, "+
			"`type`, `service_name`, "+
			"`uri_path`, `http_method`, `method`, `level`, "+
			"`app_key`, `user_id`, `user_name`, `client_ip`, "+
			"`request_header`, `request_body_digest`, `request_body_file_key`, "+
			"`response_code`, `response_body_digest`, `response_body_file_key`, "+
			"`error_message`, `duration_ms`, "+
			"`created_at`, `h_mac`, `audit_status`) VALUES ("+
			"%d, '%s', '%s', '%s', %s, "+
			"%d, '%s', "+
			"%s, %s, '%s', '%s', "+
			"%s, %s, %s, %s, "+
			"%s, '%s', '%s', "+
			"%d, '%s', '%s', "+
			"%s, %d, "+
			"'%s', '%s', %d);\n",
		entry.Id, entry.Iden, entry.TraceId, entry.SpanId, nullOrQuotedString(entry.ParentSpanId),
		entry.Type, entry.ServiceName,
		nullOrQuotedString(entry.UriPath), nullOrQuotedString(entry.HttpMethod), entry.Method, entry.Level,
		nullOrQuotedString(entry.AppKey), nullOrQuotedString(entry.UserId), nullOrQuotedString(entry.UserName), nullOrQuotedString(entry.ClientIp),
		nullOrQuotedString(entry.RequestHeader), entry.RequestBodyDigest, entry.RequestBodyFileKey,
		entry.ResponseCode, entry.ResponseBodyDigest, entry.ResponseBodyFileKey,
		nullOrQuotedString(entry.ErrorMessage), entry.DurationMs,
		entry.CreatedAt.Format("2006-01-02 15:04:05"), entry.HMac, entry.AuditStatus,
	)
}

func nullOrQuotedString(ns sql.NullString) string {
	if !ns.Valid {
		return "NULL"
	}
	s := ns.String
	if len(s) > 1 && s[0] == '\'' && s[len(s)-1] == '\'' {
		return s
	}
	// 转义单引号为两个单引号，防止SQL注入或语法错误
	safeStr := strings.ReplaceAll(s, "'", "''")
	return "'" + safeStr + "'"
}

// RestoreSQLStream streams the SQL dump file back into the database,
// executing batches of statements through execFunc.
func RestoreSQLStream(ctx context.Context, filePath string, batchSize int, execFunc func(context.Context, []string) error) error {
	// 打开文件
	file, err := os.Open(filePath)
	if err != nil {
		return err
	}
	defer func(file *os.File) {
		_ = file.Close()
	}(file)

	reader := bufio.NewReader(file)
	var buffer strings.Builder
	var batch []string

	for {
		line, err := reader.ReadString('\n')
		if err != nil && err != io.EOF {
			return err
		}

		buffer.WriteString(line)

		if strings.HasSuffix(strings.TrimSpace(line), ";") {
			stmt := strings.TrimSpace(buffer.String())
			if stmt != "" {
				batch = append(batch, stmt)
			}
			buffer.Reset()
		}

		// 到达批量大小时,执行一次批量插入
		if len(batch) >= batchSize {
			if err := execFunc(ctx, batch); err != nil {
				return err
			}
			// 清空数组
			batch = batch[:0]
		}

		if err == io.EOF {
			break
		}
	}

	// 补全插入剩下的数据
	if len(batch) > 0 {
		return execFunc(ctx, batch)
	}

	return nil
}

// ArchiveLogsToStorageAndRecord uploads the log dump to MinIO and, inside
// one transaction, deletes the archived rows and records the archive run.
// The MinIO client is guarded: when it is not loaded the caller should not
// invoke this function (checked by the business logic).
func ArchiveLogsToStorageAndRecord(
	ctx context.Context,
	logs []*model.TSystemLog,
	archiveName string,
	threshold int64,
	mode int64,
	minioClient *MinioClient,
	dbConn sqlx.SqlConn,
) (int, error) {
	// 构建归档文件
	tmpFile, objectKey, err := LogArchive(logs, archiveName)
	if err != nil {
		return 0, fmt.Errorf("日志归档文件生成失败: %w", err)
	}

	// 删除临时文件
	defer func(tmpFile *os.File) {
		_ = os.Remove(tmpFile.Name())
	}(tmpFile)

	// 上传
	fileKey, err := minioClient.UploadFile(ctx, tmpFile.Name(), objectKey)
	if err != nil {
		return 0, fmt.Errorf("归档文件上传失败: %w", err)
	}

	// 提取 ID
	var ids []int64
	for _, entry := range logs {
		ids = append(ids, entry.Id)
	}

	const deleteBatchSize = 1000

	// 入库事务
	err = dbConn.TransactCtx(ctx, func(ctx context.Context, session sqlx.Session) error {
		// 删除日志
		for i := 0; i < len(ids); i += deleteBatchSize {
			end := i + deleteBatchSize
			if end > len(ids) {
				end = len(ids)
			}
			batchIds := ids[i:end]
			placeholders := strings.Repeat("?,", len(batchIds))
			placeholders = strings.TrimSuffix(placeholders, ",")

			query := fmt.Sprintf("DELETE FROM t_system_log WHERE id IN (%s)", placeholders)

			args := make([]interface{}, len(batchIds))
			for j, v := range batchIds {
				args[j] = v
			}

			if _, err := session.ExecCtx(ctx, query, args...); err != nil {
				return err
			}
		}

		// 插入记录
		archivedBy := operatorAdmin
		archivedAt := time.Now()
		_, err := session.ExecCtx(ctx,
			"INSERT INTO t_log_archive_record (archive_name, record_count, archive_file_key, archive_mode, batch_size, archived_by, archived_at) VALUES (?,?,?,?,?,?,?)",
			archiveName, int64(len(logs)), fileKey, mode, threshold, archivedBy, archivedAt)
		return err
	})

	return len(logs), err
}