refactor: extract column projection logic into SystemColumnProjection

Move the three *_COLUMNS_OF_INTEREST constants, the three volatile cache
fields (localColumns, peersColumns, peersV2Columns), and the two static
helpers (intersectWithNeeded, buildProjectedQuery) out of ControlConnection
and into a new package-private SystemColumnProjection class.

ControlConnection now holds a single 'projection' instance and delegates
all projection state and query-building to it. resetColumnCaches() remains
on ControlConnection as a @VisibleForTesting thin wrapper over
projection.reset(), used by ControlConnectionTest between Scassandra prime
clears.

ControlConnectionUnitTest is updated to test SystemColumnProjection
directly; ScassandraCluster references updated from ControlConnection.*
to SystemColumnProjection.*.

Author: nikagra

driver-core/src/main/java/com/datastax/driver/core/ControlConnection.java

@@ -34,7 +34,6 @@
 import com.datastax.driver.core.utils.MoreFutures;
 import com.datastax.driver.core.utils.MoreObjects;
 import com.google.common.annotations.VisibleForTesting;
-import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Iterators;
 import com.google.common.util.concurrent.FutureCallback;
 import com.google.common.util.concurrent.Futures;
@@ -87,91 +86,6 @@ class ControlConnection implements Connection.Owner {
   private static final String SELECT_SCHEMA_LOCAL =
       "SELECT schema_version, host_id FROM system.local WHERE key='local'";
 
-  // IMPORTANT: Every column read from system.local rows — in updateInfo(),
-  // refreshNodeListAndTokenMap(), isValidPeer(), and DefaultEndPointFactory — MUST be listed here.
-  // If a new column read is added anywhere that consumes a system table row, add it to the
-  // appropriate set below, otherwise it will be silently excluded from projected queries.
-  @VisibleForTesting
-  static final ImmutableSet<String> LOCAL_COLUMNS_OF_INTEREST =
-      ImmutableSet.of(
-          "cluster_name",
-          "partitioner",
-          "data_center",
-          "rack",
-          "release_version",
-          "native_address",
-          "native_port",
-          "native_transport_address",
-          "native_transport_port",
-          "native_transport_port_ssl",
-          "rpc_address",
-          "broadcast_address",
-          "broadcast_port",
-          "listen_address",
-          "listen_port",
-          "tokens",
-          "host_id",
-          "schema_version",
-          "workload",
-          "graph",
-          "dse_version");
-
-  // IMPORTANT: see LOCAL_COLUMNS_OF_INTEREST note above.
-  // Includes all columns consumed by updateInfo(), refreshNodeListAndTokenMap(),
-  // isValidPeer(), and DefaultEndPointFactory.create() from system.peers rows.
-  // Columns that are absent from the actual server schema are silently excluded by
-  // intersectWithNeeded(), so listing extra columns here is safe.
-  @VisibleForTesting
-  static final ImmutableSet<String> PEERS_COLUMNS_OF_INTEREST =
-      ImmutableSet.of(
-          "peer",
-          "peer_port", // peers_v2 column; harmless to list here — absent on peers, excluded safely
-          "rpc_address",
-          "data_center",
-          "rack",
-          "release_version",
-          "tokens",
-          "listen_address",
-          "listen_port",
-          "host_id",
-          "schema_version",
-          "native_address", // may appear on some server variants; guarded by contains() in code
-          "native_port", // same
-          "native_transport_address",
-          "native_transport_port",
-          "native_transport_port_ssl",
-          "workload",
-          "graph",
-          "dse_version");
-
-  // IMPORTANT: see LOCAL_COLUMNS_OF_INTEREST note above.
-  // Includes all columns consumed by updateInfo(), refreshNodeListAndTokenMap(),
-  // isValidPeer(), and DefaultEndPointFactory.create() from system.peers_v2 rows.
-  // Columns that are absent from the actual server schema are silently excluded by
-  // intersectWithNeeded(), so listing extra columns here is safe.
-  @VisibleForTesting
-  static final ImmutableSet<String> PEERS_V2_COLUMNS_OF_INTEREST =
-      ImmutableSet.of(
-          "peer",
-          "peer_port",
-          "native_address",
-          "native_port",
-          "data_center",
-          "rack",
-          "release_version",
-          "tokens",
-          "host_id",
-          "schema_version",
-          "workload",
-          "graph",
-          "dse_version",
-          "listen_address",
-          "listen_port",
-          "rpc_address", // legacy; guarded by contains() in code — harmless if absent
-          "native_transport_address", // same
-          "native_transport_port", // same
-          "native_transport_port_ssl"); // same
-
   private static final VersionNumber _3_11 = VersionNumber.parse("3.11.0");
 
   @VisibleForTesting
@@ -188,13 +102,7 @@ class ControlConnection implements Connection.Owner {
   // from here on out.
   private volatile boolean isPeersV2 = true;
 
-  // Column projection caches. null = uninitialized: the first query to each system table issues
-  // SELECT * to discover which columns the server exposes, then subsequent queries project only
-  // the intersection of those columns with the corresponding *_COLUMNS_OF_INTEREST set.
-  // Reset to null on every new connection so that the new server's schema is re-discovered.
-  private volatile Set<String> localColumns = null;
-  private volatile Set<String> peersColumns = null;
-  private volatile Set<String> peersV2Columns = null;
+  private final SystemColumnProjection projection = new SystemColumnProjection();
 
   public ControlConnection(Cluster.Manager manager) {
     this.cluster = manager;
@@ -207,9 +115,7 @@ public ControlConnection(Cluster.Manager manager) {
    */
   @VisibleForTesting
   void resetColumnCaches() {
-    localColumns = null;
-    peersColumns = null;
-    peersV2Columns = null;
+    projection.reset();
   }
 
   // Only for the initial connection. Does not schedule retries if it fails
@@ -436,9 +342,7 @@ private Connection tryConnect(Host host, boolean isInitialConnection)
       // which columns this server exposes, rather than a projected query built for a
       // previous connection's server. The caches are populated during the queries below
       // and remain warm for the lifetime of this connection.
-      localColumns = null;
-      peersColumns = null;
-      peersV2Columns = null;
+      projection.reset();
 
       // We need to refresh the node list first so we know about the cassandra version of
       // the node we're connecting to.
@@ -567,6 +471,11 @@ void refreshNodeListAndTokenMap() {
     } catch (ExecutionException e) {
       // If we're being shutdown during refresh, this can happen. That's fine so don't scare the
       // user.
+      if (e.getCause() instanceof InvalidQueryException) {
+        // A projected query referenced a column the server no longer exposes; reset caches so
+        // the next connection re-discovers columns via SELECT *.
+        projection.reset();
+      }
       if (!isShutdown)
         logger.error(
             "[Control connection] Unexpected error while refreshing node list and token map", e);
@@ -605,10 +514,7 @@ private Row fetchNodeInfo(Host host, Connection c)
     if (isConnectedHost || host.getBroadcastSocketAddress() != null) {
       String query;
       if (isConnectedHost) {
-        query =
-            localColumns == null
-                ? SELECT_LOCAL
-                : buildProjectedQuery("system.local", localColumns, "key='local'");
+        query = projection.localQuery();
       } else {
         InetSocketAddress broadcastAddress = host.getBroadcastSocketAddress();
         // Always use SELECT * for single-row WHERE lookups. Projected queries are only used for
@@ -641,17 +547,11 @@ private Row fetchNodeInfo(Host host, Connection c)
       // the projected columns.
       if (row != null) {
         if (isConnectedHost) {
-          if (localColumns == null) {
-            localColumns = intersectWithNeeded(rs, LOCAL_COLUMNS_OF_INTEREST);
-          }
+          projection.populateLocal(rs);
         } else if (isPeersV2) {
-          if (peersV2Columns == null) {
-            peersV2Columns = intersectWithNeeded(rs, PEERS_V2_COLUMNS_OF_INTEREST);
-          }
+          projection.populatePeersV2(rs);
         } else {
-          if (peersColumns == null) {
-            peersColumns = intersectWithNeeded(rs, PEERS_COLUMNS_OF_INTEREST);
-          }
+          projection.populatePeers(rs);
         }
         return row;
       } else {
@@ -724,6 +624,11 @@ boolean refreshNodeInfo(Host host) {
     } catch (ExecutionException e) {
       // If we're being shutdown during refresh, this can happen. That's fine so don't scare the
       // user.
+      if (e.getCause() instanceof InvalidQueryException) {
+        // A projected query referenced a column the server no longer exposes; reset caches so
+        // the next connection re-discovers columns via SELECT *.
+        projection.reset();
+      }
       if (!isShutdown)
         logger.debug("[Control connection] Unexpected error while refreshing node info", e);
       signalError();
@@ -859,13 +764,9 @@ private static void updateLocationInfo(
    */
   private ListenableFuture<ResultSet> selectPeersFuture(final Connection connection) {
     if (isPeersV2) {
-      String peersV2Query =
-          peersV2Columns == null
-              ? SELECT_PEERS_V2
-              : buildProjectedQuery("system.peers_v2", peersV2Columns, null);
       DefaultResultSetFuture peersV2Future =
           new DefaultResultSetFuture(
-              null, cluster.protocolVersion(), new Requests.Query(peersV2Query));
+              null, cluster.protocolVersion(), new Requests.Query(projection.peersV2Query()));
       connection.write(peersV2Future);
       final SettableFuture<ResultSet> peersFuture = SettableFuture.create();
       // if peers v2 query fails, query peers table instead.
@@ -875,9 +776,7 @@ private ListenableFuture<ResultSet> selectPeersFuture(final Connection connectio
 
             @Override
             public void onSuccess(ResultSet result) {
-              if (peersV2Columns == null) {
-                peersV2Columns = intersectWithNeeded(result, PEERS_V2_COLUMNS_OF_INTEREST);
-              }
+              projection.populatePeersV2(result);
               peersFuture.set(result);
             }
 
@@ -891,9 +790,9 @@ public void onFailure(Throwable t) {
                   || (t instanceof ServerError
                       && t.getMessage().contains("Unknown keyspace/cf pair (system.peers_v2)"))) {
                 isPeersV2 = false;
-                // Also reset the peers cache so the first system.peers query issues SELECT *
-                // to discover which columns that table exposes on this server.
-                peersColumns = null;
+                // Reset all caches: peersV2Columns is now stale, and peers cache should be cleared
+                // so the first system.peers query re-discovers columns via SELECT *.
+                projection.reset();
                 MoreFutures.propagateFuture(peersFuture, selectPeersFuture(connection));
               } else {
                 peersFuture.setException(t);
@@ -903,28 +802,27 @@ public void onFailure(Throwable t) {
           MoreExecutors.directExecutor());
       return peersFuture;
     } else {
-      String peersQuery =
-          peersColumns == null
-              ? SELECT_PEERS
-              : buildProjectedQuery("system.peers", peersColumns, null);
       DefaultResultSetFuture rawFuture =
           new DefaultResultSetFuture(
-              null, cluster.protocolVersion(), new Requests.Query(peersQuery));
+              null, cluster.protocolVersion(), new Requests.Query(projection.peersQuery()));
       connection.write(rawFuture);
       final SettableFuture<ResultSet> peersFuture = SettableFuture.create();
       Futures.addCallback(
           rawFuture,
           new FutureCallback<ResultSet>() {
             @Override
             public void onSuccess(ResultSet result) {
-              if (peersColumns == null) {
-                peersColumns = intersectWithNeeded(result, PEERS_COLUMNS_OF_INTEREST);
-              }
+              projection.populatePeers(result);
               peersFuture.set(result);
             }
 
             @Override
             public void onFailure(Throwable t) {
+              if (t instanceof InvalidQueryException) {
+                // The projected query referenced a column the server no longer exposes; reset
+                // caches so the next query re-discovers columns via SELECT *.
+                projection.reset();
+              }
               peersFuture.setException(t);
             }
           },
@@ -933,35 +831,6 @@ public void onFailure(Throwable t) {
     }
   }
 
-  /**
-   * Returns the intersection of the columns returned by the server (from {@code rs}) with the given
-   * {@code needed} set, or {@code null} if the intersection is empty. The result is used to cache
-   * projected column lists so subsequent queries fetch only what the driver actually reads. A
-   * {@code null} return keeps the cache in the "uninitialized" sentinel state, ensuring the driver
-   * continues issuing {@code SELECT *} rather than generating an invalid empty-column projection.
-   */
-  @VisibleForTesting
-  static Set<String> intersectWithNeeded(ResultSet rs, ImmutableSet<String> needed) {
-    ImmutableSet.Builder<String> result = ImmutableSet.builder();
-    for (ColumnDefinitions.Definition def : rs.getColumnDefinitions()) {
-      if (needed.contains(def.getName())) {
-        result.add(def.getName());
-      }
-    }
-    ImmutableSet<String> built = result.build();
-    return built.isEmpty() ? null : built;
-  }
-
-  /**
-   * Builds a {@code SELECT col1, col2, ... FROM table [WHERE whereClause]} query string from the
-   * given projected column set. {@code whereClause} may be {@code null} for table-wide scans.
-   */
-  @VisibleForTesting
-  static String buildProjectedQuery(String table, Set<String> columns, String whereClause) {
-    String query = "SELECT " + String.join(", ", columns) + " FROM " + table;
-    return whereClause != null ? query + " WHERE " + whereClause : query;
-  }
-
   private void refreshNodeListAndTokenMap(
       final Connection connection,
       final Cluster.Manager cluster,
@@ -975,12 +844,9 @@ private void refreshNodeListAndTokenMap(
 
     // Make sure we're up to date on nodes and tokens
 
-    String localQuery =
-        localColumns == null
-            ? SELECT_LOCAL
-            : buildProjectedQuery("system.local", localColumns, "key='local'");
     DefaultResultSetFuture localFuture =
-        new DefaultResultSetFuture(null, cluster.protocolVersion(), new Requests.Query(localQuery));
+        new DefaultResultSetFuture(
+            null, cluster.protocolVersion(), new Requests.Query(projection.localQuery()));
     ListenableFuture<ResultSet> peersFuture = selectPeersFuture(connection);
     connection.write(localFuture);
 
@@ -990,9 +856,7 @@ private void refreshNodeListAndTokenMap(
 
     // Update cluster name, DC and rack for the one node we are connected to
     ResultSet localRs = localFuture.get();
-    if (localColumns == null) {
-      localColumns = intersectWithNeeded(localRs, LOCAL_COLUMNS_OF_INTEREST);
-    }
+    projection.populateLocal(localRs);
     Row localRow = localRs.one();
     if (localRow == null) {
       throw new IllegalStateException(