Adds MySQL support for Autocomplete tags

zipkin-autoconfigure/storage-mysql/src/main/java/zipkin2/autoconfigure/storage/mysql/ZipkinMySQLStorageAutoConfiguration.java

@@ -13,6 +13,7 @@
  */
 package zipkin2.autoconfigure.storage.mysql;
 
+import java.util.List;
 import java.util.concurrent.Executor;
 import javax.sql.DataSource;
 import org.jooq.ExecuteListenerProvider;
@@ -57,14 +58,18 @@ DataSource mysqlDataSource() {
 
   @Bean
   StorageComponent storage(
-      Executor executor,
-      DataSource dataSource,
-      @Value("${zipkin.storage.strict-trace-id:true}") boolean strictTraceId) {
+    Executor executor,
+    DataSource dataSource,
+    @Value("${zipkin.storage.strict-trace-id:true}") boolean strictTraceId,
+    @Value("${zipkin.storage.search-enabled:true}") boolean searchEnabled,
+    @Value("${zipkin.storage.autocomplete-keys:}") List<String> autocompleteKeys) {
     return MySQLStorage.newBuilder()
-        .strictTraceId(strictTraceId)
-        .executor(executor)
-        .datasource(dataSource)
-        .listenerProvider(listener)
-        .build();
+      .strictTraceId(strictTraceId)
+      .searchEnabled(searchEnabled)
+      .autocompleteKeys(autocompleteKeys)
+      .executor(executor)
+      .datasource(dataSource)
+      .listenerProvider(listener)
+      .build();
   }
 }

zipkin-autoconfigure/storage-mysql/src/test/java/zipkin2/storage/mysql/v1/ZipkinMySQLStorageAutoConfigurationTest.java

@@ -94,6 +94,33 @@ public void strictTraceId_canSetToFalse() {
     assertThat(context.getBean(MySQLStorage.class).strictTraceId).isFalse();
   }
 
