Java Code Examples for com.google.cloud.spanner.Struct

The following examples show how to use com.google.cloud.spanner.Struct. These examples are extracted from open source projects. You can vote up the ones you like or vote down the ones you don't like, and go to the original project or source file by following the links above each example. You may check out the related API usage on the sidebar.
Example 1
Source Project: DataflowTemplates   Source File: SpannerConverters.java    License: Apache License 2.0 7 votes vote down vote up
/** Function to get all column names from the table. */
private LinkedHashMap<String, String> getAllColumns(ReadContext context, String tableName) {
  LinkedHashMap<String, String> columns = Maps.newLinkedHashMap();
  ResultSet resultSet =
      context.executeQuery(
          Statement.newBuilder(
                  "SELECT COLUMN_NAME, SPANNER_TYPE FROM INFORMATION_SCHEMA.COLUMNS "
                      + "WHERE [email protected]_name ORDER BY ORDINAL_POSITION")
              .bind("table_name")
              .to(tableName)
              .build());
  LOG.info("Got schema information. Reading columns.");
  while (resultSet.next()) {
    Struct currentRow = resultSet.getCurrentRowAsStruct();
    columns.put(currentRow.getString(0), currentRow.getString(1));
  }
  return columns;
}
 
Example 2
Source Project: DataflowTemplates   Source File: SpannerConverters.java    License: Apache License 2.0 6 votes vote down vote up
/**
 * Prints Struct as a CSV String.
 *
 * @param struct Spanner Struct.
 * @return Spanner Struct encoded as a CSV String.
 */
public String print(Struct struct) {
  StringWriter stringWriter = new StringWriter();
  try {
    CSVPrinter printer =
        new CSVPrinter(stringWriter, CSVFormat.DEFAULT.withRecordSeparator("")
            .withQuoteMode(QuoteMode.ALL_NON_NULL));
    LinkedHashMap<String, BiFunction<Struct, String, String>> parsers = Maps.newLinkedHashMap();
    parsers.putAll(mapColumnParsers(struct.getType().getStructFields()));
    List<String> values = parseResultSet(struct, parsers);
    printer.printRecord(values);
  } catch (IOException e) {
    throw new RuntimeException(e);
  }
  return stringWriter.toString();
}
 
Example 3
Source Project: google-cloud-java   Source File: DatabaseClientSnippets.java    License: Apache License 2.0 6 votes vote down vote up
/** Example of read only transaction. */
// [TARGET readOnlyTransaction()]
// [VARIABLE my_singer_id]
// [VARIABLE my_album_id]
public String readOnlyTransaction(long singerId, long albumId) {
  // [START readOnlyTransaction]
  String singerColumn = "FirstName";
  String albumColumn = "AlbumTitle";
  String albumTitle = null;
  // ReadOnlyTransaction should be closed to prevent resource leak.
  try (ReadOnlyTransaction txn = dbClient.readOnlyTransaction()) {
    Struct singerRow =
        txn.readRow("Singers", Key.of(singerId), Collections.singleton(singerColumn));
    Struct albumRow =
        txn.readRow("Albums", Key.of(singerId, albumId), Collections.singleton(albumColumn));
    singerRow.getString(singerColumn);
    albumTitle = albumRow.getString(albumColumn);
  }
  // [END readOnlyTransaction]
  return albumTitle;
}
 
Example 4
Source Project: beam   Source File: NaiveSpannerRead.java    License: Apache License 2.0 6 votes vote down vote up
@Override
public PCollection<Struct> expand(PCollection<ReadOperation> input) {
  PCollectionView<Transaction> txView = getTxView();
  if (txView == null) {
    Pipeline begin = input.getPipeline();
    SpannerIO.CreateTransaction createTx =
        SpannerIO.createTransaction()
            .withSpannerConfig(getSpannerConfig())
            .withTimestampBound(getTimestampBound());
    txView = begin.apply(createTx);
  }

  return input.apply(
      "Naive read from Cloud Spanner",
      ParDo.of(new NaiveSpannerReadFn(getSpannerConfig(), txView)).withSideInputs(txView));
}
 
Example 5
@Test
public void allKnownMappingTypesTest() throws NoSuchFieldException {
	for (Method method : Struct.class.getMethods()) {
		String methodName = method.getName();
		// ignoring private methods, ones not named like a getter. Getters must also
		// only take the column index or name
		if (!Modifier.isPublic(method.getModifiers()) || !methodName.startsWith("get")
				|| method.getParameterCount() != 1
				|| DISREGARDED_METHOD_NAMES.contains(methodName)) {
			continue;
		}
		Class returnType = ConversionUtils.boxIfNeeded(method.getReturnType());
		if (ConversionUtils.isIterableNonByteArrayType(returnType)) {
			Class innerReturnType = (Class) ((ParameterizedType) method
					.getGenericReturnType()).getActualTypeArguments()[0];
			assertThat(StructAccessor.readIterableMapping.keySet()).contains(innerReturnType);
		}
		else {
			assertThat(StructAccessor.singleItemReadMethodMapping.keySet()).contains(returnType);
		}
	}
}
 
Example 6
Source Project: google-cloud-java   Source File: DatabaseClientSnippets.java    License: Apache License 2.0 6 votes vote down vote up
/** Example of using {@link TransactionManager}. */
// [TARGET transactionManager()]
// [VARIABLE my_singer_id]
public void transactionManager(final long singerId) throws InterruptedException {
  // [START transactionManager]
  try (TransactionManager manager = dbClient.transactionManager()) {
    TransactionContext txn = manager.begin();
    while (true) {
      try {
        String column = "FirstName";
        Struct row = txn.readRow("Singers", Key.of(singerId), Collections.singleton(column));
        String name = row.getString(column);
        txn.buffer(
            Mutation.newUpdateBuilder("Singers").set(column).to(name.toUpperCase()).build());
        manager.commit();
        break;
      } catch (AbortedException e) {
        Thread.sleep(e.getRetryDelayInMillis() / 1000);
        txn = manager.resetForRetry();
      }
    }
  }
  // [END transactionManager]
}
 
Example 7
@Test
public void dateTimestamp() {
  Ddl ddl = Ddl.builder()
      .createTable("users")
      .column("id").int64().notNull().endColumn()
      .column("date").date().endColumn()
      .column("ts").timestamp().endColumn()
      .primaryKey().asc("id").end()
      .endTable()
      .build();
  Schema schema = converter.convert(ddl).iterator().next();
  SpannerRecordConverter recordConverter = new SpannerRecordConverter(schema);
  Struct struct = Struct.newBuilder()
      .set("id").to(1L)
      .set("date").to(Date.fromYearMonthDay(2018, 2, 2))
      .set("ts").to(Timestamp.ofTimeMicroseconds(10))
      .build();

  GenericRecord avroRecord = recordConverter.convert(struct);

  assertThat(avroRecord.get("id"), equalTo(1L));
  assertThat(avroRecord.get("date"), equalTo("2018-02-02"));
  assertThat(avroRecord.get("ts"), equalTo("1970-01-01T00:00:00.000010000Z"));
}
 
Example 8
Source Project: spanner-jdbc   Source File: RunningOperationsStore.java    License: MIT License 6 votes vote down vote up
/**
 * 
 * @return A result set of all DDL operations that have been issued on this connection since the
 *         last clear operation.
 * @throws SQLException If a database error occurs
 */
ResultSet getOperations(CloudSpannerStatement statement) throws SQLException {
  List<Struct> rows = new ArrayList<>(operations.size());
  for (DdlOperation op : operations) {
    op.operation = op.operation.reload();
    String exception = null;
    try {
      op.operation.getResult();
    } catch (Exception e) {
      exception = e.getMessage();
    }
    for (String ddl : op.sql) {
      rows.add(Struct.newBuilder().set("NAME").to(Value.string(op.operation.getName()))
          .set("TIME_STARTED").to(Value.timestamp(op.timeStarted)).set("STATEMENT")
          .to(Value.string(ddl)).set("DONE").to(Value.bool(op.operation.isDone()))
          .set("EXCEPTION").to(Value.string(exception)).build());
    }
  }
  com.google.cloud.spanner.ResultSet rs =
      ResultSets.forRows(Type.struct(StructField.of("NAME", Type.string()),
          StructField.of("TIME_STARTED", Type.timestamp()),
          StructField.of("STATEMENT", Type.string()), StructField.of("DONE", Type.bool()),
          StructField.of("EXCEPTION", Type.string())), rows);
  return new CloudSpannerResultSet(statement, rs, null);
}
 
Example 9
Source Project: beam   Source File: SpannerIOWriteTest.java    License: Apache License 2.0 6 votes vote down vote up
private void preparePkMetadata(ReadOnlyTransaction tx, List<Struct> rows) {
  Type type =
      Type.struct(
          Type.StructField.of("table_name", Type.string()),
          Type.StructField.of("column_name", Type.string()),
          Type.StructField.of("column_ordering", Type.string()));
  when(tx.executeQuery(
          argThat(
              new ArgumentMatcher<Statement>() {

                @Override
                public boolean matches(Statement argument) {
                  if (!(argument instanceof Statement)) {
                    return false;
                  }
                  Statement st = (Statement) argument;
                  return st.getSql().contains("information_schema.index_columns");
                }
              })))
      .thenReturn(ResultSets.forRows(type, rows));
}
 
Example 10
/**
 * Creates a Cloud Spanner statement.
 * @param sql the SQL string with tags.
 * @param tags the tags that appear in the SQL string.
 * @param paramStructConvertFunc a function to use to convert params to {@link Struct}
 *     objects if they cannot be directly mapped to Cloud Spanner supported param types.
 *     If null then this last-attempt conversion is skipped.
 * @param spannerCustomConverter a converter used to convert params that aren't Cloud
 *     Spanner native types. if {@code null} then this conversion is not attempted.
 * @param params the parameters to substitute the tags. The ordering must be the same as
 *     the tags.
 * @param queryMethodParams the parameter metadata from Query Method if available.
 * @return an SQL statement ready to use with Spanner.
 * @throws IllegalArgumentException if the number of tags does not match the number of
 *     params, or if a param of an unsupported type is given.
 */
public static Statement buildStatementFromSqlWithArgs(String sql, List<String> tags,
		Function<Object, Struct> paramStructConvertFunc, SpannerCustomConverter spannerCustomConverter,
		Object[] params, Map<String, Parameter> queryMethodParams) {
	if (tags == null && params == null) {
		return Statement.of(sql);
	}
	if (tags == null || params == null || tags.size() != params.length) {
		throw new IllegalArgumentException(
				"The number of tags does not match the number of params.");
	}
	Statement.Builder builder = Statement.newBuilder(sql);
	for (int i = 0; i < tags.size(); i++) {
		bindParameter(builder.bind(tags.get(i)), paramStructConvertFunc, spannerCustomConverter,
				params[i], queryMethodParams == null ? null : queryMethodParams.get(tags.get(i)));
	}
	return builder.build();
}
 
Example 11
Source Project: beam   Source File: ReadSpannerSchemaTest.java    License: Apache License 2.0 6 votes vote down vote up
private void prepareColumnMetadata(ReadOnlyTransaction tx, List<Struct> rows) {
  Type type =
      Type.struct(
          Type.StructField.of("table_name", Type.string()),
          Type.StructField.of("column_name", Type.string()),
          Type.StructField.of("spanner_type", Type.string()),
          Type.StructField.of("cells_mutated", Type.int64()));
  when(tx.executeQuery(
          argThat(
              new ArgumentMatcher<Statement>() {

                @Override
                public boolean matches(Statement argument) {
                  if (!(argument instanceof Statement)) {
                    return false;
                  }
                  Statement st = (Statement) argument;
                  return st.getSql().contains("information_schema.columns");
                }
              })))
      .thenReturn(ResultSets.forRows(type, rows));
}
 
Example 12
Source Project: java-docs-samples   Source File: SpannerSample.java    License: Apache License 2.0 6 votes vote down vote up
static void updateUsingDmlWithStruct(DatabaseClient dbClient) {
  Struct name =
      Struct.newBuilder().set("FirstName").to("Timothy").set("LastName").to("Campbell").build();
  Statement s =
      Statement.newBuilder(
              "UPDATE Singers SET LastName = 'Grant' "
                  + "WHERE STRUCT<FirstName STRING, LastName STRING>(FirstName, LastName) "
                  + "= @name")
          .bind("name")
          .to(name)
          .build();
  dbClient
      .readWriteTransaction()
      .run(
          new TransactionCallable<Void>() {
            @Override
            public Void run(TransactionContext transaction) throws Exception {
              long rowCount = transaction.executeUpdate(s);
              System.out.printf("%d record updated.\n", rowCount);
              return null;
            }
          });
}
 
