From 986c69a812ff037c0ef891ef6fd76295c6d7a000 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Toma=C5=BE=20Jerman?= <tomaz.jerman@kendu.si>
Date: Thu, 29 Aug 2019 10:07:52 +0200
Subject: [PATCH] Provide entry count

---
 compose/decoder/decoder.go         | 28 ++++++++++++++++++++++++++--
 compose/internal/service/record.go |  2 ++
 compose/rest/record.go             | 18 +++++++++++++++---
 pkg/count/file.go                  | 28 ++++++++++++++++++++++++++++
 4 files changed, 71 insertions(+), 5 deletions(-)
 create mode 100644 pkg/count/file.go

diff --git a/compose/decoder/decoder.go b/compose/decoder/decoder.go
index 8159f8b49..4a9387644 100644
--- a/compose/decoder/decoder.go
+++ b/compose/decoder/decoder.go
@@ -2,6 +2,8 @@ package decoder
 
 import (
 	"io"
+
+	"github.com/cortezaproject/corteza-server/pkg/count"
 )
 
 type (
@@ -17,12 +19,14 @@ type (
 	}
 
 	flatReader struct {
+		f      io.ReadSeeker
 		r      FlatReader
 		header []string
 		more   bool
 	}
 
 	structuredDecoder struct {
+		f      io.ReadSeeker
 		header []string
 		d      StructuredDecoder
 		buf    []map[string]interface{}
@@ -34,13 +38,27 @@ type (
 )
 
 // flat reader
-func NewFlatReader(r FlatReader) *flatReader {
+func NewFlatReader(r FlatReader, f io.ReadSeeker) *flatReader {
 	return &flatReader{
+		f:    f,
 		r:    r,
 		more: true,
 	}
 }
 
+func (dec *flatReader) EntryCount() (uint64, error) {
+	defer dec.f.Seek(0, 0)
+
+	c, err := count.Lines(dec.f)
+	if err != nil {
+		return 0, err
+	}
+	if c <= 0 {
+		return 0, nil
+	}
+	return c - 1, nil
+}
+
 func (dec *flatReader) get(fnc fdCallback) error {
 	v, err := dec.r.Read()
 	if err == io.EOF {
@@ -76,12 +94,18 @@ func (dec *flatReader) Header() []string {
 }
 
 // structured decoder
-func NewStructuredDecoder(d StructuredDecoder) *structuredDecoder {
+func NewStructuredDecoder(d StructuredDecoder, f io.ReadSeeker) *structuredDecoder {
 	return &structuredDecoder{
+		f: f,
 		d: d,
 	}
 }
 
+func (dec *structuredDecoder) EntryCount() (uint64, error) {
+	defer dec.f.Seek(0, 0)
+	return count.Lines(dec.f)
+}
+
 func (dec *structuredDecoder) get(fnc sdCallback) error {
 	if !dec.d.More() {
 		return nil
diff --git a/compose/internal/service/record.go b/compose/internal/service/record.go
index c67472c08..8f20bc0fd 100644
--- a/compose/internal/service/record.go
+++ b/compose/internal/service/record.go
@@ -77,6 +77,7 @@ type (
 
 	Decoder interface {
 		Header() []string
+		EntryCount() (uint64, error)
 		Records(fields map[string]string, Create decoder.RecordCreator) error
 	}
 
@@ -96,6 +97,7 @@ type (
 	RecordImportProgress struct {
 		StartedAt  *time.Time `json:"startedAt"`
 		FinishedAt *time.Time `json:"finishedAt"`
+		EntryCount uint64     `json:"entryCount"`
 		Completed  uint64     `json:"completed"`
 		Failed     uint64     `json:"failed"`
 		FailReason error      `json:"failReason"`
diff --git a/compose/rest/record.go b/compose/rest/record.go
index 0de13d17e..b503f967d 100644
--- a/compose/rest/record.go
+++ b/compose/rest/record.go
@@ -175,6 +175,7 @@ func (ctrl *Record) ImportInit(ctx context.Context, r *request.RecordImportInit)
 	var (
 		err           error
 		recordDecoder service.Decoder
+		entryCount    uint64
 	)
 
 	// Access control.
@@ -204,15 +205,19 @@ func (ctrl *Record) ImportInit(ctx context.Context, r *request.RecordImportInit)
 	// determine decoder
 	switch strings.ToLower(ext) {
 	case "json", "jsonl", "ldjson", "ndjson":
-		recordDecoder = decoder.NewStructuredDecoder(json.NewDecoder(f))
+		recordDecoder = decoder.NewStructuredDecoder(json.NewDecoder(f), f)
 
 	case "csv":
-		recordDecoder = decoder.NewFlatReader(csv.NewReader(f))
+		recordDecoder = decoder.NewFlatReader(csv.NewReader(f), f)
 
 	default:
 		return nil, errors.New(fmt.Sprintf("unsupported format (\"%s\")", ext))
 
 	}
+	entryCount, err = recordDecoder.EntryCount()
+	if err != nil {
+		return nil, err
+	}
 
 	header := recordDecoder.Header()
 	hh := make(map[string]string)
@@ -220,7 +225,14 @@ func (ctrl *Record) ImportInit(ctx context.Context, r *request.RecordImportInit)
 		hh[h] = ""
 	}
 
-	return ctrl.importSession.SetRecordByID(ctx, 0, r.NamespaceID, r.ModuleID, hh, nil, recordDecoder)
+	return ctrl.importSession.SetRecordByID(
+		ctx,
+		0,
+		r.NamespaceID,
+		r.ModuleID,
+		hh,
+		&service.RecordImportProgress{EntryCount: entryCount},
+		recordDecoder)
 }
 
 func (ctrl *Record) ImportRun(ctx context.Context, r *request.RecordImportRun) (interface{}, error) {
diff --git a/pkg/count/file.go b/pkg/count/file.go
new file mode 100644
index 000000000..6c1a88597
--- /dev/null
+++ b/pkg/count/file.go
@@ -0,0 +1,28 @@
+package count
+
+import (
+	"bytes"
+	"io"
+)
+
+// Lines provides a line count
+//
+// https://stackoverflow.com/a/24563853
+func Lines(r io.ReadSeeker) (count uint64, err error) {
+	defer r.Seek(0, 0)
+	buf := make([]byte, 32*1024)
+	lineSep := []byte{'\n'}
+
+	for {
+		c, err := r.Read(buf)
+		count += uint64(bytes.Count(buf[:c], lineSep))
+
+		switch {
+		case err == io.EOF:
+			return count, nil
+
+		case err != nil:
+			return count, err
+		}
+	}
+}