Skip to content

Commit

Permalink
Use iterator for select queries
Browse files Browse the repository at this point in the history
  • Loading branch information
@trueleo
trueleo committed Sep 12, 2023
1 parent 556e507 commit a451a17
Showing 1 changed file with 152 additions and 24 deletions.
176 changes: 152 additions & 24 deletions pkg/model/query.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,9 @@ import (
"net/http"
"os"
"pb/pkg/config"
"pb/pkg/iterator"
"regexp"
"strings"
"time"

"github.com/charmbracelet/bubbles/help"
Expand Down Expand Up @@ -117,16 +120,17 @@ const (
)

type QueryModel struct {
width int
height int
table table.Model
query textarea.Model
timeRange TimeInputModel
profile config.Profile
help help.Model
status StatusBar
overlay uint
focused int
width int
height int
table table.Model
query textarea.Model
timeRange TimeInputModel
profile config.Profile
help help.Model
status StatusBar
queryIterator *iterator.QueryIterator[QueryData, FetchResult]
overlay uint
focused int
}

func (m *QueryModel) focusSelected() {
Expand All @@ -145,6 +149,43 @@ func (m *QueryModel) currentFocus() string {
return QueryNavigationMap[m.focused]
}

func (m *QueryModel) initIterator() {
startTime := m.timeRange.start.Time()
endTime := m.timeRange.end.Time()

startTime = startTime.Truncate(time.Minute)
endTime = endTime.Truncate(time.Minute).Add(time.Minute)

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
}
table := matches[1]
queryIterator := iterator.NewQueryIterator(
startTime, endTime,
false,
func(t1, t2 time.Time) (QueryData, FetchResult) {
client := &http.Client{
Timeout: time.Second * 50,
}
return fetchData(client, &m.profile, m.query.Value(), t1.UTC().Format(time.RFC3339), t2.UTC().Format(time.RFC3339))
},
func(t1, t2 time.Time) bool {
client := &http.Client{
Timeout: time.Second * 50,
}
res, err := fetchData(client, &m.profile, "select count(*) as count from "+table, m.timeRange.StartValueUtc(), m.timeRange.EndValueUtc())
if err == fetchErr {
return false
}
count := res.Records[0]["count"].(float64)
return count > 0
})
m.queryIterator = &queryIterator
}

func NewQueryModel(profile config.Profile, stream string, duration uint) QueryModel {
w, h, _ := term.GetSize(int(os.Stdout.Fd()))

Expand Down Expand Up @@ -185,20 +226,22 @@ func NewQueryModel(profile config.Profile, stream string, duration uint) QueryMo
help.Styles.FullDesc = lipgloss.NewStyle().Foreground(FocusSecondry)

return QueryModel{
width: w,
height: h,
table: table,
query: query,
timeRange: inputs,
overlay: overlayNone,
profile: profile,
help: help,
status: NewStatusBar(profile.URL, stream, w),
width: w,
height: h,
table: table,
query: query,
timeRange: inputs,
overlay: overlayNone,
profile: profile,
help: help,
queryIterator: nil,
status: NewStatusBar(profile.URL, stream, w),
}
}

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())
}

