Expose getters from CREATE statements

query-builder/src/main/java/com/datastax/oss/driver/api/querybuilder/schema/CreateTableWithOptions.java

@@ -22,6 +22,7 @@
 import com.datastax.oss.driver.shaded.guava.common.collect.Maps;
 import edu.umd.cs.findbugs.annotations.CheckReturnValue;
 import edu.umd.cs.findbugs.annotations.NonNull;
+import java.util.List;
 import java.util.Map;
 
 public interface CreateTableWithOptions
@@ -37,4 +38,25 @@ public interface CreateTableWithOptions
   default CreateTableWithOptions withExtensions(@NonNull Map<String, byte[]> extensions) {
     return withOption("extensions", Maps.transformValues(extensions, RawOptionsWrapper::of));
   }
+
+  /**
+   * Returns the regular (non-key, non-static) columns declared via {@code withColumn()}, in
+   * declaration order.
+   */
+  @NonNull
+  List<ColumnSpec> getColumns();
+
+  /**
+   * Returns the partition-key columns declared via {@code withPartitionKey()}, in declaration
+   * order.
+   */
+  @NonNull
+  List<ColumnSpec> getPartitionKeyColumns();
+
+  /**
+   * Returns the clustering columns declared via {@code withClusteringColumn()}, in declaration
+   * order.
+   */
+  @NonNull
+  List<ColumnSpec> getClusteringColumns();
 }

query-builder/src/main/java/com/datastax/oss/driver/internal/querybuilder/schema/DefaultCreateTable.java

@@ -22,15 +22,18 @@
 import com.datastax.oss.driver.api.core.CqlIdentifier;
 import com.datastax.oss.driver.api.core.metadata.schema.ClusteringOrder;
 import com.datastax.oss.driver.api.core.type.DataType;
+import com.datastax.oss.driver.api.querybuilder.schema.ColumnSpec;
 import com.datastax.oss.driver.api.querybuilder.schema.CreateTable;
 import com.datastax.oss.driver.api.querybuilder.schema.CreateTableStart;
 import com.datastax.oss.driver.api.querybuilder.schema.CreateTableWithOptions;
 import com.datastax.oss.driver.internal.querybuilder.CqlHelper;
 import com.datastax.oss.driver.internal.querybuilder.ImmutableCollections;
+import com.datastax.oss.driver.shaded.guava.common.collect.ImmutableList;
 import com.datastax.oss.driver.shaded.guava.common.collect.ImmutableMap;
 import com.datastax.oss.driver.shaded.guava.common.collect.ImmutableSet;
 import edu.umd.cs.findbugs.annotations.NonNull;
 import edu.umd.cs.findbugs.annotations.Nullable;
+import java.util.List;
 import java.util.Map;
 import net.jcip.annotations.Immutable;
 
@@ -395,4 +398,33 @@ public ImmutableSet<CqlIdentifier> getRegularColumns() {
   public ImmutableMap<CqlIdentifier, ClusteringOrder> getOrderings() {
     return orderings;
   }
+
+  @NonNull
+  @Override
+  public List<ColumnSpec> getColumns() {
+    return toColumnSpecList(regularColumns);
+  }
+
+  @NonNull
+  @Override
+  public List<ColumnSpec> getPartitionKeyColumns() {
+    return toColumnSpecList(partitionKeyColumns);
+  }
+
+  @NonNull
+  @Override
+  public List<ColumnSpec> getClusteringColumns() {
+    return toColumnSpecList(clusteringKeyColumns);
+  }
+
+  private List<ColumnSpec> toColumnSpecList(ImmutableSet<CqlIdentifier> ids) {
+    ImmutableList.Builder<ColumnSpec> result = ImmutableList.builderWithExpectedSize(ids.size());
+    // Walk columnsInOrder to preserve declaration order
+    for (Map.Entry<CqlIdentifier, DataType> entry : columnsInOrder.entrySet()) {
+      if (ids.contains(entry.getKey())) {
+        result.add(ColumnSpec.of(entry.getKey(), entry.getValue()));
+      }
+    }
+    return result.build();
+  }
 }

query-builder/src/test/java/com/datastax/oss/driver/api/querybuilder/schema/CreateTableTest.java