Example 13
Source Project: beam   Source File: SpannerIOWriteTest.java    License: Apache License 2.0 6 votes vote down vote up
private void prepareColumnMetadata(ReadOnlyTransaction tx, List<Struct> rows) {
  Type type =
      Type.struct(
          Type.StructField.of("table_name", Type.string()),
          Type.StructField.of("column_name", Type.string()),
          Type.StructField.of("spanner_type", Type.string()),
          Type.StructField.of("cells_mutated", Type.int64()));
  when(tx.executeQuery(
          argThat(
              new ArgumentMatcher<Statement>() {

                @Override
                public boolean matches(Statement argument) {
                  if (!(argument instanceof Statement)) {
                    return false;
                  }
                  Statement st = (Statement) argument;
                  return st.getSql().contains("information_schema.columns");
                }
              })))
      .thenReturn(ResultSets.forRows(type, rows));
}
 
Example 14
@Test
public void readNestedStructsAsStructsTest() {
	Struct innerStruct = Struct.newBuilder().set("value")
			.to(Value.string("inner-value")).build();
	Struct outerStruct = Struct.newBuilder().set("id").to(Value.string("key1"))
			.set("innerStructs")
			.toStructArray(Type.struct(StructField.of("value", Type.string())),
					Arrays.asList(innerStruct))
			.build();

	OuterTestHoldingStructsEntity result = this.spannerEntityReader
			.read(OuterTestHoldingStructsEntity.class, outerStruct);
	assertThat(result.id).isEqualTo("key1");
	assertThat(result.innerStructs).hasSize(1);
	assertThat(result.innerStructs.get(0).getString("value")).isEqualTo("inner-value");
}
 
Example 15
@Test
public void readSingularArrayMismatchTest() {
	this.expectedEx.expect(SpannerDataException.class);
	this.expectedEx.expectMessage("The value in column with name innerLengths could not be converted " +
			"to the corresponding property in the entity. The property's type is class java.lang.Integer.");

	Struct colStruct = Struct.newBuilder().set("string_col").to(Value.string("value"))
			.build();
	Struct rowStruct = Struct.newBuilder().set("id").to(Value.string("key1"))
			.set("innerLengths")
			.toStructArray(Type.struct(StructField.of("string_col", Type.string())),
					Arrays.asList(colStruct))
			.build();

	new ConverterAwareMappingSpannerEntityReader(new SpannerMappingContext(),
			new SpannerReadConverter(Arrays.asList(new Converter<Struct, Integer>() {
				@Nullable
				@Override
				public Integer convert(Struct source) {
					return source.getString("string_col").length();
				}
			}))).read(OuterTestEntityFlatFaulty.class, rowStruct);
}
 
Example 16
@Test
public void readConvertedNestedStructTest() {
	Struct colStruct = Struct.newBuilder().set("string_col").to(Value.string("value"))
			.build();
	Struct rowStruct = Struct.newBuilder().set("id").to(Value.string("key1"))
			.set("innerLengths")
			.toStructArray(Type.struct(StructField.of("string_col", Type.string())),
					Arrays.asList(colStruct))
			.build();

	OuterTestEntityFlat result = new ConverterAwareMappingSpannerEntityReader(
			new SpannerMappingContext(),
			new SpannerReadConverter(Arrays.asList(new Converter<Struct, Integer>() {
				@Nullable
				@Override
				public Integer convert(Struct source) {
					return source.getString("string_col").length();
				}
			}))).read(OuterTestEntityFlat.class, rowStruct);
	assertThat(result.id).isEqualTo("key1");
	assertThat(result.innerLengths).hasSize(1);
	assertThat(result.innerLengths.get(0)).isEqualTo(5);
}
 
Example 17
Source Project: google-cloud-java   Source File: DatabaseClientSnippets.java    License: Apache License 2.0 6 votes vote down vote up
/** Example of a read write transaction. */
// [TARGET readWriteTransaction()]
// [VARIABLE my_singer_id]
public void readWriteTransaction(final long singerId) {
  // [START readWriteTransaction]
  TransactionRunner runner = dbClient.readWriteTransaction();
  runner.run(
      new TransactionCallable<Void>() {

        @Override
        public Void run(TransactionContext transaction) throws Exception {
          String column = "FirstName";
          Struct row =
              transaction.readRow("Singers", Key.of(singerId), Collections.singleton(column));
          String name = row.getString(column);
          transaction.buffer(
              Mutation.newUpdateBuilder("Singers").set(column).to(name.toUpperCase()).build());
          return null;
        }
      });
  // [END readWriteTransaction]
}
 
