From c0309394be954ab6f6be8f09a374c4fdb8709437 Mon Sep 17 00:00:00 2001
From: candiduslynx <candiduslynx@gmail.com>
Date: Tue, 27 Jun 2023 12:41:48 +0300
Subject: [PATCH] put null helpers back

---
 plugin/nulls.go         | 72 +++++++++++++++++++++++++++++++++++++++++
 plugin/testing_write.go | 13 ++++----
 2 files changed, 79 insertions(+), 6 deletions(-)
 create mode 100644 plugin/nulls.go

diff --git a/plugin/nulls.go b/plugin/nulls.go
new file mode 100644
index 0000000000..e60a88a24f
--- /dev/null
+++ b/plugin/nulls.go
@@ -0,0 +1,72 @@
+package plugin
+
+import (
+	"github.com/apache/arrow/go/v13/arrow"
+	"github.com/apache/arrow/go/v13/arrow/array"
+	"github.com/apache/arrow/go/v13/arrow/memory"
+)
+
+func stripNullsFromLists(records []arrow.Record) {
+	for i := range records {
+		cols := records[i].Columns()
+		for c, col := range cols {
+			if col.DataType().ID() != arrow.LIST {
+				continue
+			}
+
+			list := col.(*array.List)
+			bldr := array.NewListBuilder(memory.DefaultAllocator, list.DataType().(*arrow.ListType).Elem())
+			for j := 0; j < list.Len(); j++ {
+				if list.IsNull(j) {
+					bldr.AppendNull()
+					continue
+				}
+				bldr.Append(true)
+				vBldr := bldr.ValueBuilder()
+				from, to := list.ValueOffsets(j)
+				slc := array.NewSlice(list.ListValues(), from, to)
+				for k := 0; k < int(to-from); k++ {
+					if slc.IsNull(k) {
+						continue
+					}
+					err := vBldr.AppendValueFromString(slc.ValueStr(k))
+					if err != nil {
+						panic(err)
+					}
+				}
+			}
+			cols[c] = bldr.NewArray()
+		}
+		records[i] = array.NewRecord(records[i].Schema(), cols, records[i].NumRows())
+	}
+}
+
+type AllowNullFunc func(arrow.DataType) bool
+
+func (f AllowNullFunc) replaceNullsByEmpty(records []arrow.Record) {
+	if f == nil {
+		return
+	}
+	for i := range records {
+		cols := records[i].Columns()
+		for c, col := range records[i].Columns() {
+			if col.NullN() == 0 || f(col.DataType()) {
+				continue
+			}
+
+			builder := array.NewBuilder(memory.DefaultAllocator, records[i].Column(c).DataType())
+			for j := 0; j < col.Len(); j++ {
+				if col.IsNull(j) {
+					builder.AppendEmptyValue()
+					continue
+				}
+
+				if err := builder.AppendValueFromString(col.ValueStr(j)); err != nil {
+					panic(err)
+				}
+			}
+			cols[c] = builder.NewArray()
+		}
+		records[i] = array.NewRecord(records[i].Schema(), cols, records[i].NumRows())
+	}
+}
diff --git a/plugin/testing_write.go b/plugin/testing_write.go
index 4f198e80a8..6be215baa3 100644
--- a/plugin/testing_write.go
+++ b/plugin/testing_write.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"testing"
 
+	"github.com/apache/arrow/go/v13/arrow"
 	"github.com/cloudquery/plugin-sdk/v4/schema"
 )
 
@@ -16,7 +17,7 @@ type WriterTestSuite struct {
 	// Destinations that have problems representing some data types should provide a custom implementation here.
 	// If this param is empty, the default is to allow all data types to be nullable.
 	// When the value returned by this func is `true` the comparison is made with the empty value instead of null.
-	// allowNull AllowNullFunc
+	allowNull AllowNullFunc
 
 	// IgnoreNullsInLists allows stripping null values from lists before comparison.
 	// Destination setups that don't support nulls in lists should set this to true.
@@ -56,11 +57,11 @@ type WriterTestSuiteTests struct {
 
 type NewPluginFunc func() *Plugin
 
-// func WithTestSourceAllowNull(allowNull func(arrow.DataType) bool) func(o *WriterTestSuite) {
-// 	return func(o *WriterTestSuite) {
-// 		o.allowNull = allowNull
-// 	}
-// }
+func WithTestSourceAllowNull(allowNull func(arrow.DataType) bool) func(o *WriterTestSuite) {
+	return func(o *WriterTestSuite) {
+		o.allowNull = allowNull
+	}
+}
 
 func WithTestIgnoreNullsInLists() func(o *WriterTestSuite) {
 	return func(o *WriterTestSuite) {