github.com/streamdal/segmentio-kafka-go@v0.4.47-streamdal/streamdal.go (about)

     1  package kafka
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"os"
     7  
     8  	streamdal "github.com/streamdal/streamdal/sdks/go"
     9  )
    10  
    11  const (
    12  	StreamdalEnvAddress     = "STREAMDAL_ADDRESS"
    13  	StreamdalEnvAuthToken   = "STREAMDAL_AUTH_TOKEN"
    14  	StreamdalEnvServiceName = "STREAMDAL_SERVICE_NAME"
    15  
    16  	StreamdalDefaultComponentName = "kafka"
    17  	StreamdalDefaultOperationName = "unknown"
    18  )
    19  
    20  // StreamdalRuntimeConfig is an optional configuration structure that can be
    21  // passed to kafka.FetchMessage() and kafka.WriteMessage() methods to influence
    22  // streamdal shim behavior.
    23  //
    24  // NOTE: This struct is intended to be passed as a value in a context.Context.
    25  // This is done this way to avoid having to change FetchMessage() and WriteMessages()
    26  // signatures.
    27  type StreamdalRuntimeConfig struct {
    28  	// StrictErrors will cause the shim to return a kafka.Error if Streamdal.Process()
    29  	// runs into an unrecoverable error. Default: swallow error and return original value.
    30  	StrictErrors bool
    31  
    32  	// Audience is used to specify a custom audience when the shim calls on
    33  	// streamdal.Process(); if nil, a default ComponentName and OperationName
    34  	// will be used. Only non-blank values will be used to override audience defaults.
    35  	Audience *streamdal.Audience
    36  }
    37  
    38  func streamdalSetup() (*streamdal.Streamdal, error) {
    39  	address := os.Getenv(StreamdalEnvAddress)
    40  	if address == "" {
    41  		return nil, errors.New(StreamdalEnvAddress + " env var is not set")
    42  	}
    43  
    44  	authToken := os.Getenv(StreamdalEnvAuthToken)
    45  	if authToken == "" {
    46  		return nil, errors.New(StreamdalEnvAuthToken + " env var is not set")
    47  	}
    48  
    49  	serviceName := os.Getenv(StreamdalEnvServiceName)
    50  	if serviceName == "" {
    51  		return nil, errors.New(StreamdalEnvServiceName + " env var is not set")
    52  	}
    53  
    54  	sc, err := streamdal.New(&streamdal.Config{
    55  		ServerURL:   address,
    56  		ServerToken: authToken,
    57  		ServiceName: serviceName,
    58  		ClientType:  streamdal.ClientTypeShim,
    59  	})
    60  
    61  	if err != nil {
    62  		return nil, errors.New("unable to create streamdal client: " + err.Error())
    63  	}
    64  
    65  	return sc, nil
    66  }
    67  
    68  func streamdalProcess(ctx context.Context, sc *streamdal.Streamdal, ot streamdal.OperationType, msg *Message, loggers ...Logger) (*Message, error) {
    69  	// Nothing to do if streamdal client is nil
    70  	if sc == nil {
    71  		return msg, nil
    72  	}
    73  
    74  	// Maybe extract runtime config from context
    75  	var src *StreamdalRuntimeConfig
    76  	if ctx != nil {
    77  		src = ctx.Value("streamdal-runtime-config").(*StreamdalRuntimeConfig)
    78  	}
    79  
    80  	// Generate an audience from the provided parameters
    81  	aud := streamdalGenerateAudience(ot, msg.Topic, src)
    82  
    83  	// Process msg payload via Streamdal
    84  	resp := sc.Process(ctx, &streamdal.ProcessRequest{
    85  		ComponentName: aud.ComponentName,
    86  		OperationType: ot,
    87  		OperationName: aud.OperationName,
    88  		Data:          msg.Value,
    89  	})
    90  
    91  	switch resp.Status {
    92  	case streamdal.ExecStatusTrue, streamdal.ExecStatusFalse:
    93  		// Process() did not error - replace kafka.Value
    94  		msg.Value = resp.Data
    95  	case streamdal.ExecStatusError:
    96  		// Process() errored - return message as-is; if strict errors are NOT
    97  		// set, return error instead of message
    98  		streamdalLogError(loggers, "streamdal.Process() error: "+ptrStr(resp.StatusMessage))
    99  
   100  		if src != nil && src.StrictErrors {
   101  			return nil, errors.New("streamdal.Process() error: " + ptrStr(resp.StatusMessage))
   102  		}
   103  	}
   104  
   105  	return msg, nil
   106  }
   107  
   108  // Helper func for generating an "audience" that can be passed to streamdal's .Process() method.
   109  //
   110  // Topic is only used if the provided runtime config is nil or the underlying
   111  // audience does not have an OperationName set.
   112  func streamdalGenerateAudience(ot streamdal.OperationType, topic string, src *StreamdalRuntimeConfig) *streamdal.Audience {
   113  	var (
   114  		componentName = StreamdalDefaultComponentName
   115  		operationName = StreamdalDefaultOperationName
   116  	)
   117  
   118  	if topic != "" {
   119  		operationName = topic
   120  	}
   121  
   122  	if src != nil && src.Audience != nil {
   123  		if src.Audience.OperationName != "" {
   124  			operationName = src.Audience.OperationName
   125  		}
   126  
   127  		if src.Audience.ComponentName != "" {
   128  			componentName = src.Audience.ComponentName
   129  		}
   130  	}
   131  
   132  	return &streamdal.Audience{
   133  		OperationType: ot,
   134  		OperationName: operationName,
   135  		ComponentName: componentName,
   136  	}
   137  }
   138  
   139  // Helper func for logging errors encountered during streamdal.Process()
   140  func streamdalLogError(loggers []Logger, msg string) {
   141  	for _, l := range loggers {
   142  		if l == nil {
   143  			continue
   144  		}
   145  
   146  		l.Printf(msg)
   147  	}
   148  }
   149  
   150  // Helper func to deref string ptrs
   151  func ptrStr(s *string) string {
   152  	if s == nil {
   153  		return ""
   154  	}
   155  
   156  	return *s
   157  }