snapshot of https://github.com/sqlc-dev/sqlc.git (main, 2026-09-16; blobs > 131072B omitted)

internal/x/expander/expander_test.go

@@ -0,0 +1,435 @@
+package expander
+
+import (
+	"context"
+	"database/sql"
+	"os"
+	"testing"
+
+	_ "github.com/go-sql-driver/mysql"
+	"github.com/jackc/pgx/v5/pgxpool"
+	"github.com/ncruces/go-sqlite3"
+
+	"github.com/sqlc-dev/sqlc/internal/engine/dolphin"
+	"github.com/sqlc-dev/sqlc/internal/engine/postgresql"
+	"github.com/sqlc-dev/sqlc/internal/engine/sqlite"
+	"github.com/sqlc-dev/sqlc/internal/sqltest/docker"
+	"github.com/sqlc-dev/sqlc/internal/sqltest/native"
+)
+
+// PostgreSQLColumnGetter implements ColumnGetter for PostgreSQL using pgxpool.
+type PostgreSQLColumnGetter struct {
+	pool *pgxpool.Pool
+}
+
+func (g *PostgreSQLColumnGetter) GetColumnNames(ctx context.Context, query string) ([]string, error) {
+	conn, err := g.pool.Acquire(ctx)
+	if err != nil {
+		return nil, err
+	}
+	defer conn.Release()
+
+	desc, err := conn.Conn().Prepare(ctx, "", query)
+	if err != nil {
+		return nil, err
+	}
+
+	columns := make([]string, len(desc.Fields))
+	for i, field := range desc.Fields {
+		columns[i] = field.Name
+	}
+
+	return columns, nil
+}
+
+// MySQLColumnGetter implements ColumnGetter for MySQL. Column names are read
+// from the result set metadata returned by executing the query; the test
+// tables are empty, so no real rows are transferred.
+//
+// An earlier implementation pulled column metadata straight out of a prepared
+// statement via a forked mysql driver exposing StmtMetadata. That fork
+// required a `replace` directive in go.mod, which broke `go install
+// github.com/sqlc-dev/sqlc/cmd/sqlc@latest` (see
+// https://github.com/sqlc-dev/sqlc/issues/4397). Reading columns from sql.Rows
+// works with the upstream driver and keeps the test covering the same
+// behavior.
+type MySQLColumnGetter struct {
+	db *sql.DB
+}
+
+func (g *MySQLColumnGetter) GetColumnNames(ctx context.Context, query string) ([]string, error) {
+	rows, err := g.db.QueryContext(ctx, query)
+	if err != nil {
+		return nil, err
+	}
+	defer rows.Close()
+	return rows.Columns()
+}
+
+// SQLiteColumnGetter implements ColumnGetter for SQLite using the native ncruces/go-sqlite3 API.
+type SQLiteColumnGetter struct {
+	conn *sqlite3.Conn
+}
+
+func (g *SQLiteColumnGetter) GetColumnNames(ctx context.Context, query string) ([]string, error) {
+	// Prepare the statement - this gives us column metadata without executing
+	stmt, _, err := g.conn.Prepare(query)
+	if err != nil {
+		return nil, err
+	}
+	defer stmt.Close()
+
+	// Get column names from the prepared statement
+	count := stmt.ColumnCount()
+	columns := make([]string, count)
+	for i := range count {
+		columns[i] = stmt.ColumnName(i)
+	}
+
+	return columns, nil
+}
+
+func TestExpandPostgreSQL(t *testing.T) {
+	ctx := context.Background()
+
+	uri := os.Getenv("POSTGRESQL_SERVER_URI")
+	if uri == "" {
+		if err := docker.Installed(); err == nil {
+			u, err := docker.StartPostgreSQLServer(ctx)
+			if err != nil {
+				t.Fatal(err)
+			}
+			uri = u
+		} else if err := native.Supported(); err == nil {
+			u, err := native.StartPostgreSQLServer(ctx)
+			if err != nil {
+				t.Fatal(err)
+			}
+			uri = u
+		} else {
+			t.Skip("POSTGRESQL_SERVER_URI is empty and neither Docker nor native installation is available")
+		}
+	}
+
+	pool, err := pgxpool.New(ctx, uri)
+	if err != nil {
+		t.Skipf("could not connect to database: %v", err)
+	}
+	defer pool.Close()
+
+	// Create a test table
+	_, err = pool.Exec(ctx, `
+		DROP TABLE IF EXISTS authors;
+		CREATE TABLE authors (
+			id SERIAL PRIMARY KEY,
+			name TEXT NOT NULL,
+			bio TEXT
+		);
+	`)
+	if err != nil {
+		t.Fatalf("failed to create test table: %v", err)
+	}
+	defer pool.Exec(ctx, "DROP TABLE IF EXISTS authors")
+
+	// Create the parser which also implements format.Dialect
+	parser := postgresql.NewParser()
+
+	// Create the expander
+	colGetter := &PostgreSQLColumnGetter{pool: pool}
+	exp := New(colGetter, parser, parser)
+
+	tests := []struct {
+		name     string
+		query    string
+		expected string
+	}{
+		{
+			name:     "simple select star",
+			query:    "SELECT * FROM authors",
+			expected: "SELECT id, name, bio FROM authors;",
+		},
+		{
+			name:     "select with no star",
+			query:    "SELECT id, name FROM authors",
+			expected: "SELECT id, name FROM authors", // No change, returns original
+		},
+		{
+			name:     "select star with where clause",
+			query:    "SELECT * FROM authors WHERE id = 1",
+			expected: "SELECT id, name, bio FROM authors WHERE id = 1;",
+		},
+		{
+			name:     "double star",
+			query:    "SELECT *, * FROM authors",
+			expected: "SELECT id, name, bio, id, name, bio FROM authors;",
+		},
+		{
+			name:     "table qualified star",
+			query:    "SELECT authors.* FROM authors",
+			expected: "SELECT authors.id, authors.name, authors.bio FROM authors;",
+		},
+		{
+			name:     "star in middle of columns",
+			query:    "SELECT id, *, name FROM authors",
+			expected: "SELECT id, id, name, bio, name FROM authors;",
+		},
+		{
+			name:     "insert returning star",
+			query:    "INSERT INTO authors (name, bio) VALUES ('John', 'A writer') RETURNING *",
+			expected: "INSERT INTO authors (name, bio) VALUES ('John', 'A writer') RETURNING id, name, bio;",
+		},
+		{
+			name:     "insert returning mixed",
+			query:    "INSERT INTO authors (name, bio) VALUES ('John', 'A writer') RETURNING id, *",
+			expected: "INSERT INTO authors (name, bio) VALUES ('John', 'A writer') RETURNING id, id, name, bio;",
+		},
+		{
+			name:     "update returning star",
+			query:    "UPDATE authors SET name = 'Jane' WHERE id = 1 RETURNING *",
+			expected: "UPDATE authors SET name = 'Jane' WHERE id = 1 RETURNING id, name, bio;",
+		},
+		{
+			name:     "delete returning star",
+			query:    "DELETE FROM authors WHERE id = 1 RETURNING *",
+			expected: "DELETE FROM authors WHERE id = 1 RETURNING id, name, bio;",
+		},
+		{
+			name:     "cte with select star",
+			query:    "WITH a AS (SELECT * FROM authors) SELECT * FROM a",
+			expected: "WITH a AS (SELECT id, name, bio FROM authors) SELECT id, name, bio FROM a;",
+		},
+		{
+			name:     "multiple ctes with dependency",
+			query:    "WITH a AS (SELECT * FROM authors), b AS (SELECT * FROM a) SELECT * FROM b",
+			expected: "WITH a AS (SELECT id, name, bio FROM authors), b AS (SELECT id, name, bio FROM a) SELECT id, name, bio FROM b;",
+		},
+		{
+			name:     "count star not expanded",
+			query:    "SELECT COUNT(*) FROM authors",
+			expected: "SELECT COUNT(*) FROM authors", // No change - COUNT(*) should not be expanded
+		},
+		{
+			name:     "count star with other columns",
+			query:    "SELECT COUNT(*), name FROM authors GROUP BY name",
+			expected: "SELECT COUNT(*), name FROM authors GROUP BY name", // No change
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			result, err := exp.Expand(ctx, tc.query)
+			if err != nil {
+				t.Fatalf("Expand failed: %v", err)
+			}
+			if result != tc.expected {
+				t.Errorf("expected %q, got %q", tc.expected, result)
+			}
+		})
+	}
+}
+
+func TestExpandMySQL(t *testing.T) {
+	ctx := context.Background()
+
+	source := os.Getenv("MYSQL_SERVER_URI")
+	if source == "" {
+		if err := docker.Installed(); err == nil {
+			u, err := docker.StartMySQLServer(ctx)
+			if err != nil {
+				t.Fatal(err)
+			}
+			source = u
+		} else if err := native.Supported(); err == nil {
+			u, err := native.StartMySQLServer(ctx)
+			if err != nil {
+				t.Fatal(err)
+			}
+			source = u
+		} else {
+			t.Skip("MYSQL_SERVER_URI is empty and neither Docker nor native installation is available")
+		}
+	}
+
+	db, err := sql.Open("mysql", source)
+	if err != nil {
+		t.Skipf("could not connect to MySQL: %v", err)
+	}
+	defer db.Close()
+
+	// Verify connection
+	if err := db.Ping(); err != nil {
+		t.Skipf("could not ping MySQL: %v", err)
+	}
+
+	// Create a test table
+	_, err = db.ExecContext(ctx, `DROP TABLE IF EXISTS authors`)
+	if err != nil {
+		t.Fatalf("failed to drop test table: %v", err)
+	}
+	_, err = db.ExecContext(ctx, `
+		CREATE TABLE authors (
+			id INT AUTO_INCREMENT PRIMARY KEY,
+			name VARCHAR(255) NOT NULL,
+			bio TEXT
+		)
+	`)
+	if err != nil {
+		t.Fatalf("failed to create test table: %v", err)
+	}
+	defer db.ExecContext(ctx, "DROP TABLE IF EXISTS authors")
+
+	// Create the parser which also implements format.Dialect
+	parser := dolphin.NewParser()
+
+	// Create the expander
+	colGetter := &MySQLColumnGetter{db: db}
+	exp := New(colGetter, parser, parser)
+
+	tests := []struct {
+		name     string
+		query    string
+		expected string
+	}{
+		{
+			name:     "simple select star",
+			query:    "SELECT * FROM authors",
+			expected: "SELECT id, name, bio FROM authors;",
+		},
+		{
+			name:     "select with no star",
+			query:    "SELECT id, name FROM authors",
+			expected: "SELECT id, name FROM authors", // No change, returns original
+		},
+		{
+			name:     "select star with where clause",
+			query:    "SELECT * FROM authors WHERE id = 1",
+			expected: "SELECT id, name, bio FROM authors WHERE id = 1;",
+		},
+		{
+			name:     "table qualified star",
+			query:    "SELECT authors.* FROM authors",
+			expected: "SELECT authors.id, authors.name, authors.bio FROM authors;",
+		},
+		{
+			name:     "double table qualified star",
+			query:    "SELECT authors.*, authors.* FROM authors",
+			expected: "SELECT authors.id, authors.name, authors.bio, authors.id, authors.name, authors.bio FROM authors;",
+		},
+		{
+			name:     "star in middle of columns table qualified",
+			query:    "SELECT id, authors.*, name FROM authors",
+			expected: "SELECT id, authors.id, authors.name, authors.bio, name FROM authors;",
+		},
+		{
+			name:     "count star not expanded",
+			query:    "SELECT COUNT(*) FROM authors",
+			expected: "SELECT COUNT(*) FROM authors", // No change - COUNT(*) should not be expanded
+		},
+		{
+			name:     "count star with other columns",
+			query:    "SELECT COUNT(*), name FROM authors GROUP BY name",
+			expected: "SELECT COUNT(*), name FROM authors GROUP BY name", // No change
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			result, err := exp.Expand(ctx, tc.query)
+			if err != nil {
+				t.Fatalf("Expand failed: %v", err)
+			}
+			if result != tc.expected {
+				t.Errorf("expected %q, got %q", tc.expected, result)
+			}
+		})
+	}
+}
+
+func TestExpandSQLite(t *testing.T) {
+	ctx := context.Background()
+
+	// Create an in-memory SQLite database using native API
+	conn, err := sqlite3.Open(":memory:")
+	if err != nil {
+		t.Fatalf("could not open SQLite: %v", err)
+	}
+	defer conn.Close()
+
+	// Create a test table
+	err = conn.Exec(`
+		CREATE TABLE authors (
+			id INTEGER PRIMARY KEY AUTOINCREMENT,
+			name TEXT NOT NULL,
+			bio TEXT
+		)
+	`)
+	if err != nil {
+		t.Fatalf("failed to create test table: %v", err)
+	}
+
+	// Create the parser which also implements format.Dialect
+	parser := sqlite.NewParser()
+
+	// Create the expander using native SQLite column getter
+	colGetter := &SQLiteColumnGetter{conn: conn}
+	exp := New(colGetter, parser, parser)
+
+	tests := []struct {
+		name     string
+		query    string
+		expected string
+	}{
+		{
+			name:     "simple select star",
+			query:    "SELECT * FROM authors",
+			expected: "SELECT id, name, bio FROM authors;",
+		},
+		{
+			name:     "select with no star",
+			query:    "SELECT id, name FROM authors",
+			expected: "SELECT id, name FROM authors", // No change, returns original
+		},
+		{
+			name:     "select star with where clause",
+			query:    "SELECT * FROM authors WHERE id = 1",
+			expected: "SELECT id, name, bio FROM authors WHERE id = 1;",
+		},
+		{
+			name:     "double star",
+			query:    "SELECT *, * FROM authors",
+			expected: "SELECT id, name, bio, id, name, bio FROM authors;",
+		},
+		{
+			name:     "table qualified star",
+			query:    "SELECT authors.* FROM authors",
+			expected: "SELECT authors.id, authors.name, authors.bio FROM authors;",
+		},
+		{
+			name:     "star in middle of columns",
+			query:    "SELECT id, *, name FROM authors",
+			expected: "SELECT id, id, name, bio, name FROM authors;",
+		},
+		{
+			name:     "count star not expanded",
+			query:    "SELECT COUNT(*) FROM authors",
+			expected: "SELECT COUNT(*) FROM authors", // No change - COUNT(*) should not be expanded
+		},
+		{
+			name:     "count star with other columns",
+			query:    "SELECT COUNT(*), name FROM authors GROUP BY name",
+			expected: "SELECT COUNT(*), name FROM authors GROUP BY name", // No change
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			result, err := exp.Expand(ctx, tc.query)
+			if err != nil {
+				t.Fatalf("Expand failed: %v", err)
+			}
+			if result != tc.expected {
+				t.Errorf("expected %q, got %q", tc.expected, result)
+			}
+		})
+	}
+}