Skip to content

Commit 4f826c5

Browse files
olavloiteolavloite
authored
test: support PostgreSQL dialect for RandomResultsGenerator (#1868)
Adds support for PostgreSQL dialect to RandomResultsGenerator
1 parent 70838bc commit 4f826c5

File tree

1 file changed

+96
-59
lines changed

1 file changed

+96
-59
lines changed

‎google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/RandomResultSetGenerator.java

Lines changed: 96 additions & 59 deletions
Original file line numberDiff line numberDiff line change
@@ -16,9 +16,10 @@
1616

1717
package com.google.cloud.spanner.connection;
1818

19-
import com.google.api.client.util.Base64;
2019
import com.google.cloud.Date;
2120
import com.google.cloud.Timestamp;
21+
import com.google.cloud.spanner.Dialect;
22+
import com.google.common.io.BaseEncoding;
2223
import com.google.protobuf.ListValue;
2324
import com.google.protobuf.NullValue;
2425
import com.google.protobuf.Value;
@@ -28,6 +29,7 @@
2829
import com.google.spanner.v1.StructType;
2930
import com.google.spanner.v1.StructType.Field;
3031
import com.google.spanner.v1.Type;
32+
import com.google.spanner.v1.TypeAnnotationCode;
3133
import com.google.spanner.v1.TypeCode;
3234
import java.math.BigDecimal;
3335
import java.util.Random;
@@ -37,86 +39,109 @@
3739
* of Cloud Spanner filled with random data.
3840
*/
3941
public class RandomResultSetGenerator {
40-
private static final Type[] TYPES =
41-
new Type[] {
42-
Type.newBuilder().setCode(TypeCode.BOOL).build(),
43-
Type.newBuilder().setCode(TypeCode.INT64).build(),
44-
Type.newBuilder().setCode(TypeCode.FLOAT64).build(),
45-
Type.newBuilder().setCode(TypeCode.NUMERIC).build(),
46-
Type.newBuilder().setCode(TypeCode.STRING).build(),
47-
Type.newBuilder().setCode(TypeCode.JSON).build(),
48-
Type.newBuilder().setCode(TypeCode.BYTES).build(),
49-
Type.newBuilder().setCode(TypeCode.DATE).build(),
50-
Type.newBuilder().setCode(TypeCode.TIMESTAMP).build(),
51-
Type.newBuilder()
52-
.setCode(TypeCode.ARRAY)
53-
.setArrayElementType(Type.newBuilder().setCode(TypeCode.BOOL))
54-
.build(),
55-
Type.newBuilder()
56-
.setCode(TypeCode.ARRAY)
57-
.setArrayElementType(Type.newBuilder().setCode(TypeCode.INT64))
58-
.build(),
59-
Type.newBuilder()
60-
.setCode(TypeCode.ARRAY)
61-
.setArrayElementType(Type.newBuilder().setCode(TypeCode.FLOAT64))
62-
.build(),
63-
Type.newBuilder()
64-
.setCode(TypeCode.ARRAY)
65-
.setArrayElementType(Type.newBuilder().setCode(TypeCode.NUMERIC))
66-
.build(),
67-
Type.newBuilder()
68-
.setCode(TypeCode.ARRAY)
69-
.setArrayElementType(Type.newBuilder().setCode(TypeCode.STRING))
70-
.build(),
71-
Type.newBuilder()
72-
.setCode(TypeCode.ARRAY)
73-
.setArrayElementType(Type.newBuilder().setCode(TypeCode.JSON))
74-
.build(),
75-
Type.newBuilder()
76-
.setCode(TypeCode.ARRAY)
77-
.setArrayElementType(Type.newBuilder().setCode(TypeCode.BYTES))
78-
.build(),
79-
Type.newBuilder()
80-
.setCode(TypeCode.ARRAY)
81-
.setArrayElementType(Type.newBuilder().setCode(TypeCode.DATE))
82-
.build(),
83-
Type.newBuilder()
84-
.setCode(TypeCode.ARRAY)
85-
.setArrayElementType(Type.newBuilder().setCode(TypeCode.TIMESTAMP))
86-
.build(),
87-
};
42+
private static Type[] generateTypes(Dialect dialect) {
43+
return new Type[] {
44+
Type.newBuilder().setCode(TypeCode.BOOL).build(),
45+
Type.newBuilder().setCode(TypeCode.INT64).build(),
46+
Type.newBuilder().setCode(TypeCode.FLOAT64).build(),
47+
dialect == Dialect.POSTGRESQL
48+
? Type.newBuilder()
49+
.setCode(TypeCode.NUMERIC)
50+
.setTypeAnnotation(TypeAnnotationCode.PG_NUMERIC)
51+
.build()
52+
: Type.newBuilder().setCode(TypeCode.NUMERIC).build(),
53+
Type.newBuilder().setCode(TypeCode.STRING).build(),
54+
Type.newBuilder()
55+
.setCode(dialect == Dialect.POSTGRESQL ? TypeCode.STRING : TypeCode.JSON)
56+
.build(),
57+
Type.newBuilder().setCode(TypeCode.BYTES).build(),
58+
Type.newBuilder().setCode(TypeCode.DATE).build(),
59+
Type.newBuilder().setCode(TypeCode.TIMESTAMP).build(),
60+
Type.newBuilder()
61+
.setCode(TypeCode.ARRAY)
62+
.setArrayElementType(Type.newBuilder().setCode(TypeCode.BOOL))
63+
.build(),
64+
Type.newBuilder()
65+
.setCode(TypeCode.ARRAY)
66+
.setArrayElementType(Type.newBuilder().setCode(TypeCode.INT64))
67+
.build(),
68+
Type.newBuilder()
69+
.setCode(TypeCode.ARRAY)
70+
.setArrayElementType(Type.newBuilder().setCode(TypeCode.FLOAT64))
71+
.build(),
72+
Type.newBuilder()
73+
.setCode(TypeCode.ARRAY)
74+
.setArrayElementType(
75+
dialect == Dialect.POSTGRESQL
76+
? Type.newBuilder()
77+
.setCode(TypeCode.NUMERIC)
78+
.setTypeAnnotation(TypeAnnotationCode.PG_NUMERIC)
79+
: Type.newBuilder().setCode(TypeCode.NUMERIC))
80+
.build(),
81+
Type.newBuilder()
82+
.setCode(TypeCode.ARRAY)
83+
.setArrayElementType(Type.newBuilder().setCode(TypeCode.STRING))
84+
.build(),
85+
Type.newBuilder()
86+
.setCode(TypeCode.ARRAY)
87+
.setArrayElementType(
88+
Type.newBuilder()
89+
.setCode(dialect == Dialect.POSTGRESQL ? TypeCode.STRING : TypeCode.JSON))
90+
.build(),
91+
Type.newBuilder()
92+
.setCode(TypeCode.ARRAY)
93+
.setArrayElementType(Type.newBuilder().setCode(TypeCode.BYTES))
94+
.build(),
95+
Type.newBuilder()
96+
.setCode(TypeCode.ARRAY)
97+
.setArrayElementType(Type.newBuilder().setCode(TypeCode.DATE))
98+
.build(),
99+
Type.newBuilder()
100+
.setCode(TypeCode.ARRAY)
101+
.setArrayElementType(Type.newBuilder().setCode(TypeCode.TIMESTAMP))
102+
.build(),
103+
};
104+
}
88105

89-
private static ResultSetMetadata generateMetadata() {
106+
private static ResultSetMetadata generateMetadata(Type[] types) {
90107
StructType.Builder rowTypeBuilder = StructType.newBuilder();
91-
for (int col = 0; col < TYPES.length; col++) {
92-
rowTypeBuilder.addFields(Field.newBuilder().setName("COL" + col).setType(TYPES[col])).build();
108+
for (int col = 0; col < types.length; col++) {
109+
rowTypeBuilder.addFields(Field.newBuilder().setName("COL" + col).setType(types[col])).build();
93110
}
94111
ResultSetMetadata.Builder builder = ResultSetMetadata.newBuilder();
95112
builder.setRowType(rowTypeBuilder.build());
96113
return builder.build();
97114
}
98115

99-
private static final ResultSetMetadata METADATA = generateMetadata();
100-
116+
private final ResultSetMetadata metadata;
117+
private final Dialect dialect;
118+
private final Type[] types;
101119
private final int rowCount;
102120
private final Random random = new Random();
103121

104122
public RandomResultSetGenerator(int rowCount) {
123+
this(rowCount, Dialect.GOOGLE_STANDARD_SQL);
124+
}
125+
126+
public RandomResultSetGenerator(int rowCount, Dialect dialect) {
105127
this.rowCount = rowCount;
128+
this.dialect = dialect;
129+
this.types = generateTypes(dialect);
130+
this.metadata = generateMetadata(types);
106131
}
107132

108133
public ResultSet generate() {
109134
ResultSet.Builder builder = ResultSet.newBuilder();
110135
for (int row = 0; row < rowCount; row++) {
111136
ListValue.Builder rowBuilder = ListValue.newBuilder();
112-
for (Type type : TYPES) {
137+
for (Type type : types) {
113138
Value.Builder valueBuilder = Value.newBuilder();
114139
setRandomValue(valueBuilder, type);
115140
rowBuilder.addValues(valueBuilder.build());
116141
}
117142
builder.addRows(rowBuilder.build());
118143
}
119-
builder.setMetadata(METADATA);
144+
builder.setMetadata(metadata);
120145
return builder.build();
121146
}
122147

@@ -142,7 +167,7 @@ private void setRandomValue(Value.Builder builder, Type type) {
142167
case BYTES:
143168
byte[] bytes = new byte[random.nextInt(200)];
144169
random.nextBytes(bytes);
145-
builder.setStringValue(Base64.encodeBase64String(bytes));
170+
builder.setStringValue(BaseEncoding.base64().encode(bytes));
146171
break;
147172
case JSON:
148173
builder.setStringValue("\"" + random.nextInt(200) + "\":\"" + random.nextInt(200) + "\"");
@@ -154,10 +179,18 @@ private void setRandomValue(Value.Builder builder, Type type) {
154179
builder.setStringValue(date.toString());
155180
break;
156181
case FLOAT64:
157-
builder.setNumberValue(random.nextDouble());
182+
if (randomNaN()) {
183+
builder.setNumberValue(Double.NaN);
184+
} else {
185+
builder.setNumberValue(random.nextDouble());
186+
}
158187
break;
159188
case NUMERIC:
160-
builder.setStringValue(BigDecimal.valueOf(random.nextDouble()).toString());
189+
if (dialect == Dialect.POSTGRESQL && randomNaN()) {
190+
builder.setStringValue("NaN");
191+
} else {
192+
builder.setStringValue(BigDecimal.valueOf(random.nextDouble()).toString());
193+
}
161194
break;
162195
case INT64:
163196
builder.setStringValue(String.valueOf(random.nextLong()));
@@ -184,4 +217,8 @@ private void setRandomValue(Value.Builder builder, Type type) {
184217
private boolean randomNull() {
185218
return random.nextInt(10) == 0;
186219
}
220+
221+
private boolean randomNaN() {
222+
return random.nextInt(10) == 0;
223+
}
187224
}

0 commit comments

Comments
 (0)