From 0aed30b06503d2cca9a137d681b6d079d7737f68 Mon Sep 17 00:00:00 2001
From: Satyam Singh <trueleo@protonmail.com>
Date: Tue, 12 Sep 2023 19:58:24 +0530
Subject: [PATCH] Fix

---
 pkg/model/query.go | 37 +++++++++++++++++++++++++++++--------
 1 file changed, 29 insertions(+), 8 deletions(-)

diff --git a/pkg/model/query.go b/pkg/model/query.go
index d80bf92..91cca66 100644
--- a/pkg/model/query.go
+++ b/pkg/model/query.go
@@ -27,6 +27,7 @@ import (
 	"pb/pkg/iterator"
 	"regexp"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/charmbracelet/bubbles/help"
@@ -155,6 +156,11 @@ func (m *QueryModel) currentFocus() string {
 }
 
 func (m *QueryModel) initIterator() {
+	iter := createIteratorFromModel(m)
+	m.queryIterator = iter
+}
+
+func createIteratorFromModel(m *QueryModel) *iterator.QueryIterator[QueryData, FetchResult] {
 	startTime := m.timeRange.start.Time()
 	endTime := m.timeRange.end.Time()
 
@@ -164,11 +170,10 @@ func (m *QueryModel) initIterator() {
 	regex := regexp.MustCompile(`^select\s+(?:\*|\w+(?:,\s*\w+)*)\s+from\s+(\w+)(?:\s+;)?$`)
 	matches := regex.FindStringSubmatch(m.query.Value())
 	if matches == nil {
-		m.queryIterator = nil
-		return
+		return nil
 	}
 	table := matches[1]
-	queryIterator := iterator.NewQueryIterator(
+	iter := iterator.NewQueryIterator(
 		startTime, endTime,
 		false,
 		func(t1, t2 time.Time) (QueryData, FetchResult) {
@@ -188,7 +193,7 @@ func (m *QueryModel) initIterator() {
 			count := res.Records[0]["count"].(float64)
 			return count > 0
 		})
-	m.queryIterator = &queryIterator
+	return &iter
 }
 
 func NewQueryModel(profile config.Profile, stream string, duration uint) QueryModel {
@@ -230,7 +235,7 @@ func NewQueryModel(profile config.Profile, stream string, duration uint) QueryMo
 	help := help.New()
 	help.Styles.FullDesc = lipgloss.NewStyle().Foreground(FocusSecondry)
 
-	return QueryModel{
+	model := QueryModel{
 		width:         w,
 		height:        h,
 		table:         table,
@@ -242,12 +247,28 @@ func NewQueryModel(profile config.Profile, stream string, duration uint) QueryMo
 		queryIterator: nil,
 		status:        NewStatusBar(profile.URL, stream, w),
 	}
+	model.queryIterator = createIteratorFromModel(&model)
+	return model
 }
 
 func (m QueryModel) Init() tea.Cmd {
-	// Just return `nil`, which means "no I/O right now, please."
-	m.initIterator()
-	return NewFetchTask(m.profile, m.query.Value(), m.timeRange.StartValueUtc(), m.timeRange.EndValueUtc())
+	return func() tea.Msg {
+		var ready sync.WaitGroup
+		ready.Add(1)
+		go func() {
+			m.initIterator()
+			for !m.queryIterator.Ready() {
+				time.Sleep(time.Millisecond * 100)
+			}
+			ready.Done()
+		}()
+		ready.Wait()
+		if m.queryIterator.Finished() {
+			return nil
+		}
+
+		return IteratorNext(m.queryIterator)()
+	}
 }
 
 func (m QueryModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {