<commit_before>/***** BEGIN LICENSE BLOCK *****
# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this file,
# You can obtain one at http://mozilla.org/MPL/2.0/.
# ***** END LICENSE BLOCK *****/

package s3splitfile

import (
	"fmt"
	"github.com/AdRoll/goamz/aws"
	"github.com/AdRoll/goamz/s3"
	"github.com/mozilla-services/heka/message"
	"github.com/mozilla-services/heka/pipeline"
	"io"
	"sync"
	"sync/atomic"
	"time"
)

type S3SplitFileInput struct {
	processFileCount          int64
	processFileFailures       int64
	processFileDiscardedBytes int64
	processMessageCount       int64
	processMessageFailures    int64
	processMessageBytes       int64

	*S3SplitFileInputConfig
	bucket   *s3.Bucket
	schema   Schema
	stop     chan bool
	listChan chan string
}

type S3SplitFileInputConfig struct {
	// So we can default to using ProtobufDecoder.
	Decoder string
	// So we can default to using HekaFramingSplitter.
	Splitter string

	SchemaFile     string `toml:"schema_file"`
	AWSKey         string `toml:"aws_key"`
	AWSSecretKey   string `toml:"aws_secret_key"`
	AWSRegion      string `toml:"aws_region"`
	S3Bucket       string `toml:"s3_bucket"`
	S3BucketPrefix string `toml:"s3_bucket_prefix"`
	S3Retries      uint32 `toml:"s3_retries"`
	S3WorkerCount  uint32 `toml:"s3_worker_count"`
}

func (input *S3SplitFileInput) ConfigStruct() interface{} {
	return &S3SplitFileInputConfig{
		Decoder:        "ProtobufDecoder",
		Splitter:       "HekaFramingSplitter",
		AWSKey:         "",
		AWSSecretKey:   "",
		AWSRegion:      "us-west-2",
		S3Bucket:       "",
		S3BucketPrefix: "",
		S3Retries:      5,
		S3WorkerCount:  10,
	}
}

func (input *S3SplitFileInput) Init(config interface{}) (err error) {
	conf := config.(*S3SplitFileInputConfig)
	input.S3SplitFileInputConfig = conf

	input.schema, err = LoadSchema(conf.SchemaFile)
	if err != nil {
		return fmt.Errorf("Parameter 'schema_file' must be a valid JSON file: %s", err)
	}

	if conf.S3Bucket != "" {
		auth, err := aws.GetAuth(conf.AWSKey, conf.AWSSecretKey, "", time.Now())
		if err != nil {
			return fmt.Errorf("Authentication error: %s\n", err)
		}
		region, ok := aws.Regions[conf.AWSRegion]
		if !ok {
			return fmt.Errorf("Parameter 'aws_region' must be a valid AWS Region")
		}
		s := s3.New(auth, region)
		// TODO: ensure we can read from the bucket.
		input.bucket = s.Bucket(conf.S3Bucket)
	} else {
		input.bucket = nil
	}

	// Remove any excess path separators from the bucket prefix.
	conf.S3BucketPrefix = CleanBucketPrefix(conf.S3BucketPrefix)

	input.stop = make(chan bool)
	input.listChan = make(chan string, 1000)

	return nil
}

func (input *S3SplitFileInput) Stop() {
	close(input.stop)
}

func (input *S3SplitFileInput) Run(runner pipeline.InputRunner, helper pipeline.PluginHelper) error {
	// Begin listing the files (either straight from S3 or from a cache)
	// Write matching filenames on a "lister" channel
	// Read from the lister channel:
	//   - fetch the filename
	//   - read records from it
	//   - write them to a "reader" channel

	var (
		wg sync.WaitGroup
		i  uint32
	)

	wg.Add(1)
	go func() {
		runner.LogMessage("Starting S3 list")
		for r := range S3Iterator(input.bucket, input.S3BucketPrefix, input.schema) {
			if r.Err != nil {
				runner.LogError(fmt.Errorf("Error getting S3 list: %s", r.Err))
			} else {
				runner.LogMessage(fmt.Sprintf("Found: %s", r.Key.Key))
				input.listChan <- r.Key.Key
			}
		}
		// All done listing, close the channel
		runner.LogMessage("All done listing. Closing channel")
		close(input.listChan)
		wg.Done()
	}()

	// Run a pool of concurrent readers.
	for i = 0; i < input.S3WorkerCount; i++ {
		wg.Add(1)
		go input.fetcher(runner, &wg, i)
	}
	wg.Wait()

	return nil
}

// TODO: handle "no such file"<commit_msg>Support filtering by object basename prefix in S3SplitFileInput<commit_after>/***** BEGIN LICENSE BLOCK *****
# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this file,
# You can obtain one at http://mozilla.org/MPL/2.0/.
# ***** END LICENSE BLOCK *****/

package s3splitfile

