mockable-omnisdk

cicd/requirements.txt

@@ -8,3 +8,4 @@ requests==2.32.3
 robotframework==7.0.1
 sqlalchemy==1.4.44
 tabulate==0.9.0
+cryptography==43.0.3

docker-compose-testing.yml

@@ -142,6 +142,29 @@ services:
         --host 0.0.0.0 \
         --port  \
         1090
+  omnisdk_mockserver:
+    networks:
+      testing-network:
+        aliases:
+          - omnisdk.com
+    image: stackql/testlib
+    build:
+      context: .
+      dockerfile: testLib.Dockerfile
+    ports:
+      - "1071:1071/tcp"
+    environment:
+      IS_DOCKER: "${IS_DOCKER:-false}"
+    command: 
+      - bash
+      - -c
+      - |
+        flask \
+        --app=/opt/testlib/test/python/stackql_test_tooling/flask/omnisdk/app \
+        run \
+        --host 0.0.0.0 \
+        --port  \
+        1071
   aws_mockserver:
     networks:
       testing-network:

go.mod

@@ -18,7 +18,7 @@ require (
 	github.com/spf13/cobra v1.10.2
 	github.com/spf13/pflag v1.0.9
 	github.com/spf13/viper v1.10.1
-	github.com/stackql-labs/omnisdk v0.1.0-alpha04
+	github.com/stackql-labs/omnisdk v0.1.0-alpha06
 	github.com/stackql/any-sdk v0.5.4-alpha01
 	github.com/stackql/go-suffix-map v0.0.1-alpha01
 	github.com/stackql/psql-wire v0.1.2-beta01

go.sum

@@ -348,8 +348,8 @@ github.com/spf13/viper v1.10.1 h1:nuJZuYpG7gTj/XqiUwg8bA0cp1+M2mC3J4g5luUYBKk=
 github.com/spf13/viper v1.10.1/go.mod h1:IGlFPqhNAPKRxohIzWpI5QEy4kuI7tcl5WvR+8qy1rU=
 github.com/spiffe/go-spiffe/v2 v2.7.0 h1:uXe1MflJoHw58wAUvxVlcM7WpKtijWG7I1UidcGh6g4=
 github.com/spiffe/go-spiffe/v2 v2.7.0/go.mod h1:47Q0Q9/AqGha8QLHp+kxpH4Wca7X7EnOtlIJy3mxZ3U=
-github.com/stackql-labs/omnisdk v0.1.0-alpha04 h1:DRiAQ5E/tc5VtHlAfyjlzeam65+OqPi7ogeLHzUA1mA=
-github.com/stackql-labs/omnisdk v0.1.0-alpha04/go.mod h1:SHUOryRWXeZip4JTWeA8iwmz4lYdBSp0nbWXDX29JrU=
+github.com/stackql-labs/omnisdk v0.1.0-alpha06 h1:F8+2GCow29YMTj+VsuOOOTqHeEvvUX8+wV3nHjiLy1Y=
+github.com/stackql-labs/omnisdk v0.1.0-alpha06/go.mod h1:SHUOryRWXeZip4JTWeA8iwmz4lYdBSp0nbWXDX29JrU=
 github.com/stackql/any-sdk v0.5.4-alpha01 h1:AyjD2Hyk7D8v1fHHtLpfeqQRxKuj7rC71dzi1Zh9LoA=
 github.com/stackql/any-sdk v0.5.4-alpha01/go.mod h1:BiE8uiAJMUa8n4U/yMlhvrVhE2M+5Rt8utBXYkwj9To=
 github.com/stackql/go-suffix-map v0.0.1-alpha01 h1:TDUDS8bySu41Oo9p0eniUeCm43mnRM6zFEd6j6VUaz8=

internal/stackql/intrinsic/omnisdk.go

@@ -5,6 +5,7 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"os"
 	"sort"
 	"strconv"
 	"strings"
@@ -19,7 +20,9 @@ import (
 
 const methodPredicate = "method"
 
-const streamBatchSize = 64
+// endpointEnvVar retargets omnisdk at a local mock. Transport configuration, so
+// it stays out of the query.
+const endpointEnvVar = "STACKQL_PREVIEW_ENDPOINT"
 
 func relationName(path string) string {
 	return strings.ReplaceAll(path, ".", "_")
@@ -147,7 +150,11 @@ func openStream(
 	if err != nil {
 		return nil, err
 	}
-	args := omnisdk.Args{Params: params, Auth: omnisdkAuth(ctx, resourcePath)}
+	args := omnisdk.Args{
+		Params:   params,
+		Auth:     omnisdkAuth(ctx, resourcePath),
+		Endpoint: os.Getenv(endpointEnvVar),
+	}
 	plan, err := omnisdk.Default().New(method.Path, args)
 	if err != nil {
 		return nil, err
@@ -175,19 +182,15 @@ func (rs *rowStream) Read() (sqldata.ISQLResult, error) {
 	if rs.done {
 		return rs.result(nil), io.EOF
 	}
-	batch := make([]omnisdk.Row, 0, streamBatchSize)
-	for len(batch) < streamBatchSize && rs.rows.Next() {
+	rs.done = true
+	var batch []omnisdk.Row
+	for rs.rows.Next() {
 		batch = append(batch, rs.rows.Row())
 	}
 	if err := rs.rows.Err(); err != nil {
-		rs.done = true
 		return rs.result(nil), err
 	}
-	if len(batch) < streamBatchSize {
-		rs.done = true
-		return rs.result(batch), io.EOF
-	}
-	return rs.result(batch), nil
+	return rs.result(batch), io.EOF
 }
 
 func (rs *rowStream) result(batch []omnisdk.Row) sqldata.ISQLResult {

internal/stackql/intrinsic/omnisdk_test.go

@@ -94,31 +94,16 @@ func TestRowStreamEmitsRowsInColumnOrder(t *testing.T) {
 	}
 }
 
-func TestRowStreamBatchesLargeCursor(t *testing.T) {
-	total := streamBatchSize*2 + 3
+func TestRowStreamDrainsEntireCursor(t *testing.T) {
+	total := 200
 	rows := make([]omnisdk.Row, 0, total)
 	for i := 0; i < total; i++ {
 		rows = append(rows, omnisdk.Row{"a": fmt.Sprintf("v%d", i)})
 	}
 	stream, _ := newTestStream(rows, []string{"a"})
-	var reads int
-	var seen int
-	for {
-		res, err := stream.Read()
-		reads++
-		seen += len(res.GetRows())
-		if err != nil {
-			if errors.Is(err, io.EOF) {
-				break
-			}
-			t.Fatalf("unexpected error: %v", err)
-		}
-	}
-	if seen != total {
-		t.Fatalf("read %d rows, want %d", seen, total)
-	}
-	if reads < 3 {
-		t.Fatalf("cursor of %d rows was drained in %d reads; batching is not happening", total, reads)
+	got := drain(t, stream)
+	if len(got) != total {
+		t.Fatalf("read %d rows, want %d", len(got), total)
 	}
 }