feat: add acl for routine (#875)

Author: Praful Makani

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

@@ -104,7 +104,8 @@ public enum Type {
       GROUP,
       USER,
       VIEW,
-      IAM_MEMBER
+      IAM_MEMBER,
+      ROUTINE
     }
 
     Entity(Type type) {
@@ -136,6 +137,9 @@ static Entity fromPb(Access access) {
       if (access.getIamMember() != null) {
         return new IamMember(access.getIamMember());
       }
+      if (access.getRoutine() != null) {
+        return new Routine(RoutineId.fromPb(access.getRoutine()));
+      }
       // Unreachable
       throw new BigQueryException(
           BigQueryException.UNKNOWN_CODE, "Unrecognized access configuration");
@@ -387,6 +391,58 @@ Access toPb() {
     }
   }
 
+  /**
+   * Class for a BigQuery Routine entity. Objects of this class represent a routine from a different
+   * dataset to grant access to. Queries executed against that routine will have read access to
+   * views/tables/routines in this dataset. Only UDF is supported for now. The role field is not
+   * required when this field is set. If that routine is updated by any user, access to the routine
+   * needs to be granted again via an update operation.
+   */
+  public static final class Routine extends Entity {
+
+    private static final long serialVersionUID = -8392885851733136262L;
+
+    private final RoutineId id;
+
+    /** Creates a Routine entity given the routine's id. */
+    public Routine(RoutineId id) {
+      super(Type.ROUTINE);
+      this.id = id;
+    }
+
+    /** Returns routine's identity. */
+    public RoutineId getId() {
+      return id;
+    }
+
+    @Override
+    public boolean equals(Object obj) {
+      if (this == obj) {
+        return true;
+      }
+      if (obj == null || getClass() != obj.getClass()) {
+        return false;
+      }
+      Routine routine = (Routine) obj;
+      return Objects.equals(getType(), routine.getType()) && Objects.equals(id, routine.id);
+    }
+
+    @Override
+    public int hashCode() {
+      return Objects.hash(getType(), id);
+    }
+
+    @Override
+    public String toString() {
+      return toPb().toString();
+    }
+
+    @Override
+    Access toPb() {
+      return new Access().setRoutine(id.toPb());
+    }
+  }
+
   /**
    * Class for a BigQuery IamMember entity. Objects of this class represent a iamMember to grant
    * access to given the IAM Policy.
@@ -465,6 +521,11 @@ public static Acl of(View view) {
     return new Acl(view, null);
   }
 
+  /** Returns an Acl object for a routine entity. */
+  public static Acl of(Routine routine) {
+    return new Acl(routine, null);
+  }
+
   @Override
   public int hashCode() {
     return Objects.hash(entity, role);

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

@@ -21,6 +21,7 @@
 
 import com.google.api.client.util.Data;
 import com.google.api.services.bigquery.model.Dataset;
+import com.google.api.services.bigquery.model.RoutineReference;
 import com.google.api.services.bigquery.model.TableReference;
 import com.google.common.base.Function;
 import com.google.common.base.MoreObjects;
@@ -509,6 +510,13 @@ DatasetInfo setProjectId(String projectId) {
             viewReferencePb.setProjectId(projectId);
           }
           acls.add(Acl.of(new Acl.View(TableId.fromPb(viewReferencePb))));
+        } else if (acl.getEntity().getType() == Acl.Entity.Type.ROUTINE) {
+          Dataset.Access accessPb = acl.toPb();
+          RoutineReference routineReferencePb = accessPb.getRoutine();
+          if (routineReferencePb.getProjectId() == null) {
+            routineReferencePb.setProjectId(projectId);
+          }
+          acls.add(Acl.of(new Acl.Routine(RoutineId.fromPb(routineReferencePb))));
         } else {
           acls.add(acl);
         }

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

@@ -88,6 +88,16 @@ public void testViewEntity() {
     assertEquals(entity, Entity.fromPb(pb));
   }
 
+  @Test
+  public void testRoutineEntity() {
+    RoutineId routineId = RoutineId.of("project", "dataset", "routine");
+    Acl.Routine entity = new Acl.Routine(routineId);
+    assertEquals(routineId, entity.getId());
+    assertEquals(Type.ROUTINE, entity.getType());
+    Dataset.Access pb = entity.toPb();
+    assertEquals(entity, Entity.fromPb(pb));
+  }
+
   @Test
   public void testIamMemberEntity() {
     IamMember entity = new IamMember("member1");
@@ -107,5 +117,9 @@ public void testOf() {
     acl = Acl.of(view);
     assertEquals(view, acl.getEntity());
     assertEquals(null, acl.getRole());
+    Acl.Routine routine = new Acl.Routine(RoutineId.of("project", "dataset", "routine"));
+    acl = Acl.of(routine);
+    assertEquals(routine, acl.getEntity());
+    assertEquals(null, acl.getRole());
   }
 }

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

@@ -31,11 +31,13 @@ public class DatasetInfoTest {
   private static final List<Acl> ACCESS_RULES =
       ImmutableList.of(
           Acl.of(Acl.Group.ofAllAuthenticatedUsers(), Acl.Role.READER),
-          Acl.of(new Acl.View(TableId.of("dataset", "table"))));
+          Acl.of(new Acl.View(TableId.of("dataset", "table"))),
+          Acl.of(new Acl.Routine(RoutineId.of("dataset", "routine"))));
   private static final List<Acl> ACCESS_RULES_COMPLETE =
       ImmutableList.of(
           Acl.of(Acl.Group.ofAllAuthenticatedUsers(), Acl.Role.READER),
-          Acl.of(new Acl.View(TableId.of("project", "dataset", "table"))));
+          Acl.of(new Acl.View(TableId.of("project", "dataset", "table"))),
+          Acl.of(new Acl.Routine(RoutineId.of("project", "dataset", "routine"))));
   private static final List<Acl> ACCESS_RULES_IAM_MEMBER =
       ImmutableList.of(Acl.of(new Acl.IamMember("allUsers"), Acl.Role.READER));
   private static final Map<String, String> LABELS =

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

@@ -49,7 +49,8 @@ public class DatasetTest {
   private static final List<Acl> ACCESS_RULES =
       ImmutableList.of(
           Acl.of(Acl.Group.ofAllAuthenticatedUsers(), Acl.Role.READER),
-          Acl.of(new Acl.View(TableId.of("dataset", "table"))));
+          Acl.of(new Acl.View(TableId.of("dataset", "table"))),
+          Acl.of(new Acl.Routine(RoutineId.of("dataset", "routine"))));
   private static final Map<String, String> LABELS =
       ImmutableMap.of(
           "example-label1", "example-value1",

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

@@ -36,8 +36,10 @@ public class SerializationTest extends BaseSerializationTest {
   private static final Acl USER_ACCESS = Acl.of(new Acl.User("user"), Acl.Role.OWNER);
   private static final Acl VIEW_ACCESS =
       Acl.of(new Acl.View(TableId.of("project", "dataset", "table")), Acl.Role.WRITER);
+  private static final Acl ROUTINE_ACCESS =
+      Acl.of(new Acl.Routine(RoutineId.of("project", "dataset", "routine")), Acl.Role.WRITER);
   private static final List<Acl> ACCESS_RULES =
-      ImmutableList.of(DOMAIN_ACCESS, GROUP_ACCESS, VIEW_ACCESS, USER_ACCESS);
+      ImmutableList.of(DOMAIN_ACCESS, GROUP_ACCESS, VIEW_ACCESS, ROUTINE_ACCESS, USER_ACCESS);
   private static final Long CREATION_TIME = System.currentTimeMillis() - 10;
   private static final Long DEFAULT_TABLE_EXPIRATION = 100L;
   private static final String DESCRIPTION = "Description";
@@ -225,6 +227,7 @@ protected Serializable[] serializableObjects() {
       GROUP_ACCESS,
       USER_ACCESS,
       VIEW_ACCESS,
+      ROUTINE_ACCESS,
       DATASET_ID,
       DATASET_INFO,
       TABLE_ID,

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

@@ -112,6 +112,7 @@
 import java.nio.ByteBuffer;
 import java.nio.charset.StandardCharsets;
 import java.nio.file.FileSystems;
+import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
@@ -1491,6 +1492,32 @@ public void testRoutineAPICreation() {
     assertEquals(routine.getRoutineType(), "SCALAR_FUNCTION");
   }
 
+  @Test
+  public void testAuthorizeRoutine() {
+    String routineName = RemoteBigQueryHelper.generateRoutineName();
+    RoutineId routineId = RoutineId.of(PROJECT_ID, ROUTINE_DATASET, routineName);
+    RoutineInfo routineInfo =
+        RoutineInfo.newBuilder(routineId)
+            .setRoutineType("SCALAR_FUNCTION")
+            .setBody("x * 3")
+            .setLanguage("SQL")
+            .setArguments(
+                ImmutableList.of(
+                    RoutineArgument.newBuilder()
+                        .setName("x")
+                        .setDataType(StandardSQLDataType.newBuilder("INT64").build())
+                        .build()))
+            .build();
+    Routine routine = bigquery.create(routineInfo);
+    assertNotNull(routine);
+    assertEquals(routine.getRoutineType(), "SCALAR_FUNCTION");
+    Dataset routineDataset = bigquery.getDataset(ROUTINE_DATASET);
+    List<Acl> routineAcl = new ArrayList<>(routineDataset.getAcl());
+    routineAcl.add(Acl.of(new Acl.Routine(routineId)));
+    routineDataset = routineDataset.toBuilder().setAcl(routineAcl).build().update();
+    assertEquals(routineAcl, routineDataset.getAcl());
+  }
+
   @Test
   public void testSingleStatementsQueryException() throws InterruptedException {
     String invalidQuery =