@@ -20,7 +20,9 @@
 import static com.datastax.oss.driver.api.querybuilder.Assertions.assertThat;
 import static com.datastax.oss.driver.api.querybuilder.SchemaBuilder.createTable;
 import static com.datastax.oss.driver.api.querybuilder.SchemaBuilder.udt;
+import static org.assertj.core.api.Assertions.assertThat;
 
+import com.datastax.oss.driver.api.core.CqlIdentifier;
 import com.datastax.oss.driver.api.core.metadata.schema.ClusteringOrder;
 import com.datastax.oss.driver.api.core.type.DataTypes;
 import com.datastax.oss.driver.api.querybuilder.SchemaBuilder;
@@ -29,6 +31,7 @@
 import com.datastax.oss.driver.api.querybuilder.schema.compaction.TimeWindowCompactionStrategy.TimestampResolution;
 import com.datastax.oss.driver.shaded.guava.common.collect.ImmutableMap;
 import java.nio.charset.StandardCharsets;
+import java.util.List;
 import org.junit.Test;
 
 public class CreateTableTest {
@@ -378,4 +381,63 @@ public void should_generate_vector_column() {
                 .withColumn("v", DataTypes.vectorOf(DataTypes.FLOAT, 3)))
         .hasCql("CREATE TABLE foo (k int PRIMARY KEY,v vector<float, 3>)");
   }
+
+  @Test
+  public void should_expose_partition_key_columns() {
+    CreateTableWithOptions table =
+        createTable("foo")
+            .withPartitionKey("pk1", DataTypes.INT)
+            .withPartitionKey("pk2", DataTypes.TEXT)
+            .withClusteringColumn("cc", DataTypes.TIMEUUID)
+            .withColumn("v", DataTypes.BLOB);
+
+    List<ColumnSpec> pk = table.getPartitionKeyColumns();
+    assertThat(pk).hasSize(2);
+    assertThat(pk.get(0).getName()).isEqualTo(CqlIdentifier.fromCql("pk1"));
+    assertThat(pk.get(0).getType()).isEqualTo(DataTypes.INT);
+    assertThat(pk.get(1).getName()).isEqualTo(CqlIdentifier.fromCql("pk2"));
+    assertThat(pk.get(1).getType()).isEqualTo(DataTypes.TEXT);
+  }
+
+  @Test
+  public void should_expose_clustering_columns() {
+    CreateTableWithOptions table =
+        createTable("foo")
+            .withPartitionKey("pk", DataTypes.INT)
+            .withClusteringColumn("cc1", DataTypes.TIMEUUID)
+            .withClusteringColumn("cc2", DataTypes.TEXT)
+            .withColumn("v", DataTypes.BLOB);
+
+    List<ColumnSpec> cc = table.getClusteringColumns();
+    assertThat(cc).hasSize(2);
+    assertThat(cc.get(0).getName()).isEqualTo(CqlIdentifier.fromCql("cc1"));
+    assertThat(cc.get(0).getType()).isEqualTo(DataTypes.TIMEUUID);
+    assertThat(cc.get(1).getName()).isEqualTo(CqlIdentifier.fromCql("cc2"));
+    assertThat(cc.get(1).getType()).isEqualTo(DataTypes.TEXT);
+  }
+
+  @Test
+  public void should_expose_regular_columns() {
+    CreateTableWithOptions table =
+        createTable("foo")
+            .withPartitionKey("pk", DataTypes.INT)
+            .withColumn("v1", DataTypes.TEXT)
+            .withColumn("v2", DataTypes.BIGINT);
+
+    List<ColumnSpec> cols = table.getColumns();
+    assertThat(cols).hasSize(2);
+    assertThat(cols.get(0).getName()).isEqualTo(CqlIdentifier.fromCql("v1"));
+    assertThat(cols.get(0).getType()).isEqualTo(DataTypes.TEXT);
+    assertThat(cols.get(1).getName()).isEqualTo(CqlIdentifier.fromCql("v2"));
+    assertThat(cols.get(1).getType()).isEqualTo(DataTypes.BIGINT);
+  }
+
+  @Test
+  public void should_return_empty_lists_when_no_columns_added() {
+    CreateTableWithOptions table = createTable("foo").withPartitionKey("pk", DataTypes.INT);
+
+    assertThat(table.getColumns()).isEmpty();
+    assertThat(table.getClusteringColumns()).isEmpty();
+    assertThat(table.getPartitionKeyColumns()).hasSize(1);
+  }
 }