feat: Include CSV headers

csv/client.go

@@ -0,0 +1,38 @@
+package csv
+
+type Options func(*Client)
+
+// Client is a csv client.
+type Client struct {
+	// fileFormat     string
+	IncludeHeaders bool
+	Delimiter      rune
+}
+
+func NewClient(options ...Options) (*Client, error) {
+	c := &Client{}
+	for _, option := range options {
+		option(c)
+	}
+	c.defaults()
+
+	return c, nil
+}
+
+func WithHeader() Options {
+	return func(c *Client) {
+		c.IncludeHeaders = true
+	}
+}
+
+func WithDelimiter(delimiter rune) Options {
+	return func(c *Client) {
+		c.Delimiter = delimiter
+	}
+}
+
+func (cl *Client) defaults() {
+	if cl.Delimiter == 0 {
+		cl.Delimiter = ','
+	}
+}

csv/read.go

@@ -9,8 +9,8 @@ import (
 	"github.com/cloudquery/plugin-sdk/schema"
 )
 
-func Read(f io.Reader, table *schema.Table, sourceName string, res chan<- []any) error {
-	reader := csv.NewReader(f)
+func (*Client) Read(r io.Reader, table *schema.Table, res chan<- []any) error {
+	reader := csv.NewReader(r)
 	sourceNameIndex := table.Columns.Index(schema.CqSourceNameColumn.Name)
 	if sourceNameIndex == -1 {
 		return fmt.Errorf("could not find column %s in table %s", schema.CqSourceNameColumn.Name, table.Name)
@@ -23,9 +23,6 @@ func Read(f io.Reader, table *schema.Table, sourceName string, res chan<- []any)
 			}
 			return err
 		}
-		if record[sourceNameIndex] != sourceName {
-			continue
-		}
 		values := make([]any, len(record))
 		for i, v := range record {
 			values[i] = v

csv/write.go

@@ -7,8 +7,18 @@ import (
 	"github.com/cloudquery/plugin-sdk/schema"
 )
 
-func WriteTableBatch(w io.Writer, _ *schema.Table, resources [][]any) error {
+func (cl *Client) WriteTableBatch(w io.Writer, table *schema.Table, resources [][]any) error {
 	writer := csv.NewWriter(w)
+	writer.Comma = cl.Delimiter
+	if cl.IncludeHeaders {
+		tableHeaders := make([]string, len(table.Columns))
+		for index, header := range table.Columns {
+			tableHeaders[index] = header.Name
+		}
+		if err := writer.Write(tableHeaders); err != nil {
+			return err
+		}
+	}
 	for _, resource := range resources {
 		record := make([]string, len(resource))
 		for i, v := range resource {

csv/write_read_test.go

@@ -10,43 +10,62 @@ import (
 )
 
 func TestWriteRead(t *testing.T) {
-	var b bytes.Buffer
-	table := testdata.TestTable("test")
-	cqtypes := testdata.GenTestData(table)
-	if err := cqtypes[0].Set("test-source"); err != nil {
-		t.Fatal(err)
+	cases := []struct {
+		options     []Options
+		outputCount int
+	}{
+		{options: []Options{WithHeader()}, outputCount: 2},
+		{outputCount: 1},
 	}
-	writer := bufio.NewWriter(&b)
-	transformer := &Transformer{}
-	transformedValues := schema.TransformWithTransformer(transformer, cqtypes)
-	// schema.TransformWithTransformer(tra)
-	if err := WriteTableBatch(writer, table, [][]any{transformedValues}); err != nil {
-		t.Fatal(err)
-	}
-	writer.Flush()
-	reader := bufio.NewReader(&b)
-	ch := make(chan []any)
-	var readErr error
-	go func() {
-		readErr = Read(reader, table, "test-source", ch)
-		close(ch)
-	}()
-	totalCount := 0
-	reverseTransformer := &ReverseTransformer{}
-	for row := range ch {
-		gotCqtypes, err := reverseTransformer.ReverseTransformValues(table, row)
+
+	for _, tc := range cases {
+		var b bytes.Buffer
+		table := testdata.TestTable("test")
+		cqtypes := testdata.GenTestData(table)
+		if err := cqtypes[0].Set("test-source"); err != nil {
+			t.Fatal(err)
+		}
+		writer := bufio.NewWriter(&b)
+		reader := bufio.NewReader(&b)
+		transformer := &Transformer{}
+		transformedValues := schema.TransformWithTransformer(transformer, cqtypes)
+		client, err := NewClient(tc.options...)
 		if err != nil {
 			t.Fatal(err)
 		}
-		if diff := cqtypes.Diff(gotCqtypes); diff != "" {
-			t.Fatalf("got diff: %s", diff)
+
+		if err := client.WriteTableBatch(writer, table, [][]any{transformedValues}); err != nil {
+			t.Fatal(err)
+		}
+		writer.Flush()
+
+		ch := make(chan []any)
+		var readErr error
+		go func() {
+			readErr = client.Read(reader, table, ch)
+			close(ch)
+		}()
+		totalCount := 0
+		reverseTransformer := &ReverseTransformer{}
+		for row := range ch {
+			if client.IncludeHeaders && totalCount == 0 {
+				totalCount++
+				continue
+			}
+			gotCqtypes, err := reverseTransformer.ReverseTransformValues(table, row)
+			if err != nil {
+				t.Fatal(err)
+			}
+			if diff := cqtypes.Diff(gotCqtypes); diff != "" {
+				t.Fatalf("got diff: %s", diff)
+			}
+			totalCount++
+		}
+		if readErr != nil {
+			t.Fatal(readErr)
+		}
+		if totalCount != tc.outputCount {
+			t.Fatalf("expected %d row, got %d", tc.outputCount, totalCount)
 		}
-		totalCount++
-	}
-	if readErr != nil {
-		t.Fatal(readErr)
-	}
-	if totalCount != 1 {
-		t.Fatalf("expected 1 row, got %d", totalCount)
 	}
 }

json/client.go

@@ -0,0 +1,14 @@
+package json
+
+type Option func(*Client)
+
+type Client struct{}
+
+func NewClient(options ...Option) (*Client, error) {
+	c := &Client{}
+	for _, option := range options {
+		option(c)
+	}
+
+	return c, nil
+}

json/read.go

@@ -11,12 +11,12 @@ import (
 
 const maxJSONSize = 1024 * 1024 * 20
 
-func Read(f io.Reader, table *schema.Table, sourceName string, res chan<- []any) error {
+func (*Client) Read(r io.Reader, table *schema.Table, sourceName string, res chan<- []any) error {
 	sourceNameIndex := table.Columns.Index(schema.CqSourceNameColumn.Name)
 	if sourceNameIndex == -1 {
 		return fmt.Errorf("could not find column %s in table %s", schema.CqSourceNameColumn.Name, table.Name)
 	}
-	scanner := bufio.NewScanner(f)
+	scanner := bufio.NewScanner(r)
 	scanner.Buffer(make([]byte, maxJSONSize), maxJSONSize)
 	for scanner.Scan() {
 		jsonObj := make(map[string]any, len(table.Columns))
@@ -33,6 +33,5 @@ func Read(f io.Reader, table *schema.Table, sourceName string, res chan<- []any)
 		}
 		res <- jsonArr
 	}
-
 	return scanner.Err()
 }

json/write.go

@@ -7,7 +7,7 @@ import (
 	"github.com/cloudquery/plugin-sdk/schema"
 )
 
-func WriteTableBatch(w io.Writer, table *schema.Table, resources [][]any) error {
+func (*Client) WriteTableBatch(w io.Writer, table *schema.Table, resources [][]any) error {
 	for _, resource := range resources {
 		jsonObj := make(map[string]any, len(table.Columns))
 		for i := range resource {

json/write_read_test.go

@@ -16,19 +16,25 @@ func TestWriteRead(t *testing.T) {
 	if err := cqtypes[0].Set("test-source"); err != nil {
 		t.Fatal(err)
 	}
-
-	writer := bufio.NewWriter(&b)
 	transformer := &Transformer{}
 	transformedValues := schema.TransformWithTransformer(transformer, cqtypes)
-	if err := WriteTableBatch(writer, table, [][]any{transformedValues}); err != nil {
+
+	writer := bufio.NewWriter(&b)
+	reader := bufio.NewReader(&b)
+
+	cl, err := NewClient()
+	if err != nil {
+		t.Fatal(err)
+	}
+	if err := cl.WriteTableBatch(writer, table, [][]any{transformedValues}); err != nil {
 		t.Fatal(err)
 	}
 	writer.Flush()
-	reader := bufio.NewReader(&b)
+
 	ch := make(chan []any)
 	var readErr error
 	go func() {
-		readErr = Read(reader, table, "test-source", ch)
+		readErr = cl.Read(reader, table, "test-source", ch)
 		close(ch)
 	}()
 	totalCount := 0