Improve support for functions

- add IsArray field for function parameters
- support struct creation for table returns from funcs

Author: jrperritt

examples/booktest/postgresql/db_test.go

@@ -5,6 +5,8 @@ package booktest
 
 import (
 	"context"
+	"database/sql"
+	"reflect"
 	"testing"
 	"time"
 
@@ -14,6 +16,11 @@ import (
 func TestBooks(t *testing.T) {
 	db, cleanup := sqltest.PostgreSQL(t, []string{"schema.sql"})
 	defer cleanup()
+	defer func() {
+		if _, err := db.Exec("DROP SCHEMA IF EXISTS sqlc_book_schema CASCADE"); err != nil {
+			t.Fatal(err)
+		}
+	}()
 
 	ctx := context.Background()
 	dq := New(db)
@@ -24,6 +31,16 @@ func TestBooks(t *testing.T) {
 		t.Fatal(err)
 	}
 
+	// for https://github.com/kyleconroy/sqlc/issues/1322
+	bIdRaw, err := dq.CreateAuthorViaFunc(ctx, "Known Master")
+	if err != nil {
+		t.Fatal(err)
+	}
+	bIdRawType := reflect.TypeOf(bIdRaw)
+	if bIdRawType.String() != reflect.TypeOf(sql.NullInt32{}).String() {
+		t.Fatalf("unexpected type: want sql.NullInt32, got %s", bIdRawType)
+	}
+
 	// create transaction
 	tx, err := db.Begin()
 	if err != nil {
@@ -143,6 +160,19 @@ func TestBooks(t *testing.T) {
 		t.Logf("Book %d: '%s', Author: '%s', ISBN: '%s' Tags: '%v'\n", ab.BookID, ab.Title, ab.Name.String, ab.Isbn, ab.Tags)
 	}
 
+	booksByTagsViaFunc, err := dq.BooksByTagsViaFunc(ctx, []string{"cool", "other", "someother"})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(booksByTagsViaFunc) != len(res) {
+		t.Fatalf("got unequal number of books. wanted %d, got %d", len(res), len(booksByTagsViaFunc))
+	}
+	for i, bookByTag := range res {
+		if booksByTagsViaFunc[i].Isbn.String != bookByTag.Isbn {
+			t.Fatalf("unexpected book: wanted (%+v), got (%+v)", res, bookByTag)
+		}
+	}
+
 	// TODO: call say_hello(varchar)
 
 	// get book 4 and delete

examples/booktest/postgresql/models.go

@@ -23,7 +23,14 @@ SELECT
   tags
 FROM books
 LEFT JOIN authors ON books.author_id = authors.author_id
-WHERE tags && $1::varchar[];
+WHERE tags && $1::varchar[]
+ORDER BY name, title ASC;
+
+-- name: BooksByTagsViaFunc :many
+SELECT * FROM sqlc_book_schema.books_by_tags_func($1::VARCHAR[]);
+
+-- name: CreateAuthorViaFunc :one
+SELECT * FROM create_author_func($1);
 
 -- name: CreateAuthor :one
 INSERT INTO authors (name) VALUES ($1)

examples/booktest/postgresql/query.sql

@@ -30,3 +30,34 @@ END;
 $$ LANGUAGE plpgsql;
 
 CREATE INDEX books_title_lower_idx ON books(title);
+
+-- to test query return type / schema interaction
+CREATE SCHEMA sqlc_book_schema;
+
+CREATE OR REPLACE FUNCTION sqlc_book_schema.books_by_tags_func(_tags VARCHAR[])
+RETURNS TABLE(book_id INT, title TEXT, name TEXT, isbn TEXT, tags VARCHAR[]) AS $$
+BEGIN
+    RETURN QUERY
+    SELECT 
+        b.book_id,
+        b.title,
+        a.name,
+        b.isbn,
+        b.tags
+    FROM books b
+    LEFT JOIN authors a ON b.author_id = a.author_id
+    WHERE b.tags && _tags
+    ORDER BY name, title ASC
+    ;
+END;
+$$ LANGUAGE plpgsql;
+
+CREATE OR REPLACE FUNCTION create_author_func(_name TEXT)
+RETURNS TABLE(id INTEGER) AS $$
+BEGIN
+    INSERT INTO authors(name)
+    VALUES (_name)
+    RETURNING author_id into id;
+    RETURN NEXT;
+END;
+$$ LANGUAGE plpgsql;

examples/booktest/postgresql/query.sql.go

@@ -32,3 +32,15 @@ data class Book (
   val tags: List<String>
 )
 
+data class CreateAuthorFunc (
+  val id: Int?
+)
+
+data class SqlcBookSchemaBooksByTagsFunc (
+  val bookId: Int?,
+  val title: String?,
+  val name: String?,
+  val isbn: String?,
+  val tags: String?
+)
+

examples/booktest/postgresql/schema.sql

@@ -14,12 +14,18 @@ interface Queries {
   @Throws(SQLException::class)
   fun booksByTags(dollar1: List<String>): List<BooksByTagsRow>
 
+  @Throws(SQLException::class)
+  fun booksByTagsViaFunc(dollar1: List<String>): List<SqlcBookSchemaBooksByTagsFunc>
+  
   @Throws(SQLException::class)
   fun booksByTitleYear(title: String, year: Int): List<Book>
 
   @Throws(SQLException::class)
   fun createAuthor(name: String): Author?
 
+  @Throws(SQLException::class)
+  fun createAuthorViaFunc(name: String): Int??
+  
   @Throws(SQLException::class)
   fun createBook(
       authorId: Int,

examples/kotlin/src/main/kotlin/com/example/booktest/postgresql/Models.kt

@@ -20,6 +20,7 @@ SELECT
 FROM books
 LEFT JOIN authors ON books.author_id = authors.author_id
 WHERE tags && ?::varchar[]
+ORDER BY name, title ASC
 """
 
 data class BooksByTagsRow (
@@ -30,6 +31,10 @@ data class BooksByTagsRow (
   val tags: List<String>
 )
 
+const val booksByTagsViaFunc = """-- name: booksByTagsViaFunc :many
+SELECT book_id, title, name, isbn, tags FROM sqlc_book_schema.books_by_tags_func(?::VARCHAR[])
+"""
+
 const val booksByTitleYear = """-- name: booksByTitleYear :many
 SELECT book_id, author_id, isbn, book_type, title, year, available, tags FROM books
 WHERE title = ? AND year = ?
@@ -40,6 +45,10 @@ INSERT INTO authors (name) VALUES (?)
 RETURNING author_id, name
 """
 
+const val createAuthorViaFunc = """-- name: createAuthorViaFunc :one
+SELECT id FROM create_author_func(?)
+"""
+
 const val createBook = """-- name: createBook :one
 INSERT INTO books (
     author_id,
@@ -110,6 +119,26 @@ class QueriesImpl(private val conn: Connection) : Queries {
     }
   }
 
+  @Throws(SQLException::class)
+  override fun booksByTagsViaFunc(dollar1: List<String>): List<SqlcBookSchemaBooksByTagsFunc> {
+    return conn.prepareStatement(booksByTagsViaFunc).use { stmt ->
+      stmt.setArray(1, conn.createArrayOf("pg_catalog.varchar", dollar1.toTypedArray()))
+
+      val results = stmt.executeQuery()
+      val ret = mutableListOf<SqlcBookSchemaBooksByTagsFunc>()
+      while (results.next()) {
+          ret.add(SqlcBookSchemaBooksByTagsFunc(
+                results.getInt(1),
+                results.getString(2),
+                results.getString(3),
+                results.getString(4),
+                results.getString(5)
+            ))
+      }
+      ret
+    }
+  }
+
   @Throws(SQLException::class)
   override fun booksByTitleYear(title: String, year: Int): List<Book> {
     return conn.prepareStatement(booksByTitleYear).use { stmt ->
@@ -154,6 +183,23 @@ class QueriesImpl(private val conn: Connection) : Queries {
     }
   }
 
+  @Throws(SQLException::class)
+  override fun createAuthorViaFunc(name: String): Int?? {
+    return conn.prepareStatement(createAuthorViaFunc).use { stmt ->
+      stmt.setString(1, name)
+
+      val results = stmt.executeQuery()
+      if (!results.next()) {
+        return null
+      }
+      val ret = results.getInt(1)
+      if (results.next()) {
+          throw SQLException("expected one row in result set, but got many")
+      }
+      ret
+    }
+  }
+
   @Throws(SQLException::class)
   override fun createBook(
       authorId: Int,

examples/kotlin/src/main/kotlin/com/example/booktest/postgresql/Queries.kt

@@ -4,7 +4,7 @@
 import dataclasses
 import datetime
 import enum
-from typing import List
+from typing import List, Optional
 
 
 class BookType(str, enum.Enum):
@@ -28,3 +28,17 @@ class Book:
     year: int
     available: datetime.datetime
     tags: List[str]
+
+
+@dataclasses.dataclass()
+class CreateAuthorFunc:
+    id: Optional[int]
+
+
+@dataclasses.dataclass()
+class SqlcBookSchemaBooksByTagsFunc:
+    book_id: Optional[int]
+    title: Optional[str]
+    name: Optional[str]
+    isbn: Optional[str]
+    tags: Optional[str]