com.example.appengine.spanner.SpannerTasks.java Source code

Java tutorial

Introduction

Here is the source code for com.example.appengine.spanner.SpannerTasks.java

Source

/*
 * Copyright 2017 Google Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.example.appengine.spanner;

import com.google.cloud.spanner.Database;
import com.google.cloud.spanner.DatabaseClient;
import com.google.cloud.spanner.Key;
import com.google.cloud.spanner.KeySet;
import com.google.cloud.spanner.Mutation;
import com.google.cloud.spanner.Operation;
import com.google.cloud.spanner.ReadOnlyTransaction;
import com.google.cloud.spanner.ResultSet;
import com.google.cloud.spanner.Statement;
import com.google.cloud.spanner.Struct;
import com.google.common.base.Stopwatch;
import com.google.spanner.admin.database.v1.UpdateDatabaseDdlMetadata;

import java.io.PrintWriter;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.concurrent.TimeUnit;

class SpannerTasks {

    enum Task {
        createDatabase, writeExampleData, query, read, addMarketingBudget, updateMarketingBudget, queryMarketingBudget, addIndex, readUsingIndex, queryUsingIndex, addStoringIndex, readStoringIndex, writeTransaction, readOnlyTransaction
    }

    /** Class to contain singer sample data. */
    static class Singer {

        final long singerId;
        final String firstName;
        final String lastName;

        Singer(long singerId, String firstName, String lastName) {
            this.singerId = singerId;
            this.firstName = firstName;
            this.lastName = lastName;
        }
    }

    /** Class to contain album sample data. */
    static class Album {

        final long singerId;
        final long albumId;
        final String albumTitle;

        Album(long singerId, long albumId, String albumTitle) {
            this.singerId = singerId;
            this.albumId = albumId;
            this.albumTitle = albumTitle;
        }
    }

    private static final List<Singer> SINGERS = Arrays.asList(new Singer(1, "Marc", "Richards"),
            new Singer(2, "Catalina", "Smith"), new Singer(3, "Alice", "Trentor"), new Singer(4, "Lea", "Martin"),
            new Singer(5, "David", "Lomond"));

    private static final List<Album> ALBUMS = Arrays.asList(new Album(1, 1, "Total Junk"),
            new Album(1, 2, "Go, Go, Go"), new Album(2, 1, "Green"), new Album(2, 2, "Forever Hold Your Peace"),
            new Album(2, 3, "Terrified"));

    private static DatabaseClient databaseClient = null;

    private static void createDatabase(PrintWriter pw) {
        Iterable<String> statements = Arrays.asList(
                "CREATE TABLE Singers (\n" + "  SingerId   INT64 NOT NULL,\n" + "  FirstName  STRING(1024),\n"
                        + "  LastName   STRING(1024),\n" + "  SingerInfo BYTES(MAX)\n" + ") PRIMARY KEY (SingerId)",
                "CREATE TABLE Albums (\n" + "  SingerId     INT64 NOT NULL,\n" + "  AlbumId      INT64 NOT NULL,\n"
                        + "  AlbumTitle   STRING(MAX)\n" + ") PRIMARY KEY (SingerId, AlbumId),\n"
                        + "  INTERLEAVE IN PARENT Singers ON DELETE CASCADE");
        Database db = SpannerClient.getDatabaseAdminClient()
                .createDatabase(SpannerClient.getInstanceId(), SpannerClient.getDatabaseId(), statements).waitFor()
                .getResult();
        pw.println("Created database [" + db.getId() + "]");
    }

    private static void writeExampleData(PrintWriter pw) {
        List<Mutation> mutations = new ArrayList<>();
        for (Singer singer : SINGERS) {
            mutations.add(Mutation.newInsertBuilder("Singers").set("SingerId").to(singer.singerId).set("FirstName")
                    .to(singer.firstName).set("LastName").to(singer.lastName).build());
        }
        for (Album album : ALBUMS) {
            mutations.add(Mutation.newInsertBuilder("Albums").set("SingerId").to(album.singerId).set("AlbumId")
                    .to(album.albumId).set("AlbumTitle").to(album.albumTitle).build());
        }
        SpannerClient.getDatabaseClient().write(mutations);
    }

    private static void query(PrintWriter pw) {
        // singleUse() can be used to execute a single read or query against Cloud Spanner.
        ResultSet resultSet = SpannerClient.getDatabaseClient().singleUse()
                .executeQuery(Statement.of("SELECT SingerId, AlbumId, AlbumTitle FROM Albums"));
        while (resultSet.next()) {
            pw.printf("%d %d %s\n", resultSet.getLong(0), resultSet.getLong(1), resultSet.getString(2));
        }
    }

    private static void read(PrintWriter pw) {
        ResultSet resultSet = SpannerClient.getDatabaseClient().singleUse().read("Albums",
                // KeySet.all() can be used to read all rows in a table. KeySet exposes other
                // methods to read only a subset of the table.
                KeySet.all(), Arrays.asList("SingerId", "AlbumId", "AlbumTitle"));
        while (resultSet.next()) {
            pw.printf("%d %d %s\n", resultSet.getLong(0), resultSet.getLong(1), resultSet.getString(2));
        }
    }

    private static void addMarketingBudgetColumnToAlbums(PrintWriter pw) {
        Operation<Void, UpdateDatabaseDdlMetadata> op = SpannerClient.getDatabaseAdminClient().updateDatabaseDdl(
                SpannerClient.getInstanceId(), SpannerClient.getDatabaseId(),
                Arrays.asList("ALTER TABLE Albums ADD COLUMN MarketingBudget INT64"), null);
        op.waitFor();
    }

    // Before executing this method, a new column MarketingBudget has to be added to the Albums
    // table by applying the DDL statement "ALTER TABLE Albums ADD COLUMN MarketingBudget INT64".
    private static void updateMarketingBudgetData() {
        // Mutation can be used to update/insert/delete a single row in a table. Here we use
        // newUpdateBuilder to create update mutations.
        List<Mutation> mutations = Arrays.asList(
                Mutation.newUpdateBuilder("Albums").set("SingerId").to(1).set("AlbumId").to(1)
                        .set("MarketingBudget").to(100000).build(),
                Mutation.newUpdateBuilder("Albums").set("SingerId").to(2).set("AlbumId").to(2)
                        .set("MarketingBudget").to(500000).build());
        // This writes all the mutations to Cloud Spanner atomically.
        SpannerClient.getDatabaseClient().write(mutations);
    }

    private static void writeWithTransaction() {
        SpannerClient.getDatabaseClient().readWriteTransaction().run((transactionContext -> {
            // Transfer marketing budget from one album to another. We do it in a transaction to
            // ensure that the transfer is atomic.
            Struct row = transactionContext.readRow("Albums", Key.of(2, 2), Arrays.asList("MarketingBudget"));
            long album2Budget = row.getLong(0);
            // Transaction will only be committed if this condition still holds at the time of
            // commit. Otherwise it will be aborted and the callable will be rerun by the
            // client library.
            if (album2Budget >= 300000) {
                long album1Budget = transactionContext
                        .readRow("Albums", Key.of(1, 1), Arrays.asList("MarketingBudget")).getLong(0);
                long transfer = 200000;
                album1Budget += transfer;
                album2Budget -= transfer;
                transactionContext.buffer(Mutation.newUpdateBuilder("Albums").set("SingerId").to(1).set("AlbumId")
                        .to(1).set("MarketingBudget").to(album1Budget).build());
                transactionContext.buffer(Mutation.newUpdateBuilder("Albums").set("SingerId").to(2).set("AlbumId")
                        .to(2).set("MarketingBudget").to(album2Budget).build());
            }
            return null;
        }));
    }

    private static void queryMarketingBudget(PrintWriter pw) {
        // Rows without an explicit value for MarketingBudget will have a MarketingBudget equal to
        // null.
        ResultSet resultSet = SpannerClient.getDatabaseClient().singleUse()
                .executeQuery(Statement.of("SELECT SingerId, AlbumId, MarketingBudget FROM Albums"));
        while (resultSet.next()) {
            pw.printf("%d %d %s\n", resultSet.getLong("SingerId"), resultSet.getLong("AlbumId"),
                    // We check that the value is non null. ResultSet getters can only be used to retrieve
                    // non null values.
                    resultSet.isNull("MarketingBudget") ? "NULL" : resultSet.getLong("MarketingBudget"));
        }
    }

    private static void addIndex() {
        Operation<Void, UpdateDatabaseDdlMetadata> op = SpannerClient.getDatabaseAdminClient().updateDatabaseDdl(
                SpannerClient.getInstanceId(), SpannerClient.getDatabaseId(),
                Arrays.asList("CREATE INDEX AlbumsByAlbumTitle ON Albums(AlbumTitle)"), null);
        op.waitFor();
    }

    // Before running this example, add the index AlbumsByAlbumTitle by applying the DDL statement
    // "CREATE INDEX AlbumsByAlbumTitle ON Albums(AlbumTitle)".
    private static void queryUsingIndex(PrintWriter pw) {
        ResultSet resultSet = SpannerClient.getDatabaseClient().singleUse().executeQuery(
                // We use FORCE_INDEX hint to specify which index to use. For more details see
                // https://cloud.google.com/spanner/docs/query-syntax#from-clause
                Statement.of("SELECT AlbumId, AlbumTitle, MarketingBudget\n"
                        + "FROM Albums@{FORCE_INDEX=AlbumsByAlbumTitle}\n"
                        + "WHERE AlbumTitle >= 'Aardvark' AND AlbumTitle < 'Goo'"));
        while (resultSet.next()) {
            pw.printf("%d %s %s\n", resultSet.getLong("AlbumId"), resultSet.getString("AlbumTitle"),
                    resultSet.isNull("MarketingBudget") ? "NULL" : resultSet.getLong("MarketingBudget"));
        }
    }

    private static void readUsingIndex(PrintWriter pw) {
        ResultSet resultSet = SpannerClient.getDatabaseClient().singleUse().readUsingIndex("Albums",
                "AlbumsByAlbumTitle", KeySet.all(), Arrays.asList("AlbumId", "AlbumTitle"));
        while (resultSet.next()) {
            pw.printf("%d %s\n", resultSet.getLong(0), resultSet.getString(1));
        }
    }

    private static void addStoringIndex() {
        Operation<Void, UpdateDatabaseDdlMetadata> op = SpannerClient.getDatabaseAdminClient()
                .updateDatabaseDdl(SpannerClient.getInstanceId(), SpannerClient.getDatabaseId(), Arrays.asList(
                        "CREATE INDEX AlbumsByAlbumTitle2 " + "ON Albums(AlbumTitle) STORING (MarketingBudget)"),
                        null);
        op.waitFor();
    }

    // Before running this example, create a storing index AlbumsByAlbumTitle2 by applying the DDL
    // statement "CREATE INDEX AlbumsByAlbumTitle2 ON Albums(AlbumTitle) STORING (MarketingBudget)".
    private static void readStoringIndex(PrintWriter pw) {
        // We can read MarketingBudget also from the index since it stores a copy of MarketingBudget.
        ResultSet resultSet = SpannerClient.getDatabaseClient().singleUse().readUsingIndex("Albums",
                "AlbumsByAlbumTitle2", KeySet.all(), Arrays.asList("AlbumId", "AlbumTitle", "MarketingBudget"));
        while (resultSet.next()) {
            pw.printf("%d %s %s\n", resultSet.getLong(0), resultSet.getString(1),
                    resultSet.isNull("MarketingBudget") ? "NULL" : resultSet.getLong("MarketingBudget"));
        }
    }

    private static void readOnlyTransaction(PrintWriter pw) {
        // ReadOnlyTransaction must be closed by calling close() on it to release resources held by it.
        // We use a try-with-resource block to automatically do so.
        try (ReadOnlyTransaction transaction = SpannerClient.getDatabaseClient().readOnlyTransaction()) {
            ResultSet queryResultSet = transaction
                    .executeQuery(Statement.of("SELECT SingerId, AlbumId, AlbumTitle FROM Albums"));
            while (queryResultSet.next()) {
                pw.printf("%d %d %s\n", queryResultSet.getLong(0), queryResultSet.getLong(1),
                        queryResultSet.getString(2));
            }
            ResultSet readResultSet = transaction.read("Albums", KeySet.all(),
                    Arrays.asList("SingerId", "AlbumId", "AlbumTitle"));
            while (readResultSet.next()) {
                pw.printf("%d %d %s\n", readResultSet.getLong(0), readResultSet.getLong(1),
                        readResultSet.getString(2));
            }
        }
    }

    static void runTask(Task task, PrintWriter pw) {
        Stopwatch stopwatch = Stopwatch.createStarted();
        switch (task) {
        case createDatabase:
            createDatabase(pw);
            break;
        case writeExampleData:
            writeExampleData(pw);
            break;
        case query:
            query(pw);
            break;
        case read:
            read(pw);
            break;
        case addMarketingBudget:
            addMarketingBudgetColumnToAlbums(pw);
            break;
        case updateMarketingBudget:
            updateMarketingBudgetData();
            break;
        case queryMarketingBudget:
            queryMarketingBudget(pw);
            break;
        case addIndex:
            addIndex();
            break;
        case readUsingIndex:
            readUsingIndex(pw);
            break;
        case queryUsingIndex:
            queryUsingIndex(pw);
            break;
        case addStoringIndex:
            addStoringIndex();
            break;
        case readStoringIndex:
            readStoringIndex(pw);
            break;
        case readOnlyTransaction:
            readOnlyTransaction(pw);
            break;
        case writeTransaction:
            writeWithTransaction();
            break;
        default:
            break;
        }
        stopwatch.stop();
        pw.println(task + " in milliseconds : " + stopwatch.elapsed(TimeUnit.MILLISECONDS));
        pw.println("====================================================================");
    }
}