Example 18
@Test
public void readNestedStructWithConstructor() {
	Struct innerStruct = Struct.newBuilder().set("value").to(Value.string("value"))
			.build();
	Struct outerStruct = Struct.newBuilder().set("id").to(Value.string("key1"))
			.set("innerTestEntities")
			.toStructArray(Type.struct(StructField.of("value", Type.string())),
					Arrays.asList(innerStruct))
			.build();

	TestEntities.OuterTestEntityWithConstructor result = this.spannerEntityReader
			.read(TestEntities.OuterTestEntityWithConstructor.class, outerStruct, null, true);
	assertThat(result.id).isEqualTo("key1");
	assertThat(result.innerTestEntities).hasSize(1);
	assertThat(result.innerTestEntities.get(0).value).isEqualTo("value");
}
 
Example 19
Source Project: DataflowTemplates   Source File: LocalSpannerIO.java    License: Apache License 2.0 5 votes vote down vote up
@Override
public PCollection<Struct> expand(PCollection<ReadOperation> input) {
  PTransform<PCollection<ReadOperation>, PCollection<Struct>> readTransform;
  if (getBatching()) {
    readTransform =
        LocalBatchSpannerRead.create(getSpannerConfig(), getTransaction(), getTimestampBound());
  } else {
    readTransform =
        NaiveSpannerRead.create(getSpannerConfig(), getTransaction(), getTimestampBound());
  }
  return input
      .apply("Reshuffle", Reshuffle.viaRandomKey())
      .apply("Read from Cloud Spanner", readTransform);
}
 
Example 20
Source Project: DataflowTemplates   Source File: LocalSpannerIO.java    License: Apache License 2.0 5 votes vote down vote up
@Override
public PCollection<Struct> expand(PBegin input) {
  getSpannerConfig().validate();
  checkArgument(
      getTimestampBound() != null,
      "SpannerIO.read() runs in a read only transaction and requires timestamp to be set "
          + "with withTimestampBound or withTimestamp method");

  if (getReadOperation().getQuery() != null) {
    // TODO: validate query?
  } else if (getReadOperation().getTable() != null) {
    // Assume read
    checkNotNull(
        getReadOperation().getColumns(),
        "For a read operation SpannerIO.read() requires a list of "
            + "columns to set with withColumns method");
    checkArgument(
        !getReadOperation().getColumns().isEmpty(),
        "For a read operation SpannerIO.read() requires a"
            + " list of columns to set with withColumns method");
  } else {
    throw new IllegalArgumentException(
        "SpannerIO.read() requires configuring query or read operation.");
  }

  ReadAll readAll =
      readAll()
          .withSpannerConfig(getSpannerConfig())
          .withTimestampBound(getTimestampBound())
          .withBatching(getBatching())
          .withTransaction(getTransaction());
  return input.apply(Create.of(getReadOperation())).apply("Execute query", readAll);
}
 
Example 21
Source Project: google-cloud-java   Source File: DatabaseClientSnippets.java    License: Apache License 2.0 5 votes vote down vote up
/** Example of single use read only transaction with timestamp bound. */
// [TARGET singleUseReadOnlyTransaction(TimestampBound)]
// [VARIABLE my_singer_id]
public Timestamp singleUseReadOnlyTransactionTimestamp(long singerId) {
  // [START singleUseReadOnlyTransactionTimestamp]
  String column = "FirstName";
  ReadOnlyTransaction txn =
      dbClient.singleUseReadOnlyTransaction(TimestampBound.ofMaxStaleness(10, TimeUnit.SECONDS));
  Struct row = txn.readRow("Singers", Key.of(singerId), Collections.singleton(column));
  row.getString(column);
  Timestamp timestamp = txn.getReadTimestamp();
  // [END singleUseReadOnlyTransactionTimestamp]
  return timestamp;
}
 
