Small Tables PR for refactoring listEntities common logic (#34185)

* Small Refactoring for listEntities * Update TableUtils.java Editing whitespaces to address failing ci
Azure · Jul 6, 2023 · 2f40427 · 2f40427
1 parent 3471fc5
commit 2f40427
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 23 deletions.
diff --git a/sdk/tables/azure-data-tables/src/main/java/com/azure/data/tables/TableAsyncClient.java b/sdk/tables/azure-data-tables/src/main/java/com/azure/data/tables/TableAsyncClient.java
@@ -98,7 +98,6 @@
  */
 @ServiceClient(builder = TableClientBuilder.class, isAsync = true)
 public final class TableAsyncClient {
-    private static final String DELIMITER_CONTINUATION_TOKEN = ";";
     private final ClientLogger logger = new ClientLogger(TableAsyncClient.class);
     private final String tableName;
     private final AzureTableImpl tablesImplementation;
@@ -950,17 +949,12 @@ private <T extends TableEntity> Mono<PagedResponse<T>> listEntitiesNextPage(Stri
             return Mono.empty();
         }
 
-        String[] split = token.split(DELIMITER_CONTINUATION_TOKEN, 2);
-
-        if (split.length == 0) {
-            return monoError(logger, new RuntimeException(
-                "Split done incorrectly, must have partition key: " + token));
+        try {
+            String[] split = TableUtils.getKeysFromToken(token);
+            return listEntities(split[0], split[1], context, options, resultType);
+        } catch (RuntimeException ex) {
+            return monoError(logger, ex);
         }
-
-        String nextPartitionKey = split[0];
-        String nextRowKey = split.length > 1 ? split[1] : null;
-
-        return listEntities(nextPartitionKey, nextRowKey, context, options, resultType);
     }
 
     private <T extends TableEntity> Mono<PagedResponse<T>> listEntities(String nextPartitionKey, String nextRowKey,

diff --git a/sdk/tables/azure-data-tables/src/main/java/com/azure/data/tables/TableClient.java b/sdk/tables/azure-data-tables/src/main/java/com/azure/data/tables/TableClient.java
@@ -103,7 +103,6 @@
 public final class TableClient {
 
     private static final ExecutorService THREAD_POOL = TableUtils.getThreadPoolWithShutdownHook();
-    private static final String DELIMITER_CONTINUATION_TOKEN = ";";
     private final ClientLogger logger = new ClientLogger(TableClient.class);
     private final String tableName;
     private final AzureTableImpl tablesImplementation;
@@ -958,18 +957,12 @@ private <T extends TableEntity> PagedResponse<T> listEntitiesNextPage(String tok
             return null;
         }
 
-        String[] split = token.split(DELIMITER_CONTINUATION_TOKEN, 2);
-
-        if (split.length == 0) {
-            throw logger.logExceptionAsError(new RuntimeException(
-                "Split done incorrectly, must have partition key: " + token));
+        try {
+            String[] keys = TableUtils.getKeysFromToken(token);
+            return listEntities(keys[0], keys[1], context, options, resultType);
+        } catch (RuntimeException ex) {
+            throw logger.logExceptionAsError(ex);
         }
-
-        String nextPartitionKey = split[0];
-        String nextRowKey = split.length > 1 ? split[1] : null;
-
-        return listEntities(nextPartitionKey, nextRowKey, context, options, resultType);
-
     }
 
     private <T extends TableEntity> PagedResponse<T> listEntities(String nextPartitionKey, String nextRowKey,

diff --git a/...bles/azure-data-tables/src/main/java/com/azure/data/tables/implementation/TableUtils.java b/...bles/azure-data-tables/src/main/java/com/azure/data/tables/implementation/TableUtils.java
@@ -60,6 +60,7 @@
  */
 public final class TableUtils {
     private static final String UTF8_CHARSET = "UTF-8";
+    private static final String DELIMITER_CONTINUATION_TOKEN = ";";
     private static final String HTTP_REST_PROXY_SYNC_PROXY_ENABLE = "com.azure.core.http.restproxy.syncproxy.enable";
     private static final String TABLES_TRACING_NAMESPACE_VALUE = "Microsoft.Tables";
     private static final long THREADPOOL_SHUTDOWN_HOOK_TIMEOUT_SECINDS = 5;
@@ -595,4 +596,19 @@ public static Exception interpretException(Exception ex) {
             return (RuntimeException) mapThrowableToTableServiceException(exception);
         }
     }
+
+    public static String[] getKeysFromToken(String token) {
+        String[] split = token.split(DELIMITER_CONTINUATION_TOKEN, 2);
+        String[] keys = new String[2];
+        if (split.length == 0) {
+            throw new RuntimeException("Split done incorrectly, must have partition key. Token: " + token);
+        } else if (split.length == 2) {
+            keys[0] = split[0];
+            keys[1] = split[1];
+        } else {
+            keys[0] = split[0];
+            keys[1] = null;
+        }
+        return keys;
+    }
 }