[FlightSQL] Add missing method for creating bitmask from GetSqlInfo o…

…ption enum (apache#148) * Add util method for creating bitmask from multiple protobuf enums for FlightSql GetSqlInfo enum * Add test cases for utility method for creating bitmask from protobuf options * Make changes regarding to reviews Co-authored-by: Rafael Telles <rafael@telles.dev>
rafael-telles · Oct 28, 2021 · 311bf36 · 311bf36
1 parent b315cae
commit 311bf36
Show file tree

Hide file tree

Showing 4 changed files with 161 additions and 53 deletions.
diff --git a/...flight/flight-sql/src/main/java/org/apache/arrow/flight/sql/util/SqlInfoOptionsUtils.java b/...flight/flight-sql/src/main/java/org/apache/arrow/flight/sql/util/SqlInfoOptionsUtils.java
@@ -17,11 +17,12 @@
 
 package org.apache.arrow.flight.sql.util;
 
+import com.google.protobuf.ProtocolMessageEnum;
+import java.util.Arrays;
+import java.util.Collection;
 import org.apache.arrow.flight.sql.FlightSqlClient;
 import org.apache.arrow.flight.sql.impl.FlightSql.SqlInfo;
 
-import com.google.protobuf.ProtocolMessageEnum;
-
 /**
  * Utility class for {@link SqlInfo} and {@link FlightSqlClient#getSqlInfo} option parsing.
  */
@@ -41,4 +42,28 @@ private SqlInfoOptionsUtils() {
   public static boolean doesBitmaskTranslateToEnum(final ProtocolMessageEnum enumInstance, final long bitmask) {
     return ((bitmask >> enumInstance.getNumber()) & 1) == 1;
   }
+
+  /**
+   * Creates a bitmask that translates to the specified {@code enums}.
+   *
+   * @param enums the {@link ProtocolMessageEnum} instances to represent as bitmask.
+   * @return the bitmask.
+   */
+  public static long createBitmaskFromEnums(final ProtocolMessageEnum... enums) {
+    return createBitmaskFromEnums(Arrays.asList(enums));
+  }
+
+  /**
+   * Creates a bitmask that translates to the specified {@code enums}.
+   *
+   * @param enums the {@link ProtocolMessageEnum} instances to represent as bitmask.
+   * @return the bitmask.
+   */
+  public static long createBitmaskFromEnums(final Collection<ProtocolMessageEnum> enums) {
+    return enums.stream()
+        .mapToInt(ProtocolMessageEnum::getNumber)
+        .map(bitIndexToSet -> 1 << bitIndexToSet)
+        .reduce((firstBitmask, secondBitmask) -> firstBitmask | secondBitmask)
+        .orElse(0);
+  }
 }
diff --git a/java/flight/flight-sql/src/test/java/org/apache/arrow/flight/sql/util/AdhocTestOption.java b/java/flight/flight-sql/src/test/java/org/apache/arrow/flight/sql/util/AdhocTestOption.java
@@ -0,0 +1,45 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.arrow.flight.sql.util;
+
+import com.google.protobuf.Descriptors.EnumDescriptor;
+import com.google.protobuf.Descriptors.EnumValueDescriptor;
+import com.google.protobuf.ProtocolMessageEnum;
+
+enum AdhocTestOption implements ProtocolMessageEnum {
+  OPTION_A, OPTION_B, OPTION_C;
+
+  @Override
+  public int getNumber() {
+    return ordinal();
+  }
+
+  @Override
+  public EnumValueDescriptor getValueDescriptor() {
+    throw getUnsupportedException();
+  }
+
+  @Override
+  public EnumDescriptor getDescriptorForType() {
+    throw getUnsupportedException();
+  }
+
+  private UnsupportedOperationException getUnsupportedException() {
+    return new UnsupportedOperationException("Unimplemented method is irrelevant for the scope of this test.");
+  }
+}
diff --git a/...rc/test/java/org/apache/arrow/flight/sql/util/SqlInfoOptionsUtilsBitmaskCreationTest.java b/...rc/test/java/org/apache/arrow/flight/sql/util/SqlInfoOptionsUtilsBitmaskCreationTest.java
@@ -0,0 +1,66 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.arrow.flight.sql.util;
+
+import static java.util.Arrays.asList;
+import static org.apache.arrow.flight.sql.util.AdhocTestOption.OPTION_A;
+import static org.apache.arrow.flight.sql.util.AdhocTestOption.OPTION_B;
+import static org.apache.arrow.flight.sql.util.AdhocTestOption.OPTION_C;
+import static org.apache.arrow.flight.sql.util.SqlInfoOptionsUtils.createBitmaskFromEnums;
+import static org.hamcrest.CoreMatchers.is;
+
+import java.util.List;
+
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.ErrorCollector;
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
+import org.junit.runners.Parameterized.Parameter;
+import org.junit.runners.Parameterized.Parameters;
+
+@RunWith(Parameterized.class)
+public final class SqlInfoOptionsUtilsBitmaskCreationTest {
+
+  @Parameter
+  public AdhocTestOption[] adhocTestOptions;
+  @Parameter(value = 1)
+  public long expectedBitmask;
+  @Rule
+  public final ErrorCollector collector = new ErrorCollector();
+
+  @Parameters
+  public static List<Object[]> provideParameters() {
+    return asList(
+        new Object[][]{
+            {new AdhocTestOption[0], 0L},
+            {new AdhocTestOption[]{OPTION_A}, 1L},
+            {new AdhocTestOption[]{OPTION_B}, 0b10L},
+            {new AdhocTestOption[]{OPTION_A, OPTION_B}, 0b11L},
+            {new AdhocTestOption[]{OPTION_C}, 0b100L},
+            {new AdhocTestOption[]{OPTION_A, OPTION_C}, 0b101L},
+            {new AdhocTestOption[]{OPTION_B, OPTION_C}, 0b110L},
+            {AdhocTestOption.values(), 0b111L},
+        });
+  }
+
+  @Test
+  public void testShouldBuildBitmaskFromEnums() {
+    collector.checkThat(createBitmaskFromEnums(adhocTestOptions), is(expectedBitmask));
+  }
+}
diff --git a/...ght/sql/util/SqlInfoOptionsUtilsTest.java → ...qlInfoOptionsUtilsBitmaskParsingTest.java b/...ght/sql/util/SqlInfoOptionsUtilsTest.java → ...qlInfoOptionsUtilsBitmaskParsingTest.java
@@ -17,16 +17,19 @@
 
 package org.apache.arrow.flight.sql.util;
 
+import static java.util.Arrays.asList;
+import static java.util.Arrays.stream;
 import static java.util.stream.Collectors.toCollection;
+import static org.apache.arrow.flight.sql.util.AdhocTestOption.OPTION_A;
+import static org.apache.arrow.flight.sql.util.AdhocTestOption.OPTION_B;
+import static org.apache.arrow.flight.sql.util.AdhocTestOption.OPTION_C;
 import static org.apache.arrow.flight.sql.util.SqlInfoOptionsUtils.doesBitmaskTranslateToEnum;
 import static org.hamcrest.CoreMatchers.is;
 
-import java.util.Arrays;
 import java.util.EnumSet;
 import java.util.List;
 import java.util.Set;
 
-import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.rules.ErrorCollector;
@@ -35,68 +38,37 @@
 import org.junit.runners.Parameterized.Parameter;
 import org.junit.runners.Parameterized.Parameters;
 
-import com.google.protobuf.Descriptors.EnumDescriptor;
-import com.google.protobuf.Descriptors.EnumValueDescriptor;
-import com.google.protobuf.ProtocolMessageEnum;
-
 @RunWith(Parameterized.class)
-public final class SqlInfoOptionsUtilsTest {
+public final class SqlInfoOptionsUtilsBitmaskParsingTest {
 
   @Parameter
   public long bitmask;
   @Parameter(value = 1)
-  public Set<TestOption> messageEnums;
-  public Set<TestOption> expectedOutcome;
+  public Set<AdhocTestOption> expectedOptions;
   @Rule
   public final ErrorCollector collector = new ErrorCollector();
 
-  @Before
-  public void setUp() {
-    expectedOutcome =
-        Arrays.stream(TestOption.values())
-            .filter(enumInstance -> doesBitmaskTranslateToEnum(enumInstance, bitmask))
-            .collect(toCollection(() -> EnumSet.noneOf(TestOption.class)));
-  }
-
   @Parameters
   public static List<Object[]> provideParameters() {
-    return Arrays.asList(new Object[][]{
-        {0, EnumSet.noneOf(TestOption.class)},
-        {1, EnumSet.of(TestOption.OPTION_A)},
-        {0b10, EnumSet.of(TestOption.OPTION_B)},
-        {0b11, EnumSet.of(TestOption.OPTION_A, TestOption.OPTION_B)},
-        {0b100, EnumSet.of(TestOption.OPTION_C)},
-        {0b101, EnumSet.of(TestOption.OPTION_A, TestOption.OPTION_C)},
-        {0b110, EnumSet.of(TestOption.OPTION_B, TestOption.OPTION_C)},
-        {0b111, EnumSet.allOf(TestOption.class)},
-    });
+    return asList(
+        new Object[][]{
+            {0L, EnumSet.noneOf(AdhocTestOption.class)},
+            {1L, EnumSet.of(OPTION_A)},
+            {0b10L, EnumSet.of(OPTION_B)},
+            {0b11L, EnumSet.of(OPTION_A, OPTION_B)},
+            {0b100L, EnumSet.of(OPTION_C)},
+            {0b101L, EnumSet.of(OPTION_A, OPTION_C)},
+            {0b110L, EnumSet.of(OPTION_B, OPTION_C)},
+            {0b111L, EnumSet.allOf(AdhocTestOption.class)},
+        });
   }
 
   @Test
   public void testShouldFilterOutEnumsBasedOnBitmask() {
-    collector.checkThat(messageEnums, is(expectedOutcome));
-  }
-
-  private enum TestOption implements ProtocolMessageEnum {
-    OPTION_A, OPTION_B, OPTION_C;
-
-    @Override
-    public int getNumber() {
-      return ordinal();
-    }
-
-    @Override
-    public EnumValueDescriptor getValueDescriptor() {
-      throw getUnsupportedException();
-    }
-
-    @Override
-    public EnumDescriptor getDescriptorForType() {
-      throw getUnsupportedException();
-    }
-
-    private UnsupportedOperationException getUnsupportedException() {
-      return new UnsupportedOperationException("Unimplemented method is irrelevant for the scope of this test.");
-    }
+    final Set<AdhocTestOption> actualOptions =
+        stream(AdhocTestOption.values())
+            .filter(enumInstance -> doesBitmaskTranslateToEnum(enumInstance, bitmask))
+            .collect(toCollection(() -> EnumSet.noneOf(AdhocTestOption.class)));
+    collector.checkThat(actualOptions, is(expectedOptions));
   }
 }