Example 22
Source Project: DataflowTemplates   Source File: LocalBatchSpannerRead.java    License: Apache License 2.0 5 votes vote down vote up
@ProcessElement
public void processElement(ProcessContext c) throws Exception {
  Transaction tx = c.sideInput(txView);

  BatchReadOnlyTransaction batchTx =
      spannerAccessor.getBatchClient().batchReadOnlyTransaction(tx.transactionId());

  Partition p = c.element();
  try (ResultSet resultSet = batchTx.execute(p)) {
    while (resultSet.next()) {
      Struct s = resultSet.getCurrentRowAsStruct();
      c.output(s);
    }
  }
}
 
Example 23
Source Project: DataflowTemplates   Source File: SpannerConverters.java    License: Apache License 2.0 5 votes vote down vote up
/**
 * Parse results given by table read. The function returns a list of strings, where each string
 * represents a single column. The list constitutes the entire result set (columns in all of the
 * records).
 */
private static List<String> parseResultSet(
    Struct struct, LinkedHashMap<String, BiFunction<Struct, String, String>> parsers) {
  List<String> result = Lists.newArrayList();

  for (String columnName : parsers.keySet()) {
    if (!parsers.containsKey(columnName)) {
      throw new RuntimeException("No parser for column: " + columnName);
    }
    result.add(parsers.get(columnName).apply(struct, columnName));
  }
  return result;
}
 
Example 24
Source Project: java-docs-samples   Source File: SpannerSample.java    License: Apache License 2.0 5 votes vote down vote up
static void queryWithArrayOfStruct(DatabaseClient dbClient) {
  // [START spanner_create_user_defined_struct]
  Type nameType =
      Type.struct(
          Arrays.asList(
              StructField.of("FirstName", Type.string()),
              StructField.of("LastName", Type.string())));
  // [END spanner_create_user_defined_struct]

  // [START spanner_create_array_of_struct_with_data]
  List<Struct> bandMembers = new ArrayList<>();
  bandMembers.add(
      Struct.newBuilder().set("FirstName").to("Elena").set("LastName").to("Campbell").build());
  bandMembers.add(
      Struct.newBuilder().set("FirstName").to("Gabriel").set("LastName").to("Wright").build());
  bandMembers.add(
      Struct.newBuilder().set("FirstName").to("Benjamin").set("LastName").to("Martinez").build());
  // [END spanner_create_array_of_struct_with_data]

  // [START spanner_query_data_with_array_of_struct]
  Statement s =
      Statement.newBuilder(
              "SELECT SingerId FROM Singers WHERE "
                  + "STRUCT<FirstName STRING, LastName STRING>(FirstName, LastName) "
                  + "IN UNNEST(@names) "
                  + "ORDER BY SingerId DESC")
          .bind("names")
          .toStructArray(nameType, bandMembers)
          .build();
  try (ResultSet resultSet = dbClient.singleUse().executeQuery(s)) {
    while (resultSet.next()) {
      System.out.printf("%d\n", resultSet.getLong("SingerId"));
    }
  }
  // [END spanner_query_data_with_array_of_struct]
}
 
Example 25
Source Project: google-cloud-java   Source File: DatabaseClientSnippets.java    License: Apache License 2.0 5 votes vote down vote up
/** Example of single use. */
// [TARGET singleUse()]
// [VARIABLE my_singer_id]
public String singleUse(long singerId) {
  // [START singleUse]
  String column = "FirstName";
  Struct row =
      dbClient.singleUse().readRow("Singers", Key.of(singerId), Collections.singleton(column));
  String firstName = row.getString(column);
  // [END singleUse]
  return firstName;
}
 
Example 26
@Test
public void testPartialConstructorWithNotEnoughArgs() {
	this.expectedEx.expect(SpannerDataException.class);
	this.expectedEx.expectMessage("Column not found: custom_col");
	Struct struct = Struct.newBuilder().set("id").to(Value.string("key1"))
			.set("booleanField").to(Value.bool(true)).set("longField")
			.to(Value.int64(3L)).set("doubleField").to(Value.float64(3.14)).build();

	this.spannerEntityReader.read(TestEntities.PartialConstructor.class, struct);
}
 
