From b24d85cee79453dbc755135fbfeeaebb164cd848 Mon Sep 17 00:00:00 2001 From: Benoit TELLIER Date: Sun, 10 May 2026 23:22:27 +0200 Subject: [PATCH] Expose getters from CREATE statements --- .../schema/CreateTableWithOptions.java | 22 +++++++ .../schema/DefaultCreateTable.java | 32 ++++++++++ .../querybuilder/schema/CreateTableTest.java | 62 +++++++++++++++++++ 3 files changed, 116 insertions(+) diff --git a/query-builder/src/main/java/com/datastax/oss/driver/api/querybuilder/schema/CreateTableWithOptions.java b/query-builder/src/main/java/com/datastax/oss/driver/api/querybuilder/schema/CreateTableWithOptions.java index c7bddf575fb..d107041689d 100644 --- a/query-builder/src/main/java/com/datastax/oss/driver/api/querybuilder/schema/CreateTableWithOptions.java +++ b/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 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 getColumns(); + + /** + * Returns the partition-key columns declared via {@code withPartitionKey()}, in declaration + * order. + */ + @NonNull + List getPartitionKeyColumns(); + + /** + * Returns the clustering columns declared via {@code withClusteringColumn()}, in declaration + * order. + */ + @NonNull + List getClusteringColumns(); } diff --git a/query-builder/src/main/java/com/datastax/oss/driver/internal/querybuilder/schema/DefaultCreateTable.java b/query-builder/src/main/java/com/datastax/oss/driver/internal/querybuilder/schema/DefaultCreateTable.java index 058aeccdd24..91c0ef0dc95 100644 --- a/query-builder/src/main/java/com/datastax/oss/driver/internal/querybuilder/schema/DefaultCreateTable.java +++ b/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 getRegularColumns() { public ImmutableMap getOrderings() { return orderings; } + + @NonNull + @Override + public List getColumns() { + return toColumnSpecList(regularColumns); + } + + @NonNull + @Override + public List getPartitionKeyColumns() { + return toColumnSpecList(partitionKeyColumns); + } + + @NonNull + @Override + public List getClusteringColumns() { + return toColumnSpecList(clusteringKeyColumns); + } + + private List toColumnSpecList(ImmutableSet ids) { + ImmutableList.Builder result = ImmutableList.builderWithExpectedSize(ids.size()); + // Walk columnsInOrder to preserve declaration order + for (Map.Entry entry : columnsInOrder.entrySet()) { + if (ids.contains(entry.getKey())) { + result.add(ColumnSpec.of(entry.getKey(), entry.getValue())); + } + } + return result.build(); + } } diff --git a/query-builder/src/test/java/com/datastax/oss/driver/api/querybuilder/schema/CreateTableTest.java b/query-builder/src/test/java/com/datastax/oss/driver/api/querybuilder/schema/CreateTableTest.java index 31efc278472..1f56ad083c6 100644 --- a/query-builder/src/test/java/com/datastax/oss/driver/api/querybuilder/schema/CreateTableTest.java +++ b/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)"); } + + @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 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 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 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); + } }