intake/core/action.go

258 lines
5.9 KiB
Go
Raw Normal View History

2025-01-17 21:49:23 +00:00
package core
import (
"bufio"
"context"
2025-01-21 03:53:22 +00:00
"database/sql/driver"
2025-01-17 21:49:23 +00:00
"encoding/json"
"errors"
2025-01-30 07:14:00 +00:00
"fmt"
2025-01-17 21:49:23 +00:00
"io"
"log"
"os"
"os/exec"
"strings"
"time"
)
2025-01-21 03:53:22 +00:00
// Type alias for storing string array as jsonb
type argList []string
func (a argList) Value() (driver.Value, error) {
return json.Marshal(a)
}
func (a *argList) Scan(value interface{}) error {
return json.Unmarshal([]byte(value.(string)), a)
}
2025-01-31 16:44:09 +00:00
func AddAction(db DB, source string, name string, argv []string) error {
2025-01-21 03:53:22 +00:00
_, err := db.Exec(`
insert into actions (source, name, argv)
values (?, ?, jsonb(?))
`, source, name, argList(argv))
return err
}
2025-01-31 16:44:09 +00:00
func UpdateAction(db DB, source string, name string, argv []string) error {
2025-01-21 03:53:22 +00:00
_, err := db.Exec(`
update actions
set argv = jsonb(?)
where source = ? and name = ?
`, argList(argv), source, name)
return err
}
2025-01-31 16:44:09 +00:00
func GetActionsForSource(db DB, source string) ([]string, error) {
2025-01-21 03:53:22 +00:00
rows, err := db.Query(`
select name
from actions
where source = ?
`, source)
if err != nil {
return nil, err
}
var names []string
for rows.Next() {
var name string
err = rows.Scan(&name)
if err != nil {
return nil, err
}
names = append(names, name)
}
return names, nil
}
2025-01-31 16:44:09 +00:00
func GetArgvForAction(db DB, source string, name string) ([]string, error) {
2025-01-21 03:53:22 +00:00
rows := db.QueryRow(`
select json(argv)
from actions
where source = ? and name = ?
`, source, name)
var argv argList
err := rows.Scan(&argv)
if err != nil {
return nil, err
}
return argv, nil
}
2025-01-31 16:44:09 +00:00
func DeleteAction(db DB, source string, name string) error {
2025-01-21 03:53:22 +00:00
_, err := db.Exec(`
delete from actions
where source = ? and name = ?
`, source, name)
2025-01-21 03:53:22 +00:00
return err
}
2025-01-23 20:26:21 +00:00
func readStdout(stdout io.ReadCloser, source string, items chan Item, cparse chan bool) {
2025-01-17 21:49:23 +00:00
var item Item
parseError := false
scanout := bufio.NewScanner(stdout)
for scanout.Scan() {
data := scanout.Bytes()
err := json.Unmarshal(data, &item)
2025-01-28 05:27:20 +00:00
if err != nil || item.Id == "" {
2025-01-23 20:26:21 +00:00
log.Printf("[%s: stdout] %s\n", source, strings.TrimSpace(string(data)))
2025-01-17 21:49:23 +00:00
parseError = true
} else {
2025-01-23 20:26:21 +00:00
item.Active = true // These fields aren't up to
item.Created = 0 // the action to set and
item.Source = source // shouldn't be overrideable
log.Printf("[%s: item] %s\n", source, item.Id)
2025-01-17 21:49:23 +00:00
items <- item
}
}
// Only send the parsing result at the end, to block main until stdout is drained
cparse <- parseError
close(items)
}
2025-01-23 20:26:21 +00:00
func readStderr(stderr io.ReadCloser, source string, done chan bool) {
2025-01-17 21:49:23 +00:00
scanerr := bufio.NewScanner(stderr)
for scanerr.Scan() {
text := strings.TrimSpace(scanerr.Text())
2025-01-23 20:26:21 +00:00
log.Printf("[%s: stderr] %s\n", source, text)
2025-01-17 21:49:23 +00:00
}
done <- true
}
func writeStdin(stdin io.WriteCloser, text string) {
defer stdin.Close()
io.WriteString(stdin, text)
}
func Execute(
2025-01-23 20:26:21 +00:00
source string,
2025-01-17 21:49:23 +00:00
argv []string,
env []string,
state []byte,
2025-01-17 21:49:23 +00:00
input string,
timeout time.Duration,
) ([]Item, []byte, error) {
2025-01-30 07:51:06 +00:00
log.Printf("executing %v", argv)
2025-01-17 21:49:23 +00:00
if len(argv) == 0 {
return nil, nil, errors.New("empty argv")
2025-01-17 21:49:23 +00:00
}
2025-01-29 15:39:00 +00:00
if source == "" {
return nil, nil, errors.New("empty source")
2025-01-29 15:39:00 +00:00
}
2025-01-17 21:49:23 +00:00
stateFile, err := os.CreateTemp("", "intake_state_*")
if err != nil {
return nil, nil, fmt.Errorf("error: failed to create temp state file: %v", err)
}
defer func() {
if err := os.Remove(stateFile.Name()); err != nil {
log.Printf("error: failed to delete %s", stateFile.Name())
}
}()
_, err = stateFile.Write(state)
if err != nil {
return nil, nil, fmt.Errorf("error: failed to write state file: %v", err)
}
env = append(env, "STATE_PATH="+stateFile.Name())
2025-01-17 21:49:23 +00:00
ctx, cancel := context.WithTimeout(context.Background(), timeout)
defer cancel()
cmd := exec.CommandContext(ctx, argv[0], argv[1:]...)
cmd.Env = append(os.Environ(), env...)
cmd.WaitDelay = time.Second * 5
// Open pipes to the command
stdin, err := cmd.StdinPipe()
if err != nil {
return nil, nil, err
2025-01-17 21:49:23 +00:00
}
stdout, err := cmd.StdoutPipe()
if err != nil {
return nil, nil, err
2025-01-17 21:49:23 +00:00
}
stderr, err := cmd.StderrPipe()
if err != nil {
return nil, nil, err
2025-01-17 21:49:23 +00:00
}
cout := make(chan Item)
cparse := make(chan bool)
cerr := make(chan bool)
// Sink routine for items produced
var items []Item
go func() {
for item := range cout {
items = append(items, item)
}
}()
// Routines handling the process i/o
go writeStdin(stdin, input)
2025-01-23 20:26:21 +00:00
go readStdout(stdout, source, cout, cparse)
go readStderr(stderr, source, cerr)
2025-01-17 21:49:23 +00:00
// Kick off the command
err = cmd.Start()
if err != nil {
return nil, nil, err
2025-01-17 21:49:23 +00:00
}
// Block until std{out,err} close
<-cerr
parseError := <-cparse
err = cmd.Wait()
if ctx.Err() == context.DeadlineExceeded {
log.Printf("Timed out after %v\n", timeout)
return nil, nil, err
2025-01-17 21:49:23 +00:00
} else if exiterr, ok := err.(*exec.ExitError); ok {
log.Printf("error: %s failed with exit code %d\n", argv[0], exiterr.ExitCode())
return nil, nil, err
2025-01-17 21:49:23 +00:00
} else if err != nil {
log.Printf("error: %s failed with error: %s\n", argv[0], err)
return nil, nil, err
2025-01-17 21:49:23 +00:00
}
if parseError {
log.Printf("error: could not parse item\n")
return nil, nil, errors.New("invalid JSON")
}
newState, err := os.ReadFile(stateFile.Name())
if err != nil {
return nil, nil, fmt.Errorf("error: failed to read state file: %v", err)
2025-01-17 21:49:23 +00:00
}
return items, newState, nil
2025-01-17 21:49:23 +00:00
}
2025-01-30 07:14:00 +00:00
// Execute an action that takes an item as input and returns the item modified.
// This is basically just a wrapper over [Execute] that handles the input and backfilling.
2025-01-30 07:14:00 +00:00
func ExecuteItemAction(
item Item,
argv []string,
2025-01-30 07:51:06 +00:00
env []string,
state []byte,
2025-01-30 07:14:00 +00:00
timeout time.Duration,
) (Item, []byte, error) {
2025-01-30 07:14:00 +00:00
itemJson, err := json.Marshal(item)
if err != nil {
return Item{}, nil, fmt.Errorf("failed to serialize item: %v", err)
2025-01-30 07:14:00 +00:00
}
res, newState, err := Execute(item.Source, argv, env, state, string(itemJson), timeout)
2025-01-30 07:14:00 +00:00
if err != nil {
return Item{}, nil, fmt.Errorf("failed to execute action for %s/%s: %v", item.Source, item.Id, err)
2025-01-30 07:14:00 +00:00
}
if len(res) != 1 {
return Item{}, nil, fmt.Errorf("expected action to produce exactly one item, got %d", len(res))
2025-01-30 07:14:00 +00:00
}
newItem := res[0]
BackfillItem(&newItem, &item)
return newItem, newState, nil
2025-01-30 07:14:00 +00:00
}