Example 27
Source Project: DataflowTemplates   Source File: SpannerConverterTest.java    License: Apache License 2.0 5 votes vote down vote up
/** Unit test for export transform. */
@Test
@Category(NeedsRunner.class)
public void testSchemaSave() throws IOException {

  ValueProvider<String> table = ValueProvider.StaticValueProvider.of(TABLE);
  SpannerConfig spannerConfig = SpannerConfig.create();
  DatabaseClient databaseClient = mock(DatabaseClient.class, withSettings().serializable());
  ReadOnlyTransaction readOnlyTransaction =
      mock(ReadOnlyTransaction.class, withSettings().serializable());
  ResultSet resultSet = mock(ResultSet.class, withSettings().serializable());
  Struct struct = mock(Struct.class, withSettings().serializable());

  when(databaseClient.readOnlyTransaction()).thenReturn(readOnlyTransaction);
  when(readOnlyTransaction.executeQuery(any(Statement.class))).thenReturn(resultSet);
  when(resultSet.next()).thenReturn(true).thenReturn(false);
  when(resultSet.getCurrentRowAsStruct()).thenReturn(struct);
  when(struct.getString(0)).thenReturn(COLUMN_NAME);
  when(struct.getString(1)).thenReturn("INT64");

  String schemaPath = "/tmp/" + UUID.randomUUID().toString();
  ValueProvider<String> textWritePrefix = ValueProvider.StaticValueProvider.of(schemaPath);
  SpannerConverters.ExportTransform exportTransform =
      SpannerConverters.ExportTransformFactory.create(table, spannerConfig, textWritePrefix);
  exportTransform.setDatabaseClient(databaseClient);

  PCollection<ReadOperation> results = pipeline.apply("Create", exportTransform);
  ReadOperation readOperation =
      ReadOperation.create().withColumns(ImmutableList.of(COLUMN_NAME)).withTable(TABLE);
  PAssert.that(results).containsInAnyOrder(readOperation);
  pipeline.run();
  ReadableByteChannel channel =
      FileSystems.open(
          FileSystems.matchNewResource(
              schemaPath + SpannerConverters.ExportTransform.ExportFn.SCHEMA_SUFFIX,
              false));
  java.util.Scanner scanner = new java.util.Scanner(channel).useDelimiter("\\A");
  assertEquals("{\"id\":\"INT64\"}", scanner.next());
}
 
Example 28
Source Project: DataflowTemplates   Source File: SpannerConverterTest.java    License: Apache License 2.0 5 votes vote down vote up
@Test
public void testBytes() {
  assertEquals(
      "\"dGVzdA==\"",
      structCsvPrinter.print(
          Struct.newBuilder().set("col").to(Value.bytes(ByteArray.copyFrom("test"))).build()));
}
 
Example 29
Source Project: DataflowTemplates   Source File: SpannerConverterTest.java    License: Apache License 2.0 5 votes vote down vote up
@Test
public void testTimestamp() {
  assertEquals(
      "\"1970-01-01T00:00:00Z\"",
      structCsvPrinter.print(
          Struct.newBuilder()
              .set("col").to(Value.timestamp(Timestamp.ofTimeMicroseconds(0)))
              .build()));
}
 
Example 30
Source Project: java-docs-samples   Source File: SpannerReadApi.java    License: Apache License 2.0 5 votes vote down vote up
public static void main(String[] args) {
  Options options = PipelineOptionsFactory.fromArgs(args).withValidation().as(Options.class);
  Pipeline p = Pipeline.create(options);

  String instanceId = options.getInstanceId();
  String databaseId = options.getDatabaseId();
  // [START spanner_dataflow_readapi]
  // Query for all the columns and rows in the specified Spanner table
  PCollection<Struct> records = p.apply(
      SpannerIO.read()
          .withInstanceId(instanceId)
          .withDatabaseId(databaseId)
          .withTable("Singers")
          .withColumns("singerId", "firstName", "lastName"));
  // [END spanner_dataflow_readapi]


  PCollection<Long> tableEstimatedSize = records
      // Estimate the size of every row
      .apply(EstimateSize.create())
      // Sum all the row sizes to get the total estimated size of the table
      .apply(Sum.longsGlobally());

  // Write the total size to a file
  tableEstimatedSize
      .apply(ToString.elements())
      .apply(TextIO.write().to(options.getOutput()).withoutSharding());

  p.run().waitUntilFinish();
}