feat: Handle panics in handlers (#575)

Handle panics in handlers by converting panics to `error`

stream.go

@@ -1,6 +1,7 @@
 package filetypes
 
 import (
+ "errors"
  "fmt"
  "io"
 
@@ -54,11 +55,22 @@ func (cl *Client) StartStream(table *schema.Table, uploadFunc func(io.Reader) er
 }
 
 // Write to the stream opened with StartStream.
-func (s *Stream) Write(records []arrow.Record) error {
+func (s *Stream) Write(records []arrow.Record) (retErr error) {
  if len(records) == 0 {
  return nil
  }
 
+ defer func() {
+ if msg := recover(); msg != nil {
+ switch v := msg.(type) {
+ case error:
+ retErr = fmt.Errorf("panic: %w [recovered]", v)
+ default:
+ retErr = fmt.Errorf("panic: %v [recovered]", msg)
+ }
+ }
+ }()
+
  return s.h.WriteContent(records)
 }
 
@@ -74,11 +86,11 @@ func (s *Stream) FinishWithError(finishError error) error {
  return <-s.done
  }
 
- if err := s.h.WriteFooter(); err != nil {
+ if err := s.writeFooter(); err != nil {
  if !s.wc.closed {
  _ = s.wc.CloseWithError(err)
  }
- return fmt.Errorf("failed to write footer: %w", <-s.done)
+ return fmt.Errorf("failed to write footer: %w", errors.Join(err, <-s.done))
  }
 
  // ParquetWriter likes to close the underlying writer, so we need to check if it's already closed
@@ -90,3 +102,18 @@ func (s *Stream) FinishWithError(finishError error) error {
 
  return <-s.done
 }
+
+func (s *Stream) writeFooter() (retErr error) {
+ defer func() {
+ if msg := recover(); msg != nil {
+ switch v := msg.(type) {
+ case error:
+ retErr = fmt.Errorf("panic: %w [recovered]", v)
+ default:
+ retErr = fmt.Errorf("panic: %v [recovered]", msg)
+ }
+ }
+ }()
+
+ return s.h.WriteFooter()
+}

stream_panic_test.go

@@ -0,0 +1,121 @@
+package filetypes
+
+import (
+ "errors"
+ "io"
+ "testing"
+
+ "github.com/apache/arrow/go/v17/arrow"
+ "github.com/apache/arrow/go/v17/arrow/array"
+ "github.com/apache/arrow/go/v17/arrow/memory"
+ "github.com/cloudquery/filetypes/v4/types"
+ "github.com/cloudquery/plugin-sdk/v4/schema"
+ "github.com/stretchr/testify/require"
+)
+
+func TestPanicOnHeader(t *testing.T) {
+ r := require.New(t)
+ cl := &Client{
+ spec: &FileSpec{
+ Compression: CompressionTypeNone,
+ },
+ filetype: &customWriter{
+ PanicOnHeader: true,
+ },
+ }
+
+ stream, err := cl.StartStream(&schema.Table{}, func(io.Reader) error {
+ return nil
+ })
+ r.Nil(stream)
+ r.Error(err)
+ r.ErrorContains(err, "panic:")
+}
+
+func TestPanicOnWrite(t *testing.T) {
+ r := require.New(t)
+ cl := &Client{
+ spec: &FileSpec{
+ Compression: CompressionTypeNone,
+ },
+ filetype: &customWriter{
+ PanicOnWrite: true,
+ },
+ }
+
+ table := &schema.Table{
+ Name: "test",
+ Columns: []schema.Column{
+ {Name: "name", Type: arrow.BinaryTypes.String},
+ },
+ }
+ bldr := array.NewRecordBuilder(memory.DefaultAllocator, table.ToArrowSchema())
+ bldr.Field(0).(*array.StringBuilder).Append("foo")
+ bldr.Field(0).(*array.StringBuilder).Append("bar")
+ record := bldr.NewRecord()
+
+ stream, err := cl.StartStream(table, func(io.Reader) error {
+ return nil
+ })
+ r.NoError(err)
+ err = stream.Write([]arrow.Record{record})
+ r.Error(err)
+ r.ErrorContains(err, "panic:")
+
+ r.NoError(stream.Finish())
+}
+
+func TestPanicOnClose(t *testing.T) {
+ r := require.New(t)
+ cl := &Client{
+ spec: &FileSpec{
+ Compression: CompressionTypeNone,
+ },
+ filetype: &customWriter{
+ PanicOnClose: true,
+ },
+ }
+
+ stream, err := cl.StartStream(&schema.Table{}, func(io.Reader) error {
+ return nil
+ })
+ r.NoError(err)
+ r.NoError(stream.Write(nil))
+
+ err = stream.Finish()
+ r.Error(err)
+ r.ErrorContains(err, "panic:")
+}
+
+type customWriter struct {
+ PanicOnHeader bool
+ PanicOnWrite bool
+ PanicOnClose bool
+}
+type customHandle struct {
+ w *customWriter
+}
+
+func (w *customWriter) WriteHeader(io.Writer, *schema.Table) (types.Handle, error) {
+ if w.PanicOnHeader {
+ panic("test panic")
+ }
+ return &customHandle{w: w}, nil
+}
+
+func (*customWriter) Read(types.ReaderAtSeeker, *schema.Table, chan<- arrow.Record) error {
+ return errors.New("not implemented")
+}
+
+func (h *customHandle) WriteContent([]arrow.Record) error {
+ if h.w.PanicOnWrite {
+ panic("test panic")
+ }
+ return nil
+}
+func (h *customHandle) WriteFooter() error {
+ if h.w.PanicOnClose {
+ panic("test panic")
+ }
+ return nil
+}

write.go

@@ -14,7 +14,18 @@ func (cl *Client) WriteTableBatchFile(w io.Writer, table *schema.Table, records
  return types.WriteAll(cl, w, table, records)
 }
 
-func (cl *Client) WriteHeader(w io.Writer, t *schema.Table) (types.Handle, error) {
+func (cl *Client) WriteHeader(w io.Writer, t *schema.Table) (h types.Handle, retErr error) {
+ defer func() {
+ if msg := recover(); msg != nil {
+ switch v := msg.(type) {
+ case error:
+ retErr = fmt.Errorf("panic: %w [recovered]", v)
+ default:
+ retErr = fmt.Errorf("panic: %v [recovered]", msg)
+ }
+ }
+ }()
+
  switch cl.spec.Compression {
  case CompressionTypeNone:
  return cl.filetype.WriteHeader(w, t)