Rewrote examples test logic + local.YDB func to work with native driver instead of database/sql

Author: Viktor Pentyukhov

examples/authors/ydb/db_test.go

@@ -6,6 +6,7 @@ import (
 
 	"github.com/sqlc-dev/sqlc/internal/sqltest/local"
 	_ "github.com/ydb-platform/ydb-go-sdk/v3"
+	"github.com/ydb-platform/ydb-go-sdk/v3/query"
 )
 
 func ptr(s string) *string {
@@ -15,10 +16,10 @@ func ptr(s string) *string {
 func TestAuthors(t *testing.T) {
 	ctx := context.Background()
 
-	test := local.YDB(t, []string{"schema.sql"})
-	defer test.DB.Close()
+	db := local.YDB(t, []string{"schema.sql"})
+	defer db.Close(ctx)
 
-	q := New(test.DB)
+	q := New(db.Query())
 
 	t.Run("InsertAuthors", func(t *testing.T) {
 		authorsToInsert := []CreateOrUpdateAuthorParams{
@@ -38,7 +39,7 @@ func TestAuthors(t *testing.T) {
 		}
 
 		for _, author := range authorsToInsert {
-			if _, err := q.CreateOrUpdateAuthor(ctx, author); err != nil {
+			if err := q.CreateOrUpdateAuthor(ctx, author, query.WithIdempotent()); err != nil {
 				t.Fatalf("failed to insert author %q: %v", author.P1, err)
 			}
 		}
@@ -52,7 +53,7 @@ func TestAuthors(t *testing.T) {
 			P2: &newBio,
 		}
 
-		returnedBio, err := q.CreateOrUpdateAuthorReturningBio(ctx, arg)
+		returnedBio, err := q.CreateOrUpdateAuthorReturningBio(ctx, arg, query.WithIdempotent())
 		if err != nil {
 			t.Fatalf("failed to create or update author: %v", err)
 		}
@@ -74,15 +75,17 @@ func TestAuthors(t *testing.T) {
 			P2: 10,
 		}
 
-		singleAuthor, err := q.UpdateAuthorByID(ctx, arg)
+		authors, err := q.UpdateAuthorByID(ctx, arg, query.WithIdempotent(), query.WithIdempotent())
 		if err != nil {
 			t.Fatal(err)
 		}
-		bio := "Null"
-		if singleAuthor.Bio != nil {
-			bio = *singleAuthor.Bio
+		for _, a := range authors {
+			bio := "Null"
+			if a.Bio != nil {
+				bio = *a.Bio
+			}
+			t.Logf("- ID: %d | Name: %s | Bio: %s", a.ID, a.Name, bio)
 		}
-		t.Logf("- ID: %d | Name: %s | Bio: %s", singleAuthor.ID, singleAuthor.Name, bio)
 	})
 
 	t.Run("ListAuthors", func(t *testing.T) {
@@ -154,7 +157,7 @@ func TestAuthors(t *testing.T) {
 	t.Run("Delete All Authors", func(t *testing.T) {
 		var i uint64
 		for i = 1; i <= 13; i++ {
-			if err := q.DeleteAuthor(ctx, i); err != nil {
+			if err := q.DeleteAuthor(ctx, i, query.WithIdempotent()); err != nil {
 				t.Fatalf("failed to delete authors: %v", err)
 			}
 		}

examples/authors/ydb/query.sql

@@ -16,11 +16,17 @@ WHERE bio IS NULL;
 -- name: Count :one
 SELECT COUNT(*) FROM authors;
 
--- name: UpsertAuthor :queryrows
-UPSERT INTO authors (id, name, bio) VALUES ($p0, $p1, $p2) RETURNING *;
+-- name: Coalesce :many
+SELECT id, name, COALESCE(bio, 'Null value!') FROM authors;
+
+-- name: CreateOrUpdateAuthor :exec 
+UPSERT INTO authors (id, name, bio) VALUES ($p0, $p1, $p2);
+
+-- name: CreateOrUpdateAuthorReturningBio :one
+UPSERT INTO authors (id, name, bio) VALUES ($p0, $p1, $p2) RETURNING bio;
 
 -- name: DeleteAuthor :exec 
 DELETE FROM authors WHERE id = $p0;
 
--- name: UpdateAuthorByID :one
-UPDATE authors SET name = $p0, bio = $p1 WHERE id = $p2 RETURNING *;
+-- name: UpdateAuthorByID :queryrows
+UPDATE authors SET name = $p0, bio = $p1 WHERE id = $p2 RETURNING *;

examples/authors/ydb/query.sql.go

@@ -2,7 +2,6 @@ package local
 
 import (
 	"context"
-	"database/sql"
 	"fmt"
 	"hash/fnv"
 	"math/rand"
@@ -14,30 +13,25 @@ import (
 	migrate "github.com/sqlc-dev/sqlc/internal/migrations"
 	"github.com/sqlc-dev/sqlc/internal/sql/sqlpath"
 	"github.com/ydb-platform/ydb-go-sdk/v3"
+	"github.com/ydb-platform/ydb-go-sdk/v3/query"
+	"github.com/ydb-platform/ydb-go-sdk/v3/sugar"
 )
 
 func init() {
 	rand.Seed(time.Now().UnixNano())
 }
 
-func YDB(t *testing.T, migrations []string) TestYDB {
+func YDB(t *testing.T, migrations []string) *ydb.Driver {
 	return link_YDB(t, migrations, true)
 }
 
-func ReadOnlyYDB(t *testing.T, migrations []string) TestYDB {
+func ReadOnlyYDB(t *testing.T, migrations []string) *ydb.Driver {
 	return link_YDB(t, migrations, false)
 }
 
-type TestYDB struct {
-	DB     *sql.DB
-	Prefix string
-}
-
-func link_YDB(t *testing.T, migrations []string, rw bool) TestYDB {
+func link_YDB(t *testing.T, migrations []string, rw bool) *ydb.Driver {
 	t.Helper()
 
-	time.Sleep(1 * time.Second) // wait for YDB to start
-
 	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
 	defer cancel()
 
@@ -72,46 +66,43 @@ func link_YDB(t *testing.T, migrations []string, rw bool) TestYDB {
 
 	var name string
 	if rw {
-		// name = fmt.Sprintf("sqlc_test_%s", id())
 		name = fmt.Sprintf("sqlc_test_%s", "test_new")
 	} else {
 		name = fmt.Sprintf("sqlc_test_%x", h.Sum(nil))
 	}
-	prefix := fmt.Sprintf("%s/%s", baseDB, name)
 
-	rootDSN := fmt.Sprintf("grpc://%s?database=%s", dbuiri, baseDB)
-	t.Logf("→ Opening root driver: %s", rootDSN)
-	driver, err := ydb.Open(ctx, rootDSN,
+	connectionString := fmt.Sprintf("grpc://%s%s", dbuiri, baseDB)
+	t.Logf("→ Opening YDB connection: %s", connectionString)
+
+	db, err := ydb.Open(ctx, connectionString,
 		ydb.WithInsecure(),
 		ydb.WithDiscoveryInterval(time.Hour),
 		ydb.WithNodeAddressMutator(func(_ string) string {
 			return host
 		}),
 	)
 	if err != nil {
-		t.Fatalf("failed to open root YDB connection: %s", err)
+		t.Fatalf("failed to open YDB connection: %s", err)
 	}
 
-	connector, err := ydb.Connector(
-		driver,
-		ydb.WithTablePathPrefix(prefix),
-		ydb.WithAutoDeclare(),
-		ydb.WithNumericArgs(),
-	)
+	prefix := fmt.Sprintf("%s/%s", baseDB, name)
+	t.Logf("→ Using prefix: %s", prefix)
+
+	err = sugar.RemoveRecursive(ctx, db, prefix)
 	if err != nil {
-		t.Fatalf("failed to create connector: %s", err)
+		t.Logf("Warning: failed to remove old data: %s", err)
 	}
 
-	db := sql.OpenDB(connector)
-
 	t.Log("→ Applying migrations to prefix: ", prefix)
 
-	schemeCtx := ydb.WithQueryMode(ctx, ydb.SchemeQueryMode)
 	for _, stmt := range seed {
-		_, err := db.ExecContext(schemeCtx, stmt)
+		err := db.Query().Exec(ctx, stmt,
+			query.WithTxControl(query.EmptyTxControl()),
+		)
 		if err != nil {
 			t.Fatalf("failed to apply migration: %s\nSQL: %s", err, stmt)
 		}
 	}
-	return TestYDB{DB: db, Prefix: prefix}
+
+	return db
 }