package database import ( "bufio" "context" "encoding/json" "fmt" "io" "os" "time" "github.com/uptrace/bun" ) // JSONLImporter handles batch import from JSONL files type JSONLImporter struct { db *bun.DB batchSize int } // NewJSONLImporter creates a new JSONL importer func NewJSONLImporter(db *bun.DB) *JSONLImporter { return &JSONLImporter{ db: db, batchSize: 100, } } // WithBatchSize sets the batch size for imports func (i *JSONLImporter) WithBatchSize(size int) *JSONLImporter { if size > 0 { i.batchSize = size } return i } // ImportResult holds import statistics type ImportResult struct { Total int `json:"total"` Imported int `json:"imported"` Updated int `json:"updated"` Failed int `json:"failed"` Errors []ImportError `json:"errors,omitempty"` Duration time.Duration `json:"duration"` } // ImportError represents a single import error type ImportError struct { Line int `json:"line"` Error string `json:"error"` Data string `json:"data,omitempty"` } // ImportAssets imports assets from a JSONL file func (i *JSONLImporter) ImportAssets(ctx context.Context, filePath, workspace, source string) (*ImportResult, error) { file, err := os.Open(filePath) if err != nil { return nil, fmt.Errorf("failed to open file: %w", err) } defer func() { _ = file.Close() }() return i.ImportAssetsFromReader(ctx, file, workspace, source) } // ImportAssetsFromReader imports assets from an io.Reader func (i *JSONLImporter) ImportAssetsFromReader(ctx context.Context, r io.Reader, workspace, source string) (*ImportResult, error) { startTime := time.Now() scanner := bufio.NewScanner(r) // Allow large lines (up to 10MB) scanner.Buffer(make([]byte, 1024*1024), 10*1024*1024) result := &ImportResult{} batch := make([]*Asset, 0, i.batchSize) for scanner.Scan() { result.Total++ line := scanner.Bytes() // Skip empty lines if len(line) == 0 { continue } asset, err := ParseAssetLine(line, workspace, source) if err != nil { result.Failed++ result.Errors = append(result.Errors, ImportError{ Line: result.Total, Error: err.Error(), Data: truncateString(string(line), 200), }) continue } batch = append(batch, asset) if len(batch) >= i.batchSize { imported, err := i.insertAssetBatch(ctx, batch) if err != nil { return result, fmt.Errorf("batch insert failed at line %d: %w", result.Total, err) } result.Imported += imported batch = batch[:0] } } // Insert remaining batch if len(batch) > 0 { imported, err := i.insertAssetBatch(ctx, batch) if err != nil { return result, fmt.Errorf("final batch insert failed: %w", err) } result.Imported += imported } if err := scanner.Err(); err != nil { return result, fmt.Errorf("scanner error: %w", err) } result.Duration = time.Since(startTime) return result, nil } // insertAssetBatch inserts a batch of assets with upsert func (i *JSONLImporter) insertAssetBatch(ctx context.Context, assets []*Asset) (int, error) { if len(assets) == 0 { return 0, nil } // Use ON CONFLICT for upsert res, err := i.db.NewInsert(). Model(&assets). On("CONFLICT (workspace, asset_value, url) DO UPDATE"). Set("status_code = EXCLUDED.status_code"). Set("title = EXCLUDED.title"). Set("tech = EXCLUDED.tech"). Set("content_type = EXCLUDED.content_type"). Set("content_length = EXCLUDED.content_length"). Set("host_ip = EXCLUDED.host_ip"). Set("a_records = EXCLUDED.a_records"). Set("tls = EXCLUDED.tls"). Set("response_time = EXCLUDED.response_time"). Set("words = EXCLUDED.words"). Set("lines = EXCLUDED.lines"). Set("remarks = EXCLUDED.remarks"). Set("raw_data = EXCLUDED.raw_data"). Set("updated_at = EXCLUDED.updated_at"). Exec(ctx) if err != nil { return 0, err } rowsAffected, _ := res.RowsAffected() return int(rowsAffected), nil } // ParseAssetLine parses a single JSONL line into an Asset func ParseAssetLine(line []byte, defaultWorkspace, source string) (*Asset, error) { var raw map[string]interface{} if err := json.Unmarshal(line, &raw); err != nil { return nil, fmt.Errorf("invalid JSON: %w", err) } now := time.Now() asset := &Asset{ Workspace: defaultWorkspace, Source: source, RawJsonData: string(line), CreatedAt: now, UpdatedAt: now, } // Map JSON fields to Asset struct // Required fields if v, ok := raw["workspace"].(string); ok && v != "" { asset.Workspace = v } if v, ok := raw["asset_value"].(string); ok { asset.AssetValue = v } // HTTP data if v, ok := raw["url"].(string); ok { asset.URL = v } if v, ok := raw["input"].(string); ok { asset.Input = v } if v, ok := raw["scheme"].(string); ok { asset.Scheme = v } if v, ok := raw["method"].(string); ok { asset.Method = v } if v, ok := raw["path"].(string); ok { asset.Path = v } // Response data if v, ok := raw["status_code"].(float64); ok { asset.StatusCode = int(v) } if v, ok := raw["content_type"].(string); ok { asset.ContentType = v } if v, ok := raw["content_length"].(float64); ok { asset.ContentLength = int64(v) } if v, ok := raw["title"].(string); ok { asset.Title = v } if v, ok := raw["words"].(float64); ok { asset.Words = int(v) } if v, ok := raw["lines"].(float64); ok { asset.Lines = int(v) } // Network data if v, ok := raw["host_ip"].(string); ok { asset.HostIP = v } if v, ok := raw["a"].([]interface{}); ok { asset.DnsRecords = interfaceSliceToStringSlice(v) } if v, ok := raw["tls"].(string); ok { asset.TLS = v } // Metadata if v, ok := raw["tech"].([]interface{}); ok { asset.Technologies = interfaceSliceToStringSlice(v) } if v, ok := raw["time"].(string); ok { asset.ResponseTime = v } if v, ok := raw["remarks"].(string); ok { asset.Labels = v } // Validate required fields if asset.AssetValue == "" { return nil, fmt.Errorf("asset_value is required") } if asset.Workspace == "" { return nil, fmt.Errorf("workspace is required") } return asset, nil } // ImportEventLogs imports event logs from a JSONL file func (i *JSONLImporter) ImportEventLogs(ctx context.Context, filePath string) (*ImportResult, error) { file, err := os.Open(filePath) if err != nil { return nil, fmt.Errorf("failed to open file: %w", err) } defer func() { _ = file.Close() }() return i.ImportEventLogsFromReader(ctx, file) } // ImportEventLogsFromReader imports event logs from an io.Reader func (i *JSONLImporter) ImportEventLogsFromReader(ctx context.Context, r io.Reader) (*ImportResult, error) { startTime := time.Now() scanner := bufio.NewScanner(r) scanner.Buffer(make([]byte, 1024*1024), 10*1024*1024) result := &ImportResult{} batch := make([]*EventLog, 0, i.batchSize) for scanner.Scan() { result.Total++ line := scanner.Bytes() if len(line) == 0 { continue } event, err := ParseEventLogLine(line) if err != nil { result.Failed++ result.Errors = append(result.Errors, ImportError{ Line: result.Total, Error: err.Error(), }) continue } batch = append(batch, event) if len(batch) >= i.batchSize { if _, err := i.db.NewInsert().Model(&batch).Exec(ctx); err != nil { return result, fmt.Errorf("batch insert failed: %w", err) } result.Imported += len(batch) batch = batch[:0] } } if len(batch) > 0 { if _, err := i.db.NewInsert().Model(&batch).Exec(ctx); err != nil { return result, fmt.Errorf("final batch insert failed: %w", err) } result.Imported += len(batch) } result.Duration = time.Since(startTime) return result, scanner.Err() } // ParseEventLogLine parses a single JSONL line into an EventLog func ParseEventLogLine(line []byte) (*EventLog, error) { var event EventLog if err := json.Unmarshal(line, &event); err != nil { return nil, fmt.Errorf("invalid JSON: %w", err) } if event.Topic == "" { return nil, fmt.Errorf("topic is required") } if event.CreatedAt.IsZero() { event.CreatedAt = time.Now() } return &event, nil } // Helper functions func interfaceSliceToStringSlice(slice []interface{}) []string { result := make([]string, 0, len(slice)) for _, v := range slice { if s, ok := v.(string); ok { result = append(result, s) } } return result } func truncateString(s string, maxLen int) string { if len(s) <= maxLen { return s } return s[:maxLen] + "..." }