1c7194176a
Add compressFileZstd, decompressZstd, and walkZstFiles helpers using klauspost/compress. Promote zstd from indirect to direct dep. Co-Authored-By: Virgil <virgil@lethean.io>
77 lines
1.7 KiB
Go
77 lines
1.7 KiB
Go
// zstd.go — zstd compression/decompression for cold storage JSONL files.
|
|
package lem
|
|
|
|
import (
|
|
"fmt"
|
|
"io"
|
|
"os"
|
|
"path/filepath"
|
|
"strings"
|
|
|
|
"github.com/klauspost/compress/zstd"
|
|
)
|
|
|
|
// decompressZstd reads a .zst file and writes the decompressed content to dst.
|
|
func decompressZstd(src, dst string) error {
|
|
f, err := os.Open(src)
|
|
if err != nil {
|
|
return fmt.Errorf("open %s: %w", src, err)
|
|
}
|
|
defer f.Close()
|
|
|
|
dec, err := zstd.NewReader(f)
|
|
if err != nil {
|
|
return fmt.Errorf("zstd reader: %w", err)
|
|
}
|
|
defer dec.Close()
|
|
|
|
out, err := os.Create(dst)
|
|
if err != nil {
|
|
return fmt.Errorf("create %s: %w", dst, err)
|
|
}
|
|
defer out.Close()
|
|
|
|
if _, err := io.Copy(out, dec); err != nil {
|
|
return fmt.Errorf("decompress %s: %w", src, err)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// compressFileZstd compresses src to dst with zstd level 3.
|
|
func compressFileZstd(src, dst string) error {
|
|
f, err := os.Open(src)
|
|
if err != nil {
|
|
return fmt.Errorf("open %s: %w", src, err)
|
|
}
|
|
defer f.Close()
|
|
|
|
out, err := os.Create(dst)
|
|
if err != nil {
|
|
return fmt.Errorf("create %s: %w", dst, err)
|
|
}
|
|
defer out.Close()
|
|
|
|
enc, err := zstd.NewWriter(out, zstd.WithEncoderLevel(zstd.SpeedDefault))
|
|
if err != nil {
|
|
return fmt.Errorf("zstd writer: %w", err)
|
|
}
|
|
|
|
if _, err := io.Copy(enc, f); err != nil {
|
|
enc.Close()
|
|
return fmt.Errorf("compress %s: %w", src, err)
|
|
}
|
|
return enc.Close()
|
|
}
|
|
|
|
// walkZstFiles walks a directory tree and calls fn for each .jsonl.zst file.
|
|
func walkZstFiles(root string, fn func(zstPath string) error) error {
|
|
return filepath.Walk(root, func(path string, info os.FileInfo, err error) error {
|
|
if err != nil || info.IsDir() {
|
|
return err
|
|
}
|
|
if strings.HasSuffix(path, ".jsonl.zst") {
|
|
return fn(path)
|
|
}
|
|
return nil
|
|
})
|
|
}
|