From e8f517574bd6229ff72f90233e090368dad30774 Mon Sep 17 00:00:00 2001
From: Akshat Jawne <akshatjawne@gmail.com>
Date: Fri, 28 Jun 2024 08:03:56 -0600
Subject: [PATCH] modify getKeyTable method

---
 .../web/client/api/JsPartitionedTable.java    | 50 +++++++++----------
 1 file changed, 25 insertions(+), 25 deletions(-)

diff --git a/web/client-api/src/main/java/io/deephaven/web/client/api/JsPartitionedTable.java b/web/client-api/src/main/java/io/deephaven/web/client/api/JsPartitionedTable.java
index 5f97a9bd9b5..c874a1feb9a 100644
--- a/web/client-api/src/main/java/io/deephaven/web/client/api/JsPartitionedTable.java
+++ b/web/client-api/src/main/java/io/deephaven/web/client/api/JsPartitionedTable.java
@@ -56,7 +56,7 @@ public class JsPartitionedTable extends HasLifecycle implements ServerObject {
     private final JsWidget widget;
     private List<String> keyColumnTypes;
     private PartitionedTableDescriptor descriptor;
-    private JsTable keys;
+    private Promise<JsTable> keys;
     private JsTable baseTable;
     private TableSubscription subscription;
 
@@ -83,18 +83,9 @@ public Promise<JsPartitionedTable> refetch() {
             return w.getExportedObjects()[0].fetch();
         }).then(result -> connection.newState((c, state, metadata) -> {
             baseTable = (JsTable) result;
-            SelectOrUpdateRequest view = new SelectOrUpdateRequest();
-            view.setSourceId(baseTable.state().getHandle().makeTableReference());
-            view.setResultId(state.getHandle().makeTicket());
-            view.setColumnSpecsList(descriptor.getKeyColumnNamesList());
-            connection.tableServiceClient().view(view, metadata, (fail, success) -> {
-                c.apply(fail, success);
-            });
-        }, "get keys table")
+        }, "get base table")
                 .refetch(this, connection.metadata())
                 .then(state -> Promise.resolve(new JsTable(connection, state)))).then(result -> {
-                    keys = result;
-
                     keyColumnTypes = new ArrayList<>();
                     InitialTableDefinition tableDefinition = WebBarrageUtils.readTableDefinition(
                             WebBarrageUtils.readSchemaMessage(descriptor.getConstituentDefinitionSchema_asU8()));
@@ -110,7 +101,7 @@ public Promise<JsPartitionedTable> refetch() {
                     JsArray<String> keyColumnNames = descriptor.getKeyColumnNamesList();
                     for (int i = 0; i < keyColumnNames.length; i++) {
                         String name = keyColumnNames.getAt(i);
-                        Column keyColumn = keys.findColumn(name);
+                        Column keyColumn = baseTable.findColumn(name);
                         keyColumnTypes.add(keyColumn.getType());
                         keyColumns[keyColumns.length] = keyColumn;
                     }
@@ -118,9 +109,9 @@ public Promise<JsPartitionedTable> refetch() {
                     this.keyColumns = JsObject.freeze(keyColumns);
 
                     // TODO(deephaven-core#3604) in case of a new session, we should do a full refetch
-                    keys.addEventListener(JsTable.EVENT_DISCONNECT, event -> fireEvent(EVENT_DISCONNECT));
-                    keys.addEventListener(JsTable.EVENT_RECONNECT, event -> {
-                        subscribeToKeys().then(ignore -> {
+                    baseTable.addEventListener(JsTable.EVENT_DISCONNECT, event -> fireEvent(EVENT_DISCONNECT));
+                    baseTable.addEventListener(JsTable.EVENT_RECONNECT, event -> {
+                        subscribeToBaseTable().then(ignore -> {
                             unsuppressEvents();
                             fireEvent(EVENT_RECONNECT);
                             return null;
@@ -133,7 +124,7 @@ public Promise<JsPartitionedTable> refetch() {
                             return null;
                         });
                     });
-                    return subscribeToKeys();
+                    return subscribeToBaseTable();
                 });
     }
 
@@ -142,14 +133,14 @@ public TypedTicket typedTicket() {
         return widget.typedTicket();
     }
 
-    private Promise<JsPartitionedTable> subscribeToKeys() {
-        subscription = keys.subscribe(
-                JsArray.asJsArray(keys.findColumns(descriptor.getKeyColumnNamesList().asArray(new String[0]))));
+    private Promise<JsPartitionedTable> subscribeToBaseTable() {
+        subscription = baseTable.subscribe(
+                JsArray.asJsArray(baseTable.findColumns(descriptor.getKeyColumnNamesList().asArray(new String[0]))));
         subscription.addEventListener(TableSubscription.EVENT_UPDATED, this::handleKeys);
 
         LazyPromise<JsPartitionedTable> promise = new LazyPromise<>();
         subscription.addEventListenerOneShot(TableSubscription.EVENT_UPDATED, data -> promise.succeed(this));
-        keys.addEventListener(JsTable.EVENT_DISCONNECT, e -> promise.fail("Underlying table disconnected"));
+        baseTable.addEventListener(JsTable.EVENT_DISCONNECT, e -> promise.fail("Underlying table disconnected"));
         return promise.asPromise();
     }
 
@@ -283,9 +274,21 @@ public Column[] getColumns() {
      *
      * @return Promise of a Table
      */
-    @JsMethod
+    @Deprecated
     public Promise<JsTable> getKeyTable() {
-        return keys.copy();
+        if (keys == null) {
+            keys = connection.newState((c, state, metadata) -> {
+                SelectOrUpdateRequest view = new SelectOrUpdateRequest();
+                view.setSourceId(baseTable.state().getHandle().makeTableReference());
+                view.setResultId(state.getHandle().makeTicket());
+                view.setColumnSpecsList(descriptor.getKeyColumnNamesList());
+                connection.tableServiceClient().view(view, metadata, c::apply);
+            }, "view only keys columns")
+                    .refetch(this, connection.metadata())
+                    .then(state -> Promise.resolve(new JsTable(state.getConnection(), state)));
+        }
+        assert keys != null;
+        return keys.then(JsTable::copy);
     }
 
     /**
@@ -300,9 +303,6 @@ public Promise<JsTable> getBaseTable() {
 
     /** Close any subscriptions to underlying tables or key tables */
     private void closeSubscriptions() {
-        if (keys != null) {
-            keys.close();
-        }
         if (baseTable != null) {
             baseTable.close();
         }