Adds the ability to read gzip-tar file for the taxonomy dump

This commit is contained in:
Eric Coissac
2025-01-24 11:47:59 +01:00
parent ffd67252c3
commit 3137c1f841
17 changed files with 305 additions and 64 deletions

20
pkg/obiutils/tar.go Normal file
View File

@ -0,0 +1,20 @@
package obiutils
import (
"archive/tar"
"fmt"
)
func TarFileReader(file *Reader, path string) (*tar.Reader, error) {
tarfile := tar.NewReader(file)
header, err := tarfile.Next()
for err == nil {
if header.Name == path {
return tarfile, nil
}
header, err = tarfile.Next()
}
return nil, fmt.Errorf("file not found: %s", path)
}

437
pkg/obiutils/xopen.go Normal file
View File

@ -0,0 +1,437 @@
// This is an integration of the xopen package originally written by Brent Pedersen
// (https://github.com/brentp/xopen).
//
// Here it can be considered as a fork of [Wei Shen](http://shenwei.me) the version :
//
// https://github.com/shenwei356/xopen
//
// Package xopen makes it easy to get buffered readers and writers.
// Ropen opens a (possibly gzipped) file/process/http site for buffered reading.
// Wopen opens a (possibly gzipped) file for buffered writing.
// Both will use gzip when appropriate and will user buffered IO.
package obiutils
import (
"bufio"
"errors"
"fmt"
"io"
"net/http"
"os"
"os/exec"
"os/user"
"path/filepath"
"strings"
"github.com/dsnet/compress/bzip2"
"github.com/klauspost/compress/zstd"
gzip "github.com/klauspost/pgzip"
"github.com/ulikunitz/xz"
)
// Level is the default compression level of gzip.
// This value will be automatically adjusted to the default value of zstd or bzip2.
var Level = gzip.DefaultCompression
// ErrNoContent means nothing in the stream/file.
var ErrNoContent = errors.New("xopen: no content")
// ErrDirNotSupported means the path is a directory.
var ErrDirNotSupported = errors.New("xopen: input is a directory")
// IsGzip returns true buffered Reader has the gzip magic.
func IsGzip(b *bufio.Reader) (bool, error) {
return CheckBytes(b, []byte{0x1f, 0x8b})
}
// IsXz returns true buffered Reader has the xz magic.
func IsXz(b *bufio.Reader) (bool, error) {
return CheckBytes(b, []byte{0xfd, 0x37, 0x7a, 0x58, 0x5a, 0x00})
}
// IsZst returns true buffered Reader has the zstd magic.
func IsZst(b *bufio.Reader) (bool, error) {
return CheckBytes(b, []byte{0x28, 0xB5, 0x2f, 0xfd})
}
// IsBzip2 returns true buffered Reader has the bzip2 magic.
func IsBzip2(b *bufio.Reader) (bool, error) {
return CheckBytes(b, []byte{0x42, 0x5a, 0x68})
}
// IsStdin checks if we are getting data from stdin.
func IsStdin() bool {
// http://stackoverflow.com/a/26567513
stat, err := os.Stdin.Stat()
if err != nil {
return false
}
return (stat.Mode() & os.ModeCharDevice) == 0
}
// ExpandUser expands ~/path and ~otheruser/path appropriately
func ExpandUser(path string) (string, error) {
if len(path) == 0 || path[0] != '~' {
return path, nil
}
var u *user.User
var err error
if len(path) == 1 || path[1] == '/' {
u, err = user.Current()
} else {
name := strings.Split(path[1:], "/")[0]
u, err = user.Lookup(name)
}
if err != nil {
return "", err
}
home := u.HomeDir
path = home + "/" + path[1:]
return path, nil
}
// Exists checks if a local file exits
func Exists(path string) bool {
path, perr := ExpandUser(path)
if perr != nil {
return false
}
_, err := os.Stat(path)
return err == nil
}
// CheckBytes peeks at a buffered stream and checks if the first read bytes match.
func CheckBytes(b *bufio.Reader, buf []byte) (bool, error) {
m, err := b.Peek(len(buf))
if err != nil {
// return false, ErrNoContent
return false, err // EOF
}
for i := range buf {
if m[i] != buf[i] {
return false, nil
}
}
return true, nil
}
// Reader is returned by Ropen
type Reader struct {
*bufio.Reader
rdr io.Reader
gz io.ReadCloser
}
// Close the associated files.
func (r *Reader) Close() error {
var err error
if r.gz != nil {
err = r.gz.Close()
if err != nil {
return err
}
}
if c, ok := r.rdr.(io.ReadCloser); ok {
err = c.Close()
if err != nil {
return err
}
}
return nil
}
// Writer is returned by Wopen
type Writer struct {
*bufio.Writer
wtr *os.File
gz *gzip.Writer
xw *xz.Writer
zw *zstd.Encoder
bz2 *bzip2.Writer
}
// Close the associated files.
func (w *Writer) Close() error {
var err error
err = w.Flush()
if err != nil {
return err
}
if w.gz != nil {
err = w.gz.Close()
if err != nil {
return err
}
}
if w.xw != nil {
err = w.xw.Close()
if err != nil {
return err
}
}
if w.zw != nil {
err = w.zw.Close()
if err != nil {
return err
}
}
if w.bz2 != nil {
err = w.bz2.Close()
if err != nil {
return err
}
}
return w.wtr.Close()
}
// Flush the writer.
func (w *Writer) Flush() error {
var err error
err = w.Writer.Flush()
if err != nil {
return err
}
if w.gz != nil {
err = w.gz.Flush()
if err != nil {
return err
}
}
if w.zw != nil {
err = w.zw.Flush()
if err != nil {
return err
}
}
return nil
}
var bufSize = 65536
// Buf returns a buffered reader from an io.Reader
// If f == "-", then it will attempt to read from os.Stdin.
// If the file is gzipped, it will be read as such.
func Buf(r io.Reader) (*Reader, error) {
b := bufio.NewReaderSize(r, bufSize)
var rd io.Reader
var rdr io.ReadCloser
if is, err := IsGzip(b); err != nil {
// check BOM
t, _, err := b.ReadRune() // no content
if err != nil {
return nil, ErrNoContent
}
if t != '\uFEFF' {
b.UnreadRune()
}
return &Reader{b, r, rdr}, nil // non-gzip file with content less than 2 bytes
} else if is {
rdr, err = gzip.NewReader(b)
if err != nil {
return nil, err
}
b = bufio.NewReaderSize(rdr, bufSize)
} else if is, err := IsZst(b); err != nil {
// check BOM
t, _, err := b.ReadRune() // no content
if err != nil {
return nil, ErrNoContent
}
if t != '\uFEFF' {
b.UnreadRune()
}
return &Reader{b, r, rdr}, nil // non-gzip/zst file with content less than 4 bytes
} else if is {
rd, err = zstd.NewReader(b)
if err != nil {
return nil, err
}
b = bufio.NewReaderSize(rd, bufSize)
} else if is, err := IsXz(b); err != nil {
// check BOM
t, _, err := b.ReadRune() // no content
if err != nil {
return nil, ErrNoContent
}
if t != '\uFEFF' {
b.UnreadRune()
}
return &Reader{b, r, rdr}, nil // non-gzip/zst/xz file with content less than 6 bytes
} else if is {
rd, err = xz.NewReader(b)
if err != nil {
return nil, err
}
b = bufio.NewReaderSize(rd, bufSize)
} else if is, err := IsBzip2(b); err != nil {
// check BOM
t, _, err := b.ReadRune() // no content
if err != nil {
return nil, ErrNoContent
}
if t != '\uFEFF' {
b.UnreadRune()
}
return &Reader{b, r, rdr}, nil // non-gzip/zst/xz file with content less than 6 bytes
} else if is {
rd, err = bzip2.NewReader(b, &bzip2.ReaderConfig{})
if err != nil {
return nil, err
}
b = bufio.NewReaderSize(rd, bufSize)
}
// other files with content >= 6 bytes
// check BOM
t, _, err := b.ReadRune()
if err != nil {
return nil, ErrNoContent
}
if t != '\uFEFF' {
b.UnreadRune()
}
return &Reader{b, r, rdr}, nil
}
// XReader returns a reader from a url string or a file.
func XReader(f string) (io.Reader, error) {
if strings.HasPrefix(f, "http://") || strings.HasPrefix(f, "https://") {
var rsp *http.Response
rsp, err := http.Get(f)
if err != nil {
return nil, err
}
if rsp.StatusCode != 200 {
return nil, fmt.Errorf("http error downloading %s. status: %s", f, rsp.Status)
}
rdr := rsp.Body
return rdr, nil
}
f, err := ExpandUser(f)
if err != nil {
return nil, err
}
fi, err := os.Stat(f)
if err != nil {
return nil, err
}
if fi.IsDir() {
return nil, ErrDirNotSupported
}
return os.Open(f)
}
// Ropen opens a buffered reader.
func Ropen(f string) (*Reader, error) {
var err error
var rdr io.Reader
if f == "-" {
if !IsStdin() {
return nil, errors.New("stdin not detected")
}
b, err := Buf(os.Stdin)
return b, err
} else if f[0] == '|' {
// TODO: use csv to handle quoted file names.
cmdStrs := strings.Split(f[1:], " ")
var cmd *exec.Cmd
if len(cmdStrs) == 2 {
cmd = exec.Command(cmdStrs[0], cmdStrs[1:]...)
} else {
cmd = exec.Command(cmdStrs[0])
}
rdr, err = cmd.StdoutPipe()
if err != nil {
return nil, err
}
err = cmd.Start()
if err != nil {
return nil, err
}
} else {
rdr, err = XReader(f)
}
if err != nil {
return nil, err
}
b, err := Buf(rdr)
return b, err
}
// Wopen opens a buffered reader.
// If f == "-", then stdout will be used.
// If f endswith ".gz", then the output will be gzipped.
// If f endswith ".xz", then the output will be zx-compressed.
// If f endswith ".zst", then the output will be zstd-compressed.
// If f endswith ".bz2", then the output will be bzip2-compressed.
func Wopen(f string) (*Writer, error) {
return WopenFile(f, os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0666)
}
// WopenFile opens a buffered reader.
// If f == "-", then stdout will be used.
// If f endswith ".gz", then the output will be gzipped.
// If f endswith ".xz", then the output will be zx-compressed.
// If f endswith ".bz2", then the output will be bzip2-compressed.
func WopenFile(f string, flag int, perm os.FileMode) (*Writer, error) {
var wtr *os.File
if f == "-" {
wtr = os.Stdout
} else {
dir := filepath.Dir(f)
fi, err := os.Stat(dir)
if err == nil && !fi.IsDir() {
return nil, fmt.Errorf("can not write file into a non-directory path: %s", dir)
}
if os.IsNotExist(err) {
os.MkdirAll(dir, 0755)
}
wtr, err = os.OpenFile(f, flag, perm)
if err != nil {
return nil, err
}
}
f2 := strings.ToLower(f)
if strings.HasSuffix(f2, ".gz") {
gz, err := gzip.NewWriterLevel(wtr, Level)
if err != nil {
err = fmt.Errorf("xopen: %s", err)
}
return &Writer{bufio.NewWriterSize(gz, bufSize), wtr, gz, nil, nil, nil}, err
}
if strings.HasSuffix(f2, ".xz") {
xw, err := xz.NewWriter(wtr)
return &Writer{bufio.NewWriterSize(xw, bufSize), wtr, nil, xw, nil, nil}, err
}
if strings.HasSuffix(f2, ".zst") {
level := Level
if level == gzip.DefaultCompression {
level = 2
}
zw, err := zstd.NewWriter(wtr, zstd.WithEncoderLevel(zstd.EncoderLevel(level)))
if err != nil {
err = fmt.Errorf("xopen: zstd: %s", err)
}
return &Writer{bufio.NewWriterSize(zw, bufSize), wtr, nil, nil, zw, nil}, err
}
if strings.HasSuffix(f2, ".bz2") {
level := Level
if level == gzip.DefaultCompression {
level = 6
}
bz2, err := bzip2.NewWriter(wtr, &bzip2.WriterConfig{Level: level})
if err != nil {
err = fmt.Errorf("xopen: %s", err)
}
return &Writer{bufio.NewWriterSize(bz2, bufSize), wtr, nil, nil, nil, bz2}, err
}
return &Writer{bufio.NewWriterSize(wtr, bufSize), wtr, nil, nil, nil, nil}, nil
}

