feat: add support for user defined TVFs

google-cloud-bigquery/clirr-ignored-differences.xml

@@ -5,6 +5,6 @@
   <difference>
     <differenceType>7013</differenceType>
     <className>com/google/cloud/bigquery/RoutineInfo$Builder</className>
-    <method>com.google.cloud.bigquery.RoutineInfo$Builder setDeterminismLevel(java.lang.String)</method>
+    <method>com.google.cloud.bigquery.RoutineInfo$Builder setReturnTableType(com.google.cloud.bigquery.StandardSQLTableType)</method>
   </difference>
 </differences>

google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/Routine.java

@@ -111,6 +111,12 @@ public Builder setReturnType(StandardSQLDataType returnType) {
       return this;
     }
 
+    @Override
+    public Builder setReturnTableType(StandardSQLTableType returnTableType) {
+      infoBuilder.setReturnTableType(returnTableType);
+      return this;
+    }
+
     @Override
     public Builder setImportedLibraries(List<String> libraries) {
       infoBuilder.setImportedLibraries(libraries);

google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/RoutineInfo.java

@@ -67,6 +67,7 @@ public Routine apply(RoutineInfo routineInfo) {
   private final String language;
   private final List<RoutineArgument> argumentList;
   private final StandardSQLDataType returnType;
+  private final StandardSQLTableType returnTableType;
   private final List<String> importedLibrariesList;
   private final String body;
 
@@ -113,6 +114,9 @@ public abstract static class Builder {
      */
     public abstract Builder setReturnType(StandardSQLDataType returnType);
 
+    /** Optional. Set only if Routine is a "TABLE_VALUED_FUNCTION". */
+    public abstract Builder setReturnTableType(StandardSQLTableType returnTableType);
+
     /**
      * Optional. If language = "JAVASCRIPT", this field stores the path of the imported JAVASCRIPT
      * libraries as a list of gs:// URLs.
@@ -159,6 +163,7 @@ static class BuilderImpl extends Builder {
     private String language;
     private List<RoutineArgument> argumentList;
     private StandardSQLDataType returnType;
+    private StandardSQLTableType returnTableType;
     private List<String> importedLibrariesList;
     private String body;
 
@@ -175,6 +180,7 @@ static class BuilderImpl extends Builder {
       this.language = routineInfo.language;
       this.argumentList = routineInfo.argumentList;
       this.returnType = routineInfo.returnType;
+      this.returnTableType = routineInfo.returnTableType;
       this.importedLibrariesList = routineInfo.importedLibrariesList;
       this.body = routineInfo.body;
     }
@@ -195,6 +201,9 @@ static class BuilderImpl extends Builder {
       if (routinePb.getReturnType() != null) {
         this.returnType = StandardSQLDataType.fromPb(routinePb.getReturnType());
       }
+      if (routinePb.getReturnTableType() != null) {
+        this.returnTableType = StandardSQLTableType.fromPb(routinePb.getReturnTableType());
+      }
       if (routinePb.getImportedLibraries() == null) {
         this.importedLibrariesList = Collections.emptyList();
       } else {
@@ -263,6 +272,12 @@ public Builder setReturnType(StandardSQLDataType returnType) {
       return this;
     }
 
+    @Override
+    public Builder setReturnTableType(StandardSQLTableType returnTableType) {
+      this.returnTableType = returnTableType;
+      return this;
+    }
+
     @Override
     public Builder setImportedLibraries(List<String> importedLibrariesList) {
       this.importedLibrariesList = importedLibrariesList;
@@ -292,6 +307,7 @@ public RoutineInfo build() {
     this.language = builder.language;
     this.argumentList = builder.argumentList;
     this.returnType = builder.returnType;
+    this.returnTableType = builder.returnTableType;
     this.importedLibrariesList = builder.importedLibrariesList;
     this.body = builder.body;
   }
@@ -350,6 +366,11 @@ public StandardSQLDataType getReturnType() {
     return returnType;
   }
 
+  /** If specified, returns the table type returned from the routine. */
+  public StandardSQLTableType getReturnTableType() {
+    return returnTableType;
+  }
+
   /**
    * Returns the list of imported libraries for the routine. Only relevant for routines implemented
    * using the JAVASCRIPT language.
@@ -381,6 +402,7 @@ public String toString() {
         .add("language", language)
         .add("arguments", argumentList)
         .add("returnType", returnType)
+        .add("returnTableType", returnTableType)
         .add("importedLibrariesList", importedLibrariesList)
         .add("body", body)
         .toString();
@@ -399,6 +421,7 @@ public int hashCode() {
         language,
         argumentList,
         returnType,
+        returnTableType,
         importedLibrariesList,
         body);
   }
@@ -448,6 +471,9 @@ Routine toPb() {
     if (getReturnType() != null) {
       routinePb.setReturnType(getReturnType().toPb());
     }
+    if (getReturnTableType() != null) {
+      routinePb.setReturnTableType(getReturnTableType().toPb());
+    }
     return routinePb;
   }
 

google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/StandardSQLTableType.java

@@ -0,0 +1,70 @@
+/*
+ * Copyright 2021 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.google.cloud.bigquery;
+
+import com.google.api.services.bigquery.model.StandardSqlTableType;
+import com.google.auto.value.AutoValue;
+import com.google.common.collect.Lists;
+import java.io.Serializable;
+import java.util.List;
+
+/** Represents Standard SQL table type information. */
+@AutoValue
+public abstract class StandardSQLTableType implements Serializable {
+
+  @AutoValue.Builder
+  public abstract static class Builder {
+
+    /** Sets the columns in this table type. */
+    public abstract Builder setColumns(List<StandardSQLField> columns);
+
+    /** Creates a {@code StandardSQLTableType} object. */
+    public abstract StandardSQLTableType build();
+  }
+
+  /** Returns the columns in this table type. */
+  public abstract List<StandardSQLField> getColumns();
+
+  public abstract Builder toBuilder();
+
+  /** Returns a builder for a {@code StandardSQLTableType} object. */
+  public static Builder newBuilder() {
+    return new AutoValue_StandardSQLTableType.Builder();
+  }
+
+  /** Returns a builder for a {@code StandardSQLTableType} object with the specified columns. */
+  public static StandardSQLTableType.Builder newBuilder(List<StandardSQLField> columns) {
+    return newBuilder().setColumns(columns);
+  }
+
+  static StandardSQLTableType fromPb(
+      com.google.api.services.bigquery.model.StandardSqlTableType tableTypePb) {
+    StandardSQLTableType.Builder builder = newBuilder();
+    if (tableTypePb.getColumns() != null) {
+      builder.setColumns(
+          Lists.transform(tableTypePb.getColumns(), StandardSQLField.FROM_PB_FUNCTION));
+    }
+    return builder.build();
+  }
+
+  StandardSqlTableType toPb() {
+    StandardSqlTableType tableType = new StandardSqlTableType();
+    if (getColumns() != null) {
+      tableType.setColumns(Lists.transform(getColumns(), StandardSQLField.TO_PB_FUNCTION));
+    }
+    return tableType;
+  }
+}

google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/RoutineTest.java

@@ -38,9 +38,11 @@
 public class RoutineTest {
 
   private static final RoutineId ROUTINE_ID = RoutineId.of("dataset", "routine");
+  private static final RoutineId ROUTINE_ID_TVF = RoutineId.of("dataset", "tvf_routine");
   private static final String DETERMINISM_LEVEL = "DETERMINISTIC";
   private static final String ETAG = "etag";
   private static final String ROUTINE_TYPE = "SCALAR_FUNCTION";
+  private static final String ROUTINE_TYPE_TVF = "TABLE_VALUED_FUNCTION";
   private static final Long CREATION_TIME = 10L;
   private static final Long LAST_MODIFIED_TIME = 20L;
   private static final String LANGUAGE = "SQL";
@@ -56,6 +58,18 @@ public class RoutineTest {
   private static final StandardSQLDataType RETURN_TYPE =
       StandardSQLDataType.newBuilder("FLOAT64").build();
 
+  private static final StandardSQLField COLUMN_1 =
+      StandardSQLField.newBuilder("COLUMN_1", StandardSQLDataType.newBuilder("STRING").build())
+          .build();
+  private static final StandardSQLField COLUMN_2 =
+      StandardSQLField.newBuilder("COLUMN_2", StandardSQLDataType.newBuilder("FLOAT64").build())
+          .build();
+
+  private static final List<StandardSQLField> COLUMN_LIST = ImmutableList.of(COLUMN_1, COLUMN_2);
+
+  private static final StandardSQLTableType RETURN_TABLE_TYPE =
+      StandardSQLTableType.newBuilder(COLUMN_LIST).build();
+
   private static final List<String> IMPORTED_LIBRARIES =
       ImmutableList.of("gs://foo", "gs://bar", "gs://baz");
 
@@ -75,11 +89,19 @@ public class RoutineTest {
           .setBody(BODY)
           .build();
 
+  private static final RoutineInfo ROUTINE_INFO_TVF =
+      RoutineInfo.newBuilder(ROUTINE_ID_TVF)
+          .setBody(BODY)
+          .setRoutineType(ROUTINE_TYPE_TVF)
+          .setReturnTableType(RETURN_TABLE_TYPE)
+          .build();
+
   @Rule public MockitoRule rule;
 
   private BigQuery bigquery;
   private BigQueryOptions mockOptions;
   private Routine expectedRoutine;
+  private Routine expectedRoutineTvf;
   private Routine routine;
 
   @Before
@@ -88,6 +110,7 @@ public void setUp() {
     mockOptions = mock(BigQueryOptions.class);
     when(bigquery.getOptions()).thenReturn(mockOptions);
     expectedRoutine = new Routine(bigquery, new RoutineInfo.BuilderImpl(ROUTINE_INFO));
+    expectedRoutineTvf = new Routine(bigquery, new RoutineInfo.BuilderImpl(ROUTINE_INFO_TVF));
     routine = new Routine(bigquery, new RoutineInfo.BuilderImpl(ROUTINE_INFO));
   }
 
@@ -114,6 +137,7 @@ public void testBuilder() {
   @Test
   public void testToBuilder() {
     compareRoutineInfo(expectedRoutine, expectedRoutine.toBuilder().build());
+    compareRoutineInfo(expectedRoutineTvf, expectedRoutineTvf.toBuilder().build());
   }
 
   @Test
@@ -200,6 +224,7 @@ public void compareRoutineInfo(RoutineInfo expected, RoutineInfo value) {
     assertEquals(expected.getLanguage(), value.getLanguage());
     assertEquals(expected.getArguments(), value.getArguments());
     assertEquals(expected.getReturnType(), value.getReturnType());
+    assertEquals(expected.getReturnTableType(), value.getReturnTableType());
     assertEquals(expected.getImportedLibraries(), value.getImportedLibraries());
     assertEquals(expected.getBody(), value.getBody());
     assertEquals(expected.hashCode(), value.hashCode());

google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/StandardSQLTableTypeTest.java

@@ -0,0 +1,60 @@
+/*
+ * Copyright 2021 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.cloud.bigquery;
+
+import static org.junit.Assert.*;
+
+import com.google.common.collect.ImmutableList;
+import java.util.List;
+import org.junit.Test;
+
+public class StandardSQLTableTypeTest {
+
+  private static final StandardSQLField COLUMN_1 =
+      StandardSQLField.newBuilder("COLUMN_1", StandardSQLDataType.newBuilder("STRING").build())
+          .build();
+  private static final StandardSQLField COLUMN_2 =
+      StandardSQLField.newBuilder("COLUMN_2", StandardSQLDataType.newBuilder("FLOAT64").build())
+          .build();
+
+  private static final List<StandardSQLField> COLUMN_LIST = ImmutableList.of(COLUMN_1, COLUMN_2);
+  private static final StandardSQLTableType TABLE_TYPE =
+      StandardSQLTableType.newBuilder(COLUMN_LIST).build();
+
+  @Test
+  public void testToBuilder() {
+    compareStandardSQLTableType(TABLE_TYPE, TABLE_TYPE.toBuilder().build());
+  }
+
+  @Test
+  public void testBuilder() {
+    assertEquals(COLUMN_1, TABLE_TYPE.getColumns().get(0));
+    assertEquals(COLUMN_2, TABLE_TYPE.getColumns().get(1));
+  }
+
+  @Test
+  public void testToAndFromPb() {
+    compareStandardSQLTableType(TABLE_TYPE, StandardSQLTableType.fromPb(TABLE_TYPE.toPb()));
+  }
+
+  private void compareStandardSQLTableType(
+      StandardSQLTableType expected, StandardSQLTableType value) {
+    assertEquals(expected, value);
+    assertEquals(expected.getColumns(), value.getColumns());
+    assertEquals(expected.hashCode(), value.hashCode());
+  }
+}

google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java

@@ -87,6 +87,8 @@
 import com.google.cloud.bigquery.RoutineInfo;
 import com.google.cloud.bigquery.Schema;
 import com.google.cloud.bigquery.StandardSQLDataType;
+import com.google.cloud.bigquery.StandardSQLField;
+import com.google.cloud.bigquery.StandardSQLTableType;
 import com.google.cloud.bigquery.StandardTableDefinition;
 import com.google.cloud.bigquery.Table;
 import com.google.cloud.bigquery.TableDataWriteChannel;
@@ -1676,6 +1678,34 @@ public void testRoutineAPICreationJavascriptUDF() {
     assertEquals(routine.getReturnType(), StandardSQLDataType.newBuilder("STRING").build());
   }
 
+  @Test
+  public void testRoutineAPICreationTVF() {
+    String routineName = RemoteBigQueryHelper.generateRoutineName();
+    RoutineId routineId = RoutineId.of(ROUTINE_DATASET, routineName);
+    List<StandardSQLField> columns =
+        ImmutableList.of(
+            StandardSQLField.newBuilder("x", StandardSQLDataType.newBuilder("INT64").build())
+                .build());
+    StandardSQLTableType returnTableType = StandardSQLTableType.newBuilder(columns).build();
+    RoutineInfo routineInfo =
+        RoutineInfo.newBuilder(routineId)
+            .setRoutineType("TABLE_VALUED_FUNCTION")
+            .setLanguage("SQL")
+            .setArguments(
+                ImmutableList.of(
+                    RoutineArgument.newBuilder()
+                        .setName("filter")
+                        .setDataType(StandardSQLDataType.newBuilder("INT64").build())
+                        .build()))
+            .setReturnTableType(returnTableType)
+            .setBody("SELECT x FROM UNNEST([1,2,3]) x WHERE x = filter")
+            .build();
+    Routine routine = bigquery.create(routineInfo);
+    assertNotNull(routine);
+    assertEquals(routine.getRoutineType(), "TABLE_VALUED_FUNCTION");
+    assertEquals(routine.getReturnTableType(), returnTableType);
+  }
+
   @Test
   public void testAuthorizeRoutine() {
     String routineName = RemoteBigQueryHelper.generateRoutineName();