Makes CASSANDRA_ENSURE_SCHEMA work when schema is absent (#1128)

Our cassandra docker image comes with the schema installed, which hid a
bug where we failed to properly guard state 0 (keyspace didn't exist).

This fixes the bug and also backfills known upgrade scenarios. Tested
on Cassandra 2.1 and 2.2.

Thanks to @gena01 for reporting

zipkin-storage/cassandra/src/main/java/zipkin/storage/cassandra/Schema.java

@@ -63,10 +63,8 @@ static final class Metadata {
   }
 
   static KeyspaceMetadata getKeyspaceMetadata(Session session) {
-    return getKeyspaceMetadata(session.getLoggedKeyspace(), session.getCluster());
-  }
-
-  private static KeyspaceMetadata getKeyspaceMetadata(String keyspace, Cluster cluster) {
+    String keyspace = session.getLoggedKeyspace();
+    Cluster cluster = session.getCluster();
     KeyspaceMetadata keyspaceMetadata = cluster.getMetadata().getKeyspace(keyspace);
 
     if (keyspaceMetadata == null) {
@@ -78,10 +76,12 @@ private static KeyspaceMetadata getKeyspaceMetadata(String keyspace, Cluster clu
   }
 
   static void ensureExists(String keyspace, Session session) {
-    KeyspaceMetadata keyspaceMetadata = getKeyspaceMetadata(keyspace, session.getCluster());
-    if (keyspaceMetadata.getTable("traces") == null) {
+    KeyspaceMetadata keyspaceMetadata = session.getCluster().getMetadata().getKeyspace(keyspace);
+    if (keyspaceMetadata == null || keyspaceMetadata.getTable("traces") == null) {
       LOG.info("Installing schema {}", SCHEMA);
       applyCqlFile(keyspace, session, SCHEMA);
+      // refresh metadata since we've installed the schema
+      keyspaceMetadata = session.getCluster().getMetadata().getKeyspace(keyspace);
     }
     if (!hasUpgrade1_defaultTtl(keyspaceMetadata)) {
       LOG.info("Upgrading schema {}", SCHEMA);

zipkin-storage/cassandra/src/test/java/zipkin/storage/cassandra/EnsureSchemaTest.java

@@ -0,0 +1,94 @@
+/**
+ * Copyright 2015-2016 The OpenZipkin Authors
+ *
+ * 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 zipkin.storage.cassandra;
+
+import com.datastax.driver.core.Cluster;
+import com.datastax.driver.core.KeyspaceMetadata;
+import com.datastax.driver.core.Session;
+import com.google.common.io.Closer;
+import java.io.IOException;
+import org.junit.After;
+import org.junit.AssumptionViolatedException;
+import org.junit.Before;
+import org.junit.BeforeClass;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TestName;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+public class EnsureSchemaTest {
+
+  @Rule
+  public TestName name = new TestName();
+
+  @BeforeClass public static void checkCassandraIsUp() {
+    try (Cluster cluster = new SessionFactory.Default().buildCluster(
+        CassandraStorage.builder().build());
+         Session session = cluster.newSession()) {
+      session.execute("SELECT now() FROM system.local");
+    } catch (RuntimeException e) {
+      throw new AssumptionViolatedException(e.getMessage(), e);
+    }
+  }
+
+  Closer closer = Closer.create();
+  String keyspace;
+  Cluster cluster;
+  Session session;
+
+  @Before
+  public void connectAndDropKeyspace() {
+    keyspace = name.getMethodName().toLowerCase();
+    cluster = closer.register(new SessionFactory.Default()
+        .buildCluster(CassandraStorage.builder().keyspace(keyspace).build()));
+    session = closer.register(cluster.newSession());
+    session.execute("DROP KEYSPACE IF EXISTS " + keyspace);
+    assertThat(session.getCluster().getMetadata().getKeyspace(keyspace)).isNull();
+  }
+
+  @After
+  public void close() throws IOException {
+    closer.close();
+  }
+
+  @Test public void installsKeyspaceWhenMissing() {
+    Schema.ensureExists(keyspace, session);
+
+    KeyspaceMetadata metadata = session.getCluster().getMetadata().getKeyspace(keyspace);
+    assertThat(metadata).isNotNull();
+    assertThat(Schema.hasUpgrade1_defaultTtl(metadata)).isTrue();
+  }
+
+  @Test public void installsTablesWhenMissing() {
+    session.execute("CREATE KEYSPACE " + keyspace
+        + " WITH replication = {'class': 'SimpleStrategy', 'replication_factor': '1'};");
+
+    Schema.ensureExists(keyspace, session);
+
+    KeyspaceMetadata metadata = session.getCluster().getMetadata().getKeyspace(keyspace);
+    assertThat(metadata).isNotNull();
+    assertThat(Schema.hasUpgrade1_defaultTtl(metadata)).isTrue();
+  }
+
+  @Test public void upgradesOldSchema() {
+    Schema.applyCqlFile(keyspace, session, "/cassandra-schema-cql3-original.txt");
+
+    Schema.ensureExists(keyspace, session);
+
+    KeyspaceMetadata metadata = session.getCluster().getMetadata().getKeyspace(keyspace);
+    assertThat(metadata).isNotNull();
+    assertThat(Schema.hasUpgrade1_defaultTtl(metadata)).isTrue();
+  }
+}