import (
	"fmt"
	"github.com/AdRoll/goamz/aws"
	"github.com/AdRoll/goamz/s3"
	"github.com/mozilla-services/heka/message"
	"github.com/mozilla-services/heka/pipeline"
	"io"
	"strings"
	"sync"
	"sync/atomic"
	"time"
)

type S3SplitFileInput struct {
	processFileCount          int64
	processFileFailures       int64
	processFileDiscardedBytes int64
	processMessageCount       int64
	processMessageFailures    int64
	processMessageBytes       int64

	*S3SplitFileInputConfig
	bucket   *s3.Bucket
	schema   Schema
	stop     chan bool
	listChan chan string
}

type S3SplitFileInputConfig struct {
	// So we can default to using ProtobufDecoder.
	Decoder string
	// So we can default to using HekaFramingSplitter.
	Splitter string

	SchemaFile             string `toml:"schema_file"`
	AWSKey                 string `toml:"aws_key"`
	AWSSecretKey           string `toml:"aws_secret_key"`
	AWSRegion              string `toml:"aws_region"`
	S3Bucket               string `toml:"s3_bucket"`
	S3BucketPrefix         string `toml:"s3_bucket_prefix"`
	S3ObjectBasenamePrefix string `toml:"s3_object_basename_prefix"`
	S3Retries              uint32 `toml:"s3_retries"`
	S3WorkerCount          uint32 `toml:"s3_worker_count"`
}

func (input *S3SplitFileInput) ConfigStruct() interface{} {
	return &S3SplitFileInputConfig{
		Decoder:                "ProtobufDecoder",
		Splitter:               "HekaFramingSplitter",
		AWSKey:                 "",
		AWSSecretKey:           "",
		AWSRegion:              "us-west-2",
		S3Bucket:               "",
		S3BucketPrefix:         "",
		S3ObjectBasenamePrefix: "",
		S3Retries:              5,
		S3WorkerCount:          10,
	}
}

func (input *S3SplitFileInput) Init(config interface{}) (err error) {
	conf := config.(*S3SplitFileInputConfig)
	input.S3SplitFileInputConfig = conf

	input.schema, err = LoadSchema(conf.SchemaFile)
	if err != nil {
		return fmt.Errorf("Parameter 'schema_file' must be a valid JSON file: %s", err)
	}

	if conf.S3Bucket != "" {
		auth, err := aws.GetAuth(conf.AWSKey, conf.AWSSecretKey, "", time.Now())
		if err != nil {
			return fmt.Errorf("Authentication error: %s\n", err)
		}
		region, ok := aws.Regions[conf.AWSRegion]
		if !ok {
			return fmt.Errorf("Parameter 'aws_region' must be a valid AWS Region")
		}
		s := s3.New(auth, region)
		// TODO: ensure we can read from the bucket.
		input.bucket = s.Bucket(conf.S3Bucket)
	} else {
		input.bucket = nil
	}

	// Remove any excess path separators from the bucket prefix.
	conf.S3BucketPrefix = CleanBucketPrefix(conf.S3BucketPrefix)

	input.stop = make(chan bool)
	input.listChan = make(chan string, 1000)

	return nil
}

func (input *S3SplitFileInput) Stop() {
	close(input.stop)
}

func (input *S3SplitFileInput) Run(runner pipeline.InputRunner, helper pipeline.PluginHelper) error {
	// Begin listing the files (either straight from S3 or from a cache)
	// Write matching filenames on a "lister" channel
	// Read from the lister channel:
	//   - fetch the filename
	//   - read records from it
	//   - write them to a "reader" channel

	var (
		wg sync.WaitGroup
		i  uint32
	)

	wg.Add(1)
	go func() {
		runner.LogMessage("Starting S3 list")
		for r := range S3Iterator(input.bucket, input.S3BucketPrefix, input.schema) {
			if r.Err != nil {
				runner.LogError(fmt.Errorf("Error getting S3 list: %s", r.Err))
			} else {
				basename := r.Key.Key[strings.LastIndex(r.Key.Key, "/")+1:]
				if strings.HasPrefix(basename, input.S3ObjectBasenamePrefix) {
					runner.LogMessage(fmt.Sprintf("Found: %s", r.Key.Key))
					input.listChan <- r.Key.Key
				} else {
					runner.LogMessage(fmt.Sprintf("Skipping: %s", r.Key.Key))
				}
			}
		}
		// All done listing, close the channel
		runner.LogMessage("All done listing. Closing channel")
		close(input.listChan)
		wg.Done()
	}()

	// Run a pool of concurrent readers.
	for i = 0; i < input.S3WorkerCount; i++ {
		wg.Add(1)
		go input.fetcher(runner, &wg, i)
	}
	wg.Wait()

	return nil
}

// TODO: handle "no such file"
{'fraction_non_alphanumeric': 0.1187683284457478, 'fraction_numerical': 0.017595307917888565, 'mean_word_length': 3.2853403141361257, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5991781', 'n_tokens_mistral': 3030, 'n_tokens_neox': 2805, 'n_words': 737}