<commit_before>/*
 * Copyright (c) 2016 <NAME>
 *
 * 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/.
 *
 */
package main

import (
	"context"
	log "github.com/Sirupsen/logrus"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/credentials/stscreds"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/firehose"
	"math"
	"sync"
	"time"
)

var (
	REQUEST_SIZE_LIMIT int = 2e+6 // 2MB
	BATCH_LIMIT            = 500
)

type FirehoseStream struct {
	svc          *firehose.Firehose
	stream       string
	dataChan     chan []byte
	interval     time.Duration
	recordFormat []Attribute
	ctx          context.Context
	wg           sync.WaitGroup
}

func NewFirehoseStream(ctx context.Context, recordFormat []Attribute, accessKey, secretAccessKey, awsRegion, awsSTSRole, streamName string) *FirehoseStream {

	if awsRegion == "" {
		log.Fatal("Please Specify the region your firehose is.")
	}

	s := &FirehoseStream{}
	sess := &session.Session{}
	awsConfig := &aws.Config{Region: aws.String(awsRegion)}

	if awsSTSRole != "" {
		sess = session.Must(session.NewSession())
		creds := stscreds.NewCredentials(sess, awsSTSRole)
		awsConfig.Credentials = creds
	} else if accessKey == "" || secretAccessKey == "" {
		// try IAM
		sess = session.New(nil)
	} else {
		creds := credentials.NewStaticCredentials(accessKey, secretAccessKey, "")
		config := &aws.Config{
			Region:      aws.String(awsRegion),
			Credentials: creds,
		}
		sess = session.New(config)
	}

	s.svc = firehose.New(sess, awsConfig)
	s.stream = streamName
	s.dataChan = make(chan []byte, BATCH_LIMIT*5)
	s.interval = 5 * time.Second
	s.recordFormat = recordFormat
	s.ctx = ctx

	go s.intervalStreamer()

	return s
}

func (s *FirehoseStream) Stream(r *Record) error {
	s.dataChan <- r.RecordToCSV()
	return nil
}

func (s *FirehoseStream) Close() {<commit_msg>fixes #4<commit_after>/*
 * Copyright (c) 2016 <NAME>
 *
 * 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/.
 *
 */
package main

import (
	"context"
	log "github.com/Sirupsen/logrus"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/credentials/stscreds"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/firehose"
	"math"
	"sync"
	"time"
)

var (
	REQUEST_SIZE_LIMIT int = 2e+6 // 2MB
	BATCH_LIMIT            = 500
)

type FirehoseStream struct {
	svc          *firehose.Firehose
	stream       string
	dataChan     chan []byte
	interval     time.Duration
	recordFormat []Attribute
	ctx          context.Context
	wg           sync.WaitGroup
}

func NewFirehoseStream(ctx context.Context, recordFormat []Attribute, accessKey, secretAccessKey, awsRegion, awsSTSRole, streamName string) *FirehoseStream {

	if awsRegion == "" {
		log.Fatal("Please Specify the region your firehose is.")
	}

	s := &FirehoseStream{}
	sess := &session.Session{}
	awsConfig := &aws.Config{Region: aws.String(awsRegion)}

	if awsSTSRole != "" {
		sess = session.Must(session.NewSession())
		creds := stscreds.NewCredentials(sess, awsSTSRole)
		awsConfig.Credentials = creds
	} else if accessKey == "" || secretAccessKey == "" {
		// try IAM
		sess = session.New(nil)
	} else {
		creds := credentials.NewStaticCredentials(accessKey, secretAccessKey, "")
		config := &aws.Config{
			Region:      aws.String(awsRegion),
			Credentials: creds,
		}
		sess = session.New(config)
	}

	s.svc = firehose.New(sess, awsConfig)
	s.stream = streamName
	s.dataChan = make(chan []byte, BATCH_LIMIT*5)
	s.interval = 5 * time.Second
	s.recordFormat = recordFormat
	s.ctx = ctx

	s.wg.Add(1)
	go s.intervalStreamer()

	return s
}

func (s *FirehoseStream) Stream(r *Record) error {
	s.dataChan <- r.RecordToCSV()
	return nil
}

func (s *FirehoseStream) Close() {
{'fraction_non_alphanumeric': 0.13266796494644595, 'fraction_numerical': 0.008276533592989289, 'mean_word_length': 3.4956236323851204, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 5, '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': '312956', 'n_tokens_mistral': 1565, 'n_tokens_neox': 1509, 'n_words': 340}