diff --git a/astra/src/main/java/com/slack/astra/graphApi/GraphConfig.java b/astra/src/main/java/com/slack/astra/graphApi/GraphConfig.java new file mode 100644 index 0000000000..501bf8bbd3 --- /dev/null +++ b/astra/src/main/java/com/slack/astra/graphApi/GraphConfig.java @@ -0,0 +1,187 @@ +package com.slack.astra.graphApi; + +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.PropertyNamingStrategies; +import com.fasterxml.jackson.dataformat.yaml.YAMLFactory; +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +public final class GraphConfig { + private static final Logger LOG = LoggerFactory.getLogger(GraphConfig.class); + + public static final GraphConfig DEFAULT = new GraphConfig(Map.of()); + + /** + * Represents how a single logical field on a node should be mapped to span tags. + * + *

Each field has: - a default key to look up in tags - a default fallback value if the key + * isn’t found - an optional list of rules that can override the default key + */ + public static final class TagConfig { + private final String defaultKey; + private final String defaultValue; + private final List rules; + + @JsonCreator + public TagConfig( + @JsonProperty("default_key") String defaultKey, + @JsonProperty("default_value") String defaultValue, + @JsonProperty("rules") List rules) { + this.defaultKey = defaultKey; + this.defaultValue = defaultValue; + this.rules = (rules == null) ? Collections.emptyList() : List.copyOf(rules); + } + + public String getDefaultKey() { + return defaultKey; + } + + public String getDefaultValue() { + return defaultValue; + } + + public List getRules() { + return rules; + } + } + + /** + * Represents a conditional rule for overriding which tag key to use. Note: This logic does not + * currently support multiple field matches under a single rule. + */ + public static class RuleConfig { + private final String field; + private final String value; + private final String overrideKey; + + @JsonCreator + public RuleConfig( + @JsonProperty("field") String field, + @JsonProperty("value") String value, + @JsonProperty("override_key") String overrideKey) { + this.field = field; + this.value = value; + this.overrideKey = overrideKey; + } + + public String getField() { + return field; + } + + public String getValue() { + return value; + } + + public String getOverrideKey() { + return overrideKey; + } + } + + // Holds the entire mapping for logical field names to their configuration of defaults and rules. + private final Map nodeMetadataTagMapping; + + @JsonCreator + public GraphConfig( + @JsonProperty("node_metadata_tag_mapping") Map nodeMetadataTagMapping) { + this.nodeMetadataTagMapping = Map.copyOf(nodeMetadataTagMapping); + } + + public Map getNodeMetadataTagMapping() { + return nodeMetadataTagMapping; + } + + /** + * Loads a GraphConfig from a YAML file on disk. On failure, log and return a null config. + * + * @param filePath Path of the config to load + * @return GraphConfig containing mappings of node metadata fields -> corresponding span tags. + */ + public static GraphConfig load(Path filePath) throws IOException { + try { + String yaml = Files.readString(filePath); + return load(yaml); + } catch (Exception e) { + LOG.warn( + "Failed to read dependency graph config from file path. Returning default config", e); + } + + return DEFAULT; + } + + /** + * Loads a GraphConfig from a YAML string. On failure, log and return a null config. + * + * @param configYAML YAML string of config contents. + * @return GraphConfig containing mappings of node metadata fields -> corresponding span tags. + */ + public static GraphConfig load(String configYAML) throws IOException { + if (!configYAML.isEmpty()) { + try { + ObjectMapper mapper = new ObjectMapper(new YAMLFactory()); + mapper.setPropertyNamingStrategy(PropertyNamingStrategies.SNAKE_CASE); + return mapper.readValue(configYAML, GraphConfig.class); + } catch (Exception e) { + LOG.warn( + "Failed to parse dependency graph config file contents. Returning default config", e); + } + } + + return DEFAULT; + } + + /** + * Resolves the actual tag value for a given logical field, using the provided span tags. + * + *

Steps: 1. Look up the TagConfig for this logical field (e.g. "resource"). 2. Default to + * using its defaultKey + defaultValue. 3. If rules are defined: - Iterate through each rule in + * reverse order. - If a rule’s field/value condition matches, switch keyToUse to overrideKey. 4. + * Finally, look up the chosen key in tags. If missing, fall back to defaultValue. + * + *

Note: This logic does not currently support multiple field matches for a single rule. + * + * @param tags Map of tags from the span. + * @param logicalField the node metadata field to resolve. + * @return String the value of the logical metadata field after applying all GraphConfig rules. + */ + public String resolve(Map tags, String logicalField) { + TagConfig baseCfg = nodeMetadataTagMapping.get(logicalField); + + // If the config doesn't define this logical field, just return from raw tags or fallback. + if (baseCfg == null) { + return tags.getOrDefault(logicalField, "unknown_" + logicalField); + } + + // Later rules override earlier ones, so start from the back of the list and use the first one + // that matches. + String keyToUse = + baseCfg.getRules().reversed().stream() + .filter( + rule -> + rule.getValue() + .equals(tags.getOrDefault(rule.getField(), "unknown_" + rule.getField()))) + .map(RuleConfig::getOverrideKey) + .filter(tags::containsKey) + .findFirst() + .orElse(baseCfg.getDefaultKey()); + + return tags.getOrDefault(keyToUse, baseCfg.getDefaultValue()); + } + + @Override + public String toString() { + try { + ObjectMapper mapper = new ObjectMapper(); + return mapper.writerWithDefaultPrettyPrinter().writeValueAsString(this); + } catch (Exception e) { + return "GraphConfig{error serializing to string}"; + } + } +} diff --git a/astra/src/main/java/com/slack/astra/graphApi/GraphService.java b/astra/src/main/java/com/slack/astra/graphApi/GraphService.java index 0826b2196f..933bcb5e80 100644 --- a/astra/src/main/java/com/slack/astra/graphApi/GraphService.java +++ b/astra/src/main/java/com/slack/astra/graphApi/GraphService.java @@ -21,6 +21,7 @@ public class GraphService { private static final Logger LOG = LoggerFactory.getLogger(GraphService.class); private final AstraQueryServiceBase searcher; + private final GraphConfig graphConfig; private static final ObjectMapper objectMapper = JsonMapper.builder() @@ -30,8 +31,11 @@ public class GraphService { .serializationInclusion(JsonInclude.Include.NON_EMPTY) .build(); - public GraphService(AstraQueryServiceBase searcher) { + public GraphService(AstraQueryServiceBase searcher, GraphConfig graphConfig) { this.searcher = searcher; + this.graphConfig = graphConfig; + + LOG.info("Started GraphService with config: {}", this.graphConfig); } @Get diff --git a/astra/src/main/java/com/slack/astra/server/Astra.java b/astra/src/main/java/com/slack/astra/server/Astra.java index 7d83715bb1..e51cf50c61 100644 --- a/astra/src/main/java/com/slack/astra/server/Astra.java +++ b/astra/src/main/java/com/slack/astra/server/Astra.java @@ -23,6 +23,7 @@ import com.slack.astra.clusterManager.ReplicaRestoreService; import com.slack.astra.clusterManager.SnapshotDeletionService; import com.slack.astra.elasticsearchApi.ElasticsearchApiService; +import com.slack.astra.graphApi.GraphConfig; import com.slack.astra.graphApi.GraphService; import com.slack.astra.logstore.LogMessage; import com.slack.astra.logstore.schema.ReservedFields; @@ -240,6 +241,16 @@ private static Set getServices( // https://github.com/slackhq/astra/pull/564) final int serverPort = astraConfig.getQueryConfig().getServerConfig().getServerPort(); + GraphConfig graphConfig = GraphConfig.DEFAULT; + String depGraphConfigFile = astraConfig.getQueryConfig().getDepGraphConfigFile(); + + if (!depGraphConfigFile.isEmpty()) { + LOG.info("Loading dependency graph config file: {}", depGraphConfigFile); + graphConfig = GraphConfig.load(Path.of(depGraphConfigFile)); + } else { + LOG.info("No dependency graph config file provided, using empty config"); + } + ArmeriaService armeriaService = new ArmeriaService.Builder(serverPort, "astraQuery", meterRegistry) .withRequestTimeout(requestTimeout) @@ -252,7 +263,7 @@ private static Set getServices( astraConfig.getQueryConfig().getZipkinDefaultMaxSpans(), astraConfig.getQueryConfig().getZipkinDefaultLookbackMins(), astraConfig.getQueryConfig().getZipkinDefaultDataFreshnessSecs())) - .withAnnotatedService(new GraphService(astraDistributedQueryService)) + .withAnnotatedService(new GraphService(astraDistributedQueryService, graphConfig)) .withGrpcService(astraDistributedQueryService) .build(); services.add(armeriaService); diff --git a/astra/src/main/proto/astra_configs.proto b/astra/src/main/proto/astra_configs.proto index 465abf9ed0..190483e36c 100644 --- a/astra/src/main/proto/astra_configs.proto +++ b/astra/src/main/proto/astra_configs.proto @@ -98,6 +98,7 @@ message QueryServiceConfig { int32 zipkin_default_max_spans = 4; int32 zipkin_default_lookback_mins = 5; int64 zipkin_default_data_freshness_secs = 6; + string dep_graph_config_file = 7; } message RedactionUpdateServiceConfig { diff --git a/astra/src/test/java/com/slack/astra/graphApi/GraphConfigTest.java b/astra/src/test/java/com/slack/astra/graphApi/GraphConfigTest.java new file mode 100644 index 0000000000..612b623cda --- /dev/null +++ b/astra/src/test/java/com/slack/astra/graphApi/GraphConfigTest.java @@ -0,0 +1,272 @@ +package com.slack.astra.graphApi; + +import static org.assertj.core.api.Assertions.assertThat; + +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.HashMap; +import java.util.Map; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +public class GraphConfigTest { + @Test + public void testLoadValidYamlConfig(@TempDir Path tempDir) throws IOException { + String yamlContent = + """ + node_metadata_tag_mapping: + service: + default_key: service.name + default_value: unknown_service + rules: + - field: cluster.name + value: prod + override_key: prod.service.name + - field: cluster.name + value: staging + override_key: test.service.name + cluster: + default_key: cluster.name + default_value: unknown_cluster + """; + + Path configFile = tempDir.resolve("test-config.yaml"); + Files.writeString(configFile, yamlContent); + + GraphConfig config = GraphConfig.load(configFile); + + assertThat(config).isNotNull(); + assertThat(config.getNodeMetadataTagMapping()).hasSize(2); + + GraphConfig.TagConfig serviceConfig = config.getNodeMetadataTagMapping().get("service"); + assertThat(serviceConfig.getDefaultKey()).isEqualTo("service.name"); + assertThat(serviceConfig.getDefaultValue()).isEqualTo("unknown_service"); + assertThat(serviceConfig.getRules()).hasSize(2); + + GraphConfig.RuleConfig rule1 = serviceConfig.getRules().getFirst(); + assertThat(rule1.getOverrideKey()).isEqualTo("prod.service.name"); + assertThat(rule1.getField()).isEqualTo("cluster.name"); + assertThat(rule1.getValue()).isEqualTo("prod"); + + GraphConfig.RuleConfig rule2 = serviceConfig.getRules().get(1); + assertThat(rule2.getOverrideKey()).isEqualTo("test.service.name"); + assertThat(rule2.getField()).isEqualTo("cluster.name"); + assertThat(rule2.getValue()).isEqualTo("staging"); + + GraphConfig.TagConfig clusterConfig = config.getNodeMetadataTagMapping().get("cluster"); + assertThat(clusterConfig.getDefaultKey()).isEqualTo("cluster.name"); + assertThat(clusterConfig.getDefaultValue()).isEqualTo("unknown_cluster"); + assertThat(clusterConfig.getRules()).hasSize(0); + } + + @Test + public void testLoadNonExistentFile() throws IOException { + GraphConfig config = GraphConfig.load(Path.of("/non/existent/file.yaml")); + assertThat(config).isEqualTo(GraphConfig.DEFAULT); + } + + @Test + public void testLoadEmptyConfigFile() throws IOException { + GraphConfig config = GraphConfig.load(Path.of("")); + assertThat(config).isEqualTo(GraphConfig.DEFAULT); + } + + @Test + public void testLoadInvalidYaml(@TempDir Path tempDir) throws IOException { + String invalidYaml = + """ + invalid: yaml: content: + - broken + """; + + Path configFile = tempDir.resolve("invalid-config.yaml"); + Files.writeString(configFile, invalidYaml); + + GraphConfig config = GraphConfig.load(configFile); + assertThat(config).isEqualTo(GraphConfig.DEFAULT); + } + + @Test + public void testResolveWithDefaultValue() throws IOException { + GraphConfig config = + GraphConfig.load( + """ + node_metadata_tag_mapping: + app: + default_key: app.name + default_value: unknown_app + namespace: + default_key: namespace.name + default_value: unknown_namespace + """); + Map tags = new HashMap<>(); + + String result = config.resolve(tags, "app"); + assertThat(result).isEqualTo("unknown_app"); + } + + @Test + public void testResolveWithUnknownField() throws IOException { + GraphConfig config = + GraphConfig.load( + """ + node_metadata_tag_mapping: + app: + default_key: app.name + default_value: unknown_app + namespace: + default_key: namespace.name + default_value: unknown_namespace + """); + Map tags = Map.of("some.tag", "some-value"); + + String result = config.resolve(tags, "some_field"); + assertThat(result).isEqualTo("unknown_some_field"); + } + + @Test + public void testResolveWithMatchingRule() throws IOException { + GraphConfig config = + GraphConfig.load( + """ + node_metadata_tag_mapping: + app: + default_key: app.name + default_value: unknown_app + rules: + - field: namespace.name + value: prod-ns + override_key: prod.app.name + - field: cluster.name + value: east + override_key: east.app.name + namespace: + default_key: namespace.name + default_value: unknown_namespace + """); + Map tags = + Map.of( + "app.name", "my-app", "namespace.name", "prod-ns", "prod.app.name", "my-app-in-prod"); + + String result = config.resolve(tags, "app"); + assertThat(result).isEqualTo("my-app-in-prod"); + } + + @Test + public void testResolveWithMatchingRuleButMissingOverrideKey() throws IOException { + GraphConfig config = + GraphConfig.load( + """ + node_metadata_tag_mapping: + app: + default_key: app.name + default_value: unknown_app + rules: + - field: namespace.name + value: prod-ns + override_key: prod.app.name + - field: cluster.name + value: east + override_key: east.app.name + namespace: + default_key: namespace.name + default_value: unknown_namespace + """); + Map tags = Map.of("app.name", "my-app", "namespace.name", "prod-ns"); + + String result = config.resolve(tags, "app"); + assertThat(result).isEqualTo("my-app"); + } + + @Test + public void testResolveWithNonMatchingRule() throws IOException { + GraphConfig config = + GraphConfig.load( + """ + node_metadata_tag_mapping: + app: + default_key: app.name + default_value: unknown_app + rules: + - field: namespace.name + value: prod-ns + override_key: prod.app.name + - field: cluster.name + value: east + override_key: east.app.name + namespace: + default_key: namespace.name + default_value: unknown_namespace + """); + Map tags = + Map.of( + "app.name", "my-app", + "namespace.name", "dev-ns"); + + String result = config.resolve(tags, "app"); + assertThat(result).isEqualTo("my-app"); + } + + @Test + public void testResolveWithMultipleRules() throws IOException { + GraphConfig config = + GraphConfig.load( + """ + node_metadata_tag_mapping: + app: + default_key: app.name + default_value: unknown_app + rules: + - field: namespace.name + value: prod-ns + override_key: prod.app.name + - field: cluster.name + value: east + override_key: east.app.name + namespace: + default_key: namespace.name + default_value: unknown_namespace + """); + Map tags = + Map.of( + "app.name", "my-app", + "namespace.name", "prod-ns", + "cluster.name", "east", + "prod.app.name", "my-app-in-prod", + "east.app.name", "my-app-east"); + + String result = config.resolve(tags, "app"); + assertThat(result).isEqualTo("my-app-east"); + } + + @Test + public void testResolveWithMultipleRulesNoMatch() throws IOException { + GraphConfig config = + GraphConfig.load( + """ + node_metadata_tag_mapping: + app: + default_key: app.name + default_value: unknown_app + rules: + - field: namespace.name + value: prod-ns + override_key: prod.app.name + - field: cluster.name + value: east + override_key: east.app.name + namespace: + default_key: namespace.name + default_value: unknown_namespace + """); + Map tags = + Map.of( + "app.name", "my-app", + "namespace.name", "dev-ns", + "cluster.name", "west"); + + String result = config.resolve(tags, "app"); + assertThat(result).isEqualTo("my-app"); + } +} diff --git a/astra/src/test/java/com/slack/astra/graphApi/GraphServiceTest.java b/astra/src/test/java/com/slack/astra/graphApi/GraphServiceTest.java index 0b024ad7ed..f461948575 100644 --- a/astra/src/test/java/com/slack/astra/graphApi/GraphServiceTest.java +++ b/astra/src/test/java/com/slack/astra/graphApi/GraphServiceTest.java @@ -18,7 +18,7 @@ public class GraphServiceTest { @BeforeEach public void setup() throws IOException { - graphService = spy(new GraphService(searcher)); + graphService = spy(new GraphService(searcher, GraphConfig.DEFAULT)); } @Test diff --git a/config/config.yaml b/config/config.yaml index f404fbd406..bf4c283989 100644 --- a/config/config.yaml +++ b/config/config.yaml @@ -54,6 +54,7 @@ queryConfig: zipkinDefaultLookbackMins: ${ASTRA_QUERY_ZIPKIN_DEFAULT_LOOKBACK_MINS:-10080} zipkinDefaultDataFreshnessSecs: ${ASTRA_QUERY_DEFAULT_DATA_FRESHNESS_SECONDS:-900} managerConnectString: ${ASTRA_MANAGER_CONNECTION_STRING:-localhost:8083} + depGraphConfigFile: ${DEP_GRAPH_CONFIG_FILE:-} metadataStoreConfig: mode: ${ASTRA_METADATA_MODE:-ZOOKEEPER_EXCLUSIVE}