+  @Test
+  public void searchEnabled_canSetToFalse() {
+    context = new AnnotationConfigApplicationContext();
+    TestPropertyValues.of(
+      "zipkin.storage.type:mysql",
+      "zipkin.storage.search-enabled:false")
+      .applyTo(context);
+    Access.registerMySQL(context);
+    context.refresh();
+
+    assertThat(context.getBean(MySQLStorage.class).searchEnabled).isFalse();
+  }
+
+  @Test
+  public void autocompleteKeys_list() {
+    context = new AnnotationConfigApplicationContext();
+    TestPropertyValues.of(
+      "zipkin.storage.type:mysql",
+      "zipkin.storage.autocomplete-keys:environment")
+      .applyTo(context);
+    Access.registerMySQL(context);
+    context.refresh();
+
+    assertThat(context.getBean(MySQLStorage.class).autocompleteKeys)
+      .containsOnly("environment");
+  }
+
   @Test
   public void usesJdbcUrl_whenPresent() {
     context = new AnnotationConfigApplicationContext();

zipkin-storage/mysql-v1/src/main/java/zipkin2/storage/mysql/v1/MySQLAutocompleteTags.java

@@ -0,0 +1,47 @@
+/*
+ * Copyright 2015-2018 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 zipkin2.storage.mysql.v1;
+
+import java.util.LinkedHashSet;
+import java.util.List;
+import zipkin2.Call;
+import zipkin2.storage.AutocompleteTags;
+
+final class MySQLAutocompleteTags implements AutocompleteTags {
+  final DataSourceCall.Factory dataSourceCallFactory;
+  final Schema schema;
+  final boolean enabled;
+  final LinkedHashSet<String> autocompleteKeys;
+  final Call<List<String>> keysCall;
+
+  MySQLAutocompleteTags(MySQLStorage storage, Schema schema) {
+    this.dataSourceCallFactory = storage.dataSourceCallFactory;
+    this.schema = schema;
+    enabled = storage.searchEnabled && !storage.autocompleteKeys.isEmpty();
+    autocompleteKeys = new LinkedHashSet<>(storage.autocompleteKeys);
+    keysCall = Call.create(storage.autocompleteKeys);
+  }
+
+  @Override public Call<List<String>> getKeys() {
+    if (!enabled) return Call.emptyList();
+    return keysCall.clone();
+  }
+
+  @Override public Call<List<String>> getValues(String key) {
+    if (key == null) throw new NullPointerException("key == null");
+    if (key.isEmpty()) throw new IllegalArgumentException("key was empty");
+    if (!enabled || !autocompleteKeys.contains(key)) return Call.emptyList();
+    return dataSourceCallFactory.create(new SelectAutocompleteValues(schema, key));
+  }
+}

zipkin-storage/mysql-v1/src/main/java/zipkin2/storage/mysql/v1/MySQLSpanStore.java

@@ -29,63 +29,60 @@ final class MySQLSpanStore implements SpanStore {
 
   final DataSourceCall.Factory dataSourceCallFactory;
   final Schema schema;
-  final boolean strictTraceId;
+  final boolean strictTraceId, searchEnabled;
   final SelectSpansAndAnnotations.Factory selectFromSpansAndAnnotationsFactory;
   final Call.Mapper<List<Span>, List<List<Span>>> groupByTraceId;
   final DataSourceCall<List<String>> getServiceNamesCall;
 
-  MySQLSpanStore(
-      DataSourceCall.Factory dataSourceCallFactory, Schema schema, boolean strictTraceId) {
-    this.dataSourceCallFactory = dataSourceCallFactory;
+  MySQLSpanStore(MySQLStorage storage, Schema schema) {
+    this.dataSourceCallFactory = storage.dataSourceCallFactory;
     this.schema = schema;
-    this.strictTraceId = strictTraceId;
+    this.strictTraceId = storage.strictTraceId;
+    this.searchEnabled = storage.searchEnabled;
     this.selectFromSpansAndAnnotationsFactory =
-        new SelectSpansAndAnnotations.Factory(schema, strictTraceId);
+      new SelectSpansAndAnnotations.Factory(schema, strictTraceId);
     this.groupByTraceId = GroupByTraceId.create(strictTraceId);
     this.getServiceNamesCall = dataSourceCallFactory.create(new SelectAnnotationServiceNames());
   }
 
-  @Override
-  public Call<List<List<Span>>> getTraces(QueryRequest request) {
+  @Override public Call<List<List<Span>>> getTraces(QueryRequest request) {
+    if (!searchEnabled) return Call.emptyList();
+
     Call<List<List<Span>>> result =
-        dataSourceCallFactory
-            .create(selectFromSpansAndAnnotationsFactory.create(request))
-            .map(groupByTraceId);
+      dataSourceCallFactory
+        .create(selectFromSpansAndAnnotationsFactory.create(request))
+        .map(groupByTraceId);
 
     return strictTraceId ? result.map(StrictTraceId.filterTraces(request)) : result;
   }
 
-  @Override
-  public Call<List<Span>> getTrace(String hexTraceId) {
+  @Override public Call<List<Span>> getTrace(String hexTraceId) {
     // make sure we have a 16 or 32 character trace ID
     hexTraceId = Span.normalizeTraceId(hexTraceId);
     long traceIdHigh = hexTraceId.length() == 32 ? lowerHexToUnsignedLong(hexTraceId, 0) : 0L;
     long traceId = lowerHexToUnsignedLong(hexTraceId);
 
     DataSourceCall<List<Span>> result =
-        dataSourceCallFactory.create(
-            selectFromSpansAndAnnotationsFactory.create(traceIdHigh, traceId));
+      dataSourceCallFactory.create(
+        selectFromSpansAndAnnotationsFactory.create(traceIdHigh, traceId));
     return strictTraceId ? result.map(StrictTraceId.filterSpans(hexTraceId)) : result;
   }
 
-  @Override
-  public Call<List<String>> getServiceNames() {
+  @Override public Call<List<String>> getServiceNames() {
+    if (!searchEnabled) return Call.emptyList();
     return getServiceNamesCall.clone();
   }
 
-  @Override
-  public Call<List<String>> getSpanNames(String serviceName) {
-    if (serviceName.isEmpty()) return Call.emptyList();
-
+  @Override public Call<List<String>> getSpanNames(String serviceName) {
+    if (serviceName.isEmpty() || !searchEnabled) return Call.emptyList();
     return dataSourceCallFactory.create(new SelectSpanNames(schema, serviceName));
   }
 
-  @Override
-  public Call<List<DependencyLink>> getDependencies(long endTs, long lookback) {
+  @Override public Call<List<DependencyLink>> getDependencies(long endTs, long lookback) {
     if (schema.hasPreAggregatedDependencies) {
       return dataSourceCallFactory.create(new SelectDependencies(schema, getDays(endTs, lookback)));
     }
     return dataSourceCallFactory.create(
-        new AggregateDependencies(schema, endTs * 1000 - lookback * 1000, endTs * 1000));
+      new AggregateDependencies(schema, endTs * 1000 - lookback * 1000, endTs * 1000));
   }
 }

zipkin-storage/mysql-v1/src/main/java/zipkin2/storage/mysql/v1/MySQLStorage.java

@@ -15,14 +15,15 @@
 
 import java.sql.Connection;
 import java.sql.SQLException;
+import java.util.ArrayList;
+import java.util.List;
 import java.util.concurrent.Executor;
 import javax.sql.DataSource;
 import org.jooq.ExecuteListenerProvider;
-import org.jooq.Record;
-import org.jooq.TableField;
 import org.jooq.conf.Settings;
 import zipkin2.CheckResult;
 import zipkin2.internal.Nullable;
+import zipkin2.storage.AutocompleteTags;
 import zipkin2.storage.SpanConsumer;
 import zipkin2.storage.SpanStore;
 import zipkin2.storage.StorageComponent;
@@ -42,20 +43,24 @@ public static final class Builder extends StorageComponent.Builder {
     private Settings settings = new Settings().withRenderSchema(false);
     private ExecuteListenerProvider listenerProvider;
     private Executor executor;
+    List<String> autocompleteKeys = new ArrayList<>();
 
-    /** {@inheritDoc} */
-    @Override
-    public Builder strictTraceId(boolean strictTraceId) {
+    @Override public Builder strictTraceId(boolean strictTraceId) {
       this.strictTraceId = strictTraceId;
       return this;
     }
 
-    @Override
-    public Builder searchEnabled(boolean searchEnabled) {
+    @Override public Builder searchEnabled(boolean searchEnabled) {
       this.searchEnabled = searchEnabled;
       return this;
     }
 
+    @Override public Builder autocompleteKeys(List<String> keys) {
+      if (keys == null) throw new NullPointerException("keys == null");
+      this.autocompleteKeys = keys;
+      return this;
+    }
+
     public Builder datasource(DataSource datasource) {
       if (datasource == null) throw new NullPointerException("datasource == null");
       this.datasource = datasource;
@@ -79,8 +84,7 @@ public Builder executor(Executor executor) {
       return this;
     }
 
-    @Override
-    public MySQLStorage build() {
+    @Override public MySQLStorage build() {
       return new MySQLStorage(this);
     }
 
@@ -95,7 +99,7 @@ public MySQLStorage build() {
   final DataSourceCall.Factory dataSourceCallFactory;
   final DSLContexts context;
   final boolean strictTraceId, searchEnabled;
-
+  final List<String> autocompleteKeys;
   volatile Schema schema;
 
   MySQLStorage(MySQLStorage.Builder builder) {
@@ -107,6 +111,7 @@ public MySQLStorage build() {
     dataSourceCallFactory = new DataSourceCall.Factory(datasource, context, executor);
     strictTraceId = builder.strictTraceId;
     searchEnabled = builder.searchEnabled;
+    autocompleteKeys = builder.autocompleteKeys;
   }
 
   /** Returns the session in use by this storage component. */
@@ -128,7 +133,11 @@ Schema schema() {
 
   @Override
   public SpanStore spanStore() {
-    return new MySQLSpanStore(dataSourceCallFactory, schema(), strictTraceId);
+    return new MySQLSpanStore(this, schema());
+  }
+
+  @Override public AutocompleteTags autocompleteTags() {
+    return new MySQLAutocompleteTags(this, schema());
   }
 
   @Override

zipkin-storage/mysql-v1/src/main/java/zipkin2/storage/mysql/v1/SelectAutocompleteValues.java

@@ -0,0 +1,59 @@
+/*
+ * Copyright 2015-2018 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 zipkin2.storage.mysql.v1;
+
+import java.util.List;
+import java.util.function.Function;
+import org.jooq.Converter;
+import org.jooq.DSLContext;
+import zipkin2.v1.V1BinaryAnnotation;
+
+import static java.nio.charset.StandardCharsets.UTF_8;
+import static zipkin2.storage.mysql.v1.internal.generated.tables.ZipkinAnnotations.ZIPKIN_ANNOTATIONS;
+
+final class SelectAutocompleteValues implements Function<DSLContext, List<String>> {
+  final Schema schema;
+  final String autocompleteKey;
+
+  SelectAutocompleteValues(Schema schema, String autocompleteKey) {
+    this.schema = schema;
+    this.autocompleteKey = autocompleteKey;
+  }
+
+  @Override public List<String> apply(DSLContext context) {
+    return context.selectDistinct(ZIPKIN_ANNOTATIONS.A_VALUE)
+      .from(ZIPKIN_ANNOTATIONS)
+      .where(ZIPKIN_ANNOTATIONS.A_TYPE.eq(V1BinaryAnnotation.TYPE_STRING)
+        .and(ZIPKIN_ANNOTATIONS.A_KEY.eq(autocompleteKey)))
+      .fetch(ZIPKIN_ANNOTATIONS.A_VALUE, STRING_CONVERTER);
+  }
+
+  static final Converter<byte[], String> STRING_CONVERTER = new Converter<byte[], String>() {
+    @Override public String from(byte[] bytes) {
+      return new String(bytes, UTF_8);
+    }
+
+    @Override public byte[] to(String input) {
+      return input.getBytes(UTF_8);
+    }
+
+    @Override public Class<byte[]> fromType() {
+      return byte[].class;
+    }
+
+    @Override public Class<String> toType() {
+      return String.class;
+    }
+  };
+}

zipkin-storage/mysql-v1/src/main/resources/mysql.sql

@@ -33,8 +33,8 @@ ALTER TABLE zipkin_annotations ADD UNIQUE KEY(`trace_id_high`, `trace_id`, `span
 ALTER TABLE zipkin_annotations ADD INDEX(`trace_id_high`, `trace_id`, `span_id`) COMMENT 'for joining with zipkin_spans';
 ALTER TABLE zipkin_annotations ADD INDEX(`trace_id_high`, `trace_id`) COMMENT 'for getTraces/ByIds';
 ALTER TABLE zipkin_annotations ADD INDEX(`endpoint_service_name`) COMMENT 'for getTraces and getServiceNames';
-ALTER TABLE zipkin_annotations ADD INDEX(`a_type`) COMMENT 'for getTraces';
-ALTER TABLE zipkin_annotations ADD INDEX(`a_key`) COMMENT 'for getTraces';
+ALTER TABLE zipkin_annotations ADD INDEX(`a_type`) COMMENT 'for getTraces and autocomplete values';
+ALTER TABLE zipkin_annotations ADD INDEX(`a_key`) COMMENT 'for getTraces and autocomplete values';
 ALTER TABLE zipkin_annotations ADD INDEX(`trace_id`, `span_id`, `a_key`) COMMENT 'for dependencies job';
 
 CREATE TABLE IF NOT EXISTS zipkin_dependencies (