features

#!/usr/bin/env python3
import os, sys, tomlkit, yaml
from yaml.loader import SafeLoader

try:
  _, filename = sys.argv
except:
  print("usage: {} CONFIG.toml".format(sys.argv[0]))
  exit(1)

_, ext = os.path.splitext(filename.lower())
if ext in ['.yaml', '.yml']:
  with open(filename) as fh:
    config = yaml.load(fh, Loader=SafeLoader)
elif ext == '.toml':
  with open(filename, 'rb') as fh:
    config = tomlkit.load(fh)
else:
  print("{}: Unknown config filename extension.".format(sys.argv[0]))
  exit(1)

base_features = set()

for key in ['api', 'enterprise']:
  if config.get(key) is not None:
    base_features.add(key)

# Extract the set of features for a particular key from the config,
# using the mapping to rewrite component names to their feature names.
def get_features(config, key, mapping):
  components = set( component['type'] for component in config.get(key, {}).values() )
  components = set( mapping.get(component, component) for component in components )
  return set( f'{key}-{component}' for component in components )

# Mapping of source names to feature names.
source_feature_map = {
  'generator': 'demo_logs',
  'logplex': 'heroku_logs',
  'prometheus_scrape': 'prometheus',
  'prometheus_remote_write': 'prometheus',
}

# Mapping of transform names to feature names.
transform_feature_map = {
  'sampler': 'sample',
  'swimlanes': 'route',
}

# Mapping of sink names to feature names.
sink_feature_map = {
  'gcp_pubsub': 'gcp',
  'gcp_cloud_storage': 'gcp',
  'gcp_stackdriver_logs': 'gcp',
  'gcp_stackdriver_metrics': 'gcp',
  'prometheus_exporter': 'prometheus',
  'prometheus_remote_write': 'prometheus',
  'splunk_hec_logs': 'splunk_hec',
}

# Set of always-compiled components, in terms of their computed feature flag, that should not be
# emitted as they don't actually have a feature flag... because we always compile them.
invalid_feature_flags = {
  'transforms-log_to_metric',
}

features = base_features \
  .union(get_features(config, 'sources', source_feature_map)) \
  .union(get_features(config, 'transforms', transform_feature_map)) \
  .union(get_features(config, 'sinks', sink_feature_map)) \
  .difference(invalid_feature_flags)

print(','.join(sorted(features)))