Files
chaisql/internal/stream/operator_test.go
Asdine El Hrychy 613ca304f4 remove rowid
2025-09-07 23:29:43 +08:00

342 lines
7.6 KiB
Go

package stream_test
import (
"fmt"
"strconv"
"testing"
"github.com/chaisql/chai/internal/database"
"github.com/chaisql/chai/internal/environment"
"github.com/chaisql/chai/internal/expr"
"github.com/chaisql/chai/internal/row"
"github.com/chaisql/chai/internal/sql/parser"
"github.com/chaisql/chai/internal/stream"
"github.com/chaisql/chai/internal/stream/path"
"github.com/chaisql/chai/internal/stream/rows"
"github.com/chaisql/chai/internal/stream/table"
"github.com/chaisql/chai/internal/testutil"
"github.com/cockroachdb/errors"
"github.com/stretchr/testify/require"
)
func TestFilter(t *testing.T) {
tests := []struct {
e expr.Expr
in []expr.Row
out []row.Row
fails bool
}{
{
parser.MustParseExpr("1"),
testutil.MakeRowExprs(t, `{"a": 1}`),
testutil.MakeRows(t, `{"a": 1}`),
false,
},
{
parser.MustParseExpr("a > 1"),
testutil.MakeRowExprs(t, `{"a": 1}`),
nil,
false,
},
{
parser.MustParseExpr("a >= 1"),
testutil.MakeRowExprs(t, `{"a": 1}`),
testutil.MakeRows(t, `{"a": 1}`),
false,
},
{
parser.MustParseExpr("null"),
testutil.MakeRowExprs(t, `{"a": 1}`),
nil,
false,
},
}
for _, test := range tests {
t.Run(test.e.String(), func(t *testing.T) {
s := stream.New(rows.Emit([]string{"a"}, test.in...)).Pipe(rows.Filter(test.e))
i := 0
err := s.Iterate(new(environment.Environment), func(r database.Row) error {
testutil.RequireRowEqual(t, test.out[i], r)
i++
return nil
})
if test.fails {
require.Error(t, err)
} else {
require.NoError(t, err)
}
})
}
t.Run("String", func(t *testing.T) {
require.Equal(t, rows.Filter(parser.MustParseExpr("1")).String(), "rows.Filter(1)")
})
}
func TestTake(t *testing.T) {
tests := []struct {
inNumber int
n int
output int
fails bool
}{
{5, 1, 1, false},
{5, 7, 5, false},
{5, -1, 0, false},
}
for _, test := range tests {
t.Run(fmt.Sprintf("%d/%d", test.inNumber, test.n), func(t *testing.T) {
var ds []expr.Row
for i := 0; i < test.inNumber; i++ {
ds = append(ds, testutil.MakeRowExpr(t, `{"a": `+strconv.Itoa(i)+`}`))
}
s := stream.New(rows.Emit([]string{"a"}, ds...))
s = s.Pipe(rows.Take(parser.MustParseExpr(strconv.Itoa(test.n))))
var count int
err := s.Iterate(new(environment.Environment), func(r database.Row) error {
count++
return nil
})
if test.fails {
require.Error(t, err)
} else {
if errors.Is(err, stream.ErrStreamClosed) {
err = nil
}
require.NoError(t, err)
require.Equal(t, test.output, count)
}
})
}
t.Run("String", func(t *testing.T) {
require.Equal(t, "rows.Take(1)", rows.Take(parser.MustParseExpr("1")).String())
})
}
func TestSkip(t *testing.T) {
tests := []struct {
inNumber int
n int
output int
fails bool
}{
{5, 1, 4, false},
{5, 7, 0, false},
{5, -1, 5, false},
}
for _, test := range tests {
t.Run(fmt.Sprintf("%d/%d", test.inNumber, test.n), func(t *testing.T) {
var ds []expr.Row
for i := 0; i < test.inNumber; i++ {
ds = append(ds, testutil.MakeRowExpr(t, `{"a": `+strconv.Itoa(i)+`}`))
}
s := stream.New(rows.Emit([]string{"a"}, ds...))
s = s.Pipe(rows.Skip(parser.MustParseExpr(strconv.Itoa(test.n))))
var count int
err := s.Iterate(new(environment.Environment), func(r database.Row) error {
count++
return nil
})
if test.fails {
require.Error(t, err)
} else {
require.NoError(t, err)
require.Equal(t, test.output, count)
}
})
}
t.Run("String", func(t *testing.T) {
require.Equal(t, "rows.Skip(1)", rows.Skip(parser.MustParseExpr("1")).String())
})
}
func TestTableInsert(t *testing.T) {
tests := []struct {
name string
in stream.Operator
out []row.Row
rowid int
fails bool
}{
{
"doc with no key",
rows.Emit([]string{"a"}, testutil.MakeRowExpr(t, `{"a": 10}`), testutil.MakeRowExpr(t, `{"a": 11}`)),
[]row.Row{testutil.MakeRow(t, `{"a": 10}`), testutil.MakeRow(t, `{"a": 11}`)},
1,
false,
},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
db, tx, cleanup := testutil.NewTestTx(t)
defer cleanup()
testutil.MustExec(t, db, tx, "CREATE TABLE test (a INTEGER PRIMARY KEY)")
in := environment.New(nil, tx, nil, nil)
s := stream.New(test.in).Pipe(table.Insert("test"))
var i int
err := s.Iterate(in, func(r database.Row) error {
testutil.RequireRowEqual(t, test.out[i], r)
i++
return nil
})
if test.fails {
require.Error(t, err)
} else {
require.NoError(t, err)
}
})
}
t.Run("String", func(t *testing.T) {
require.Equal(t, "table.Insert(\"test\")", table.Insert("test").String())
})
}
func TestTableReplace(t *testing.T) {
tests := []struct {
name string
a, b any
op stream.Operator
expected testutil.Rows
fails bool
}{
{
"row with key",
1, 1,
path.Set("b", testutil.ParseExpr(t, "2")),
testutil.MakeRows(t, `{"a": 1, "b": 2}`),
false,
},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
db, tx, cleanup := testutil.NewTestTx(t)
defer cleanup()
testutil.MustExec(t, db, tx, "CREATE TABLE test (a INTEGER PRIMARY KEY, b INTEGER)")
testutil.MustExec(t, db, tx, "INSERT INTO test VALUES ($1, $2)", environment.Param{Value: test.a}, environment.Param{Value: test.b})
in := environment.New(nil, tx, nil, nil)
s := stream.New(table.Scan("test")).
Pipe(test.op).
Pipe(table.Replace("test"))
var i int
err := s.Iterate(in, func(r database.Row) error {
got, err := row.MarshalJSON(r)
require.NoError(t, err)
want, err := row.MarshalJSON(test.expected[i])
require.NoError(t, err)
require.JSONEq(t, string(want), string(got))
i++
return nil
})
if test.fails {
require.Error(t, err)
return
}
require.NoError(t, err)
res := testutil.MustQuery(t, db, tx, "SELECT * FROM test")
defer res.Close()
var got []row.Row
err = res.Iterate(func(r database.Row) error {
var fb row.ColumnBuffer
err = fb.Copy(r)
require.NoError(t, err)
got = append(got, &fb)
return nil
})
require.NoError(t, err)
test.expected.RequireEqual(t, got)
})
}
t.Run("String", func(t *testing.T) {
require.Equal(t, table.Replace("test").String(), "table.Replace(\"test\")")
})
}
func TestTableDelete(t *testing.T) {
tests := []struct {
name string
a []int
op stream.Operator
expected testutil.Rows
fails bool
}{
{
"doc with key",
[]int{1, 2, 3},
rows.Filter(testutil.ParseExpr(t, `a > 1`)),
testutil.MakeRows(t, `{"a": 1}`),
false,
},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
db, tx, cleanup := testutil.NewTestTx(t)
defer cleanup()
testutil.MustExec(t, db, tx, "CREATE TABLE test (a INTEGER PRIMARY KEY)")
for _, a := range test.a {
testutil.MustExec(t, db, tx, "INSERT INTO test VALUES ($1)", environment.Param{Value: a})
}
env := environment.New(nil, tx, nil, nil)
s := stream.New(table.Scan("test")).Pipe(test.op).Pipe(table.Delete("test"))
err := s.Iterate(env, func(r database.Row) error {
return nil
})
if test.fails {
require.Error(t, err)
} else {
require.NoError(t, err)
}
res := testutil.MustQuery(t, db, tx, "SELECT * FROM test")
defer res.Close()
var got []row.Row
err = res.Iterate(func(r database.Row) error {
var fb row.ColumnBuffer
err = fb.Copy(r)
require.NoError(t, err)
got = append(got, &fb)
return nil
})
require.NoError(t, err)
test.expected.RequireEqual(t, got)
})
}
t.Run("String", func(t *testing.T) {
require.Equal(t, table.Delete("test").String(), "table.Delete('test')")
})
}