Expand Down Expand Up @@ -256,7 +299,21 @@ func (m QueryModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
// common keybind
if msg.Type == tea.KeyCtrlR {
m.overlay = overlayNone
return m, NewFetchTask(m.profile, m.query.Value(), m.timeRange.StartValueUtc(), m.timeRange.EndValueUtc())
if m.queryIterator == nil {
return m, NewFetchTask(m.profile, m.query.Value(), m.timeRange.StartValueUtc(), m.timeRange.EndValueUtc())
}
if m.queryIterator.Ready() && !m.queryIterator.Finished() {
return m, IteratorNext(m.queryIterator)
}
return m, nil
}

if msg.Type == tea.KeyCtrlB {
m.overlay = overlayNone
if m.queryIterator.CanFetchPrev() {
return m, IteratorPrev(m.queryIterator)
}
return m, nil
}

switch msg.Type {
Expand All @@ -269,12 +326,14 @@ func (m QueryModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
switch m.currentFocus() {
case "query":
m.query, cmd = m.query.Update(msg)
m.initIterator()
case "table":
m.table, cmd = m.table.Update(msg)
}
cmds = append(cmds, cmd)
case overlayInputs:
m.timeRange, cmd = m.timeRange.Update(msg)
m.initIterator()
cmds = append(cmds, cmd)
}
}
Expand Down Expand Up @@ -314,12 +373,33 @@ func (m QueryModel) View() string {
BorderForeground(FocusPrimary)
}

mainViewRenderElements := []string{lipgloss.JoinHorizontal(lipgloss.Top, queryOuter.Render(m.query.View()), timeOuter.Render(time)), tableOuter.Render(m.table.View())}

if m.queryIterator != nil {
inactiveStyle := lipgloss.NewStyle().Foreground(StandardPrimary)
activeStyle := lipgloss.NewStyle().Foreground(FocusPrimary)
var line strings.Builder

if m.queryIterator.CanFetchPrev() {
line.WriteString(activeStyle.Render("<<"))
} else {
line.WriteString(inactiveStyle.Render("<<"))
}

fmt.Fprintf(&line, " %d of many ", m.table.TotalRows())

if m.queryIterator.Ready() && !m.queryIterator.Finished() {
line.WriteString(activeStyle.Render(">>"))
} else {
line.WriteString(inactiveStyle.Render(">>"))
}

mainViewRenderElements = append(mainViewRenderElements, line.String())
}

switch m.overlay {
case overlayNone:
mainView = lipgloss.JoinVertical(lipgloss.Left,
lipgloss.JoinHorizontal(lipgloss.Top, queryOuter.Render(m.query.View()), timeOuter.Render(time)),
tableOuter.Render(m.table.View()),
)
mainView = lipgloss.JoinVertical(lipgloss.Left, mainViewRenderElements...)
switch m.currentFocus() {
case "query":
helpKeys = TextAreaHelpKeys{}.FullHelp()
Expand All @@ -334,7 +414,15 @@ func (m QueryModel) View() string {
mainView = m.timeRange.View()
helpKeys = m.timeRange.FullHelp()
}

helpKeys = append(helpKeys, additionalKeyBinds)
if m.queryIterator != nil {
helpKeys = append(helpKeys, []key.Binding{
key.NewBinding(key.WithKeys("ctrl+r"), key.WithHelp("ctrl enter", "Fetch Next Minute")),
key.NewBinding(key.WithKeys("ctrl+b"), key.WithHelp("ctrl b", "Fetch Prev Minute")),
})
}

helpView = m.help.FullHelpView(helpKeys)

helpHeight := lipgloss.Height(helpView)
Expand Down Expand Up @@ -377,6 +465,46 @@ func NewFetchTask(profile config.Profile, query string, startTime string, endTim
}
}

func IteratorNext(iter *iterator.QueryIterator[QueryData, FetchResult]) func() tea.Msg {
return func() tea.Msg {
res := FetchData{
status: fetchErr,
schema: []string{},
data: []map[string]interface{}{},
}

data, status := iter.Next()

if status == fetchOk {
res.data = data.Records
res.schema = data.Fields
res.status = fetchOk
}

return res
}
}

func IteratorPrev(iter *iterator.QueryIterator[QueryData, FetchResult]) func() tea.Msg {
return func() tea.Msg {
res := FetchData{
status: fetchErr,
schema: []string{},
data: []map[string]interface{}{},
}

data, status := iter.Prev()

if status == fetchOk {
res.data = data.Records
res.schema = data.Fields
res.status = fetchOk
}

return res
}
}

func fetchData(client *http.Client, profile *config.Profile, query string, startTime string, endTime string) (data QueryData, res FetchResult) {
data = QueryData{}
res = fetchErr
Expand Down

0 comments on commit a451a17

Please sign in to comment.