148
pkg/obiutils/xopen_test.go Normal file
View File

@ -0,0 +1,148 @@
package obiutils
import (
"bufio"
"bytes"
"compress/gzip"
"fmt"
"io"
"os"
"strings"
"testing"
. "gopkg.in/check.v1"
)
func Test(t *testing.T) { TestingT(t) }
type XopenTest struct{}
var _ = Suite(&XopenTest{})
func gzFromString(s string) string {
var c bytes.Buffer
gz := gzip.NewWriter(&c)
gz.Write([]byte(s))
return c.String()
}
var gzTests = []struct {
isGz bool
data string
}{
{false, "asdf"},
{true, gzFromString("asdf")},
}
func (s *XopenTest) TestIsGzip(c *C) {
for _, t := range gzTests {
isGz, err := IsGzip(bufio.NewReader(strings.NewReader(t.data)))
c.Assert(err, IsNil)
c.Assert(t.isGz, Equals, isGz)
}
}
func (s *XopenTest) TestIsStdin(c *C) {
r := IsStdin()
c.Assert(r, Equals, false)
}
func (s *XopenTest) TestRopen(c *C) {
rdr, err := Ropen("-")
c.Assert(err, ErrorMatches, "stdin not detected")
c.Assert(rdr, IsNil)
}
func (s *XopenTest) TestWopen(c *C) {
for _, f := range []string{"t.gz", "t"} {
testString := "ASDF1234"
wtr, err := Wopen(f)
c.Assert(err, IsNil)
_, err = os.Stat(f)
c.Assert(err, IsNil)
c.Assert(wtr.wtr, NotNil)
fmt.Fprint(wtr, testString)
wtr.Close()
rdr, err := Ropen(f)
c.Assert(err, IsNil)
str, err := rdr.ReadString(99)
c.Assert(str, Equals, testString)
c.Assert(err, Equals, io.EOF)
str, _ = rdr.ReadString(99)
c.Assert(str, Equals, "")
rdr.Close()
os.Remove(f)
}
}
var httpTests = []struct {
url string
expectError bool
}{
{"https://raw.githubusercontent.com/brentp/xopen/master/README.md", false},
{"http://raw.githubusercontent.com/brentp/xopen/master/README.md", false},
{"http://raw.githubusercontent.com/brentp/xopen/master/BAD.md", true},
}
func (s *XopenTest) TestReadHttp(c *C) {
for _, t := range httpTests {
rdr, err := Ropen(t.url)
if !t.expectError {
c.Assert(err, IsNil)
v, err := rdr.ReadString(byte('\n'))
c.Assert(err, IsNil)
c.Assert(len(v), Not(Equals), 0)
} else {
c.Assert(err, ErrorMatches, ".* 404 Not Found")
}
}
}
// func (s *XopenTest) TestReadProcess(c *C) {
// for _, cmd := range []string{"|ls -lh", "|ls", "|ls -lh xopen_test.go"} {
// rdr, err := Ropen(cmd)
// c.Assert(err, IsNil)
// b := make([]byte, 1000)
// _, err = rdr.Read(b)
// if err != io.EOF {
// c.Assert(err, IsNil)
// }
// lines := strings.Split(string(b), "\n")
// has := false
// for _, line := range lines {
// if strings.Contains(line, "xopen_test.go") {
// has = true
// }
// }
// c.Assert(has, Equals, true)
// }
// }
func (s *XopenTest) TestOpenStdout(c *C) {
w, err := Wopen("-")
c.Assert(err, IsNil)
c.Assert(w.wtr, Equals, os.Stdout)
}
func (s *XopenTest) TestOpenBadFile(c *C) {
r, err := Ropen("XXXXXXXXXXXXXXXXXXXXXXX")
c.Assert(r, IsNil)
c.Assert(err, ErrorMatches, ".*no such file.*")
}
func (s *XopenTest) TestExists(c *C) {
c.Assert(Exists("xopen.go"), Equals, true)
c.Assert(Exists("____xx"), Equals, false)
}
func (s *XopenTest) TestUser(c *C) {
c.Assert(Exists("~"), Equals, true)
}
func (s *XopenTest) TestExpand(c *C) {
_, err := ExpandUser("~baduser66")
c.Assert(err, Not(IsNil))
}