github.com/haalcala/mattermost-server-change-repo@v0.0.0-20210713015153-16753fbeee5f/services/awsmeter/awsmeter.go (about)

     1  // Copyright (c) 2015-present Mattermost, Inc. All Rights Reserved.
     2  // See LICENSE.txt for license information.
     3  
     4  package awsmeter
     5  
     6  import (
     7  	"encoding/json"
     8  	"os"
     9  	"time"
    10  
    11  	"github.com/aws/aws-sdk-go/aws"
    12  	"github.com/aws/aws-sdk-go/aws/credentials"
    13  	"github.com/aws/aws-sdk-go/aws/credentials/ec2rolecreds"
    14  	"github.com/aws/aws-sdk-go/aws/ec2metadata"
    15  	"github.com/aws/aws-sdk-go/aws/session"
    16  	"github.com/aws/aws-sdk-go/service/marketplacemetering"
    17  	"github.com/aws/aws-sdk-go/service/marketplacemetering/marketplacemeteringiface"
    18  	"github.com/pkg/errors"
    19  
    20  	"github.com/mattermost/mattermost-server/v5/mlog"
    21  	"github.com/mattermost/mattermost-server/v5/model"
    22  	"github.com/mattermost/mattermost-server/v5/store"
    23  )
    24  
    25  type AwsMeter struct {
    26  	store   store.Store
    27  	service *AWSMeterService
    28  	config  *model.Config
    29  }
    30  
    31  type AWSMeterService struct {
    32  	AwsDryRun      bool
    33  	AwsProductCode string
    34  	AwsMeteringSvc marketplacemeteringiface.MarketplaceMeteringAPI
    35  }
    36  
    37  type AWSMeterReport struct {
    38  	Dimension string    `json:"dimension"`
    39  	Value     int64     `json:"value"`
    40  	Timestamp time.Time `json:"timestamp"`
    41  }
    42  
    43  func (o *AWSMeterReport) ToJSON() string {
    44  	b, _ := json.Marshal(o)
    45  	return string(b)
    46  }
    47  
    48  func New(store store.Store, config *model.Config) *AwsMeter {
    49  	svc := &AWSMeterService{
    50  		AwsDryRun:      false,
    51  		AwsProductCode: "12345", //TODO
    52  	}
    53  
    54  	service, err := newAWSMarketplaceMeteringService()
    55  	if err != nil {
    56  		mlog.Debug("Could not create AWS metering service", mlog.String("error", err.Error()))
    57  		return nil
    58  	}
    59  
    60  	svc.AwsMeteringSvc = service
    61  	return &AwsMeter{
    62  		store:   store,
    63  		service: svc,
    64  		config:  config,
    65  	}
    66  }
    67  
    68  func newAWSMarketplaceMeteringService() (*marketplacemetering.MarketplaceMetering, error) {
    69  	region := os.Getenv("AWS_REGION")
    70  	s, err := session.NewSession(&aws.Config{Region: &region})
    71  	if err != nil {
    72  		return nil, err
    73  	}
    74  
    75  	creds := credentials.NewChainCredentials(
    76  		[]credentials.Provider{
    77  			&ec2rolecreds.EC2RoleProvider{
    78  				Client: ec2metadata.New(s),
    79  			},
    80  		})
    81  
    82  	_, err = creds.Get()
    83  	if err != nil {
    84  		return nil, errors.Wrap(err, "cannot obtain credentials")
    85  	}
    86  
    87  	return marketplacemetering.New(session.Must(session.NewSession(&aws.Config{
    88  		Credentials: creds,
    89  	}))), nil
    90  }
    91  
    92  // a report entry is for all metrics
    93  func (awsm *AwsMeter) GetUserCategoryUsage(dimensions []string, startTime time.Time, endTime time.Time) []*AWSMeterReport {
    94  	reports := make([]*AWSMeterReport, 0)
    95  
    96  	for _, dimension := range dimensions {
    97  		var userCount int64
    98  		var err error
    99  
   100  		switch dimension {
   101  		case model.AWS_METERING_DIMENSION_USAGE_HRS:
   102  			userCount, err = awsm.store.User().AnalyticsActiveCountForPeriod(model.GetMillisForTime(startTime), model.GetMillisForTime(endTime), model.UserCountOptions{})
   103  			if err != nil {
   104  				mlog.Warn("Failed to obtain usage data", mlog.String("dimension", dimension), mlog.String("start", startTime.String()), mlog.Int64("count", userCount), mlog.Err(err))
   105  				continue
   106  			}
   107  		default:
   108  			mlog.Debug("Dimension does not exist!", mlog.String("dimension", dimension))
   109  			continue
   110  		}
   111  
   112  		report := &AWSMeterReport{
   113  			Dimension: dimension,
   114  			Value:     userCount,
   115  			Timestamp: startTime,
   116  		}
   117  
   118  		reports = append(reports, report)
   119  	}
   120  
   121  	return reports
   122  }
   123  
   124  func (awsm *AwsMeter) ReportUserCategoryUsage(reports []*AWSMeterReport) error {
   125  	for _, report := range reports {
   126  		err := sendReportToMeteringService(awsm.service, report)
   127  		if err != nil {
   128  			return err
   129  		}
   130  	}
   131  	return nil
   132  }
   133  
   134  func sendReportToMeteringService(ams *AWSMeterService, report *AWSMeterReport) error {
   135  	params := &marketplacemetering.MeterUsageInput{
   136  		DryRun:         aws.Bool(ams.AwsDryRun),
   137  		ProductCode:    aws.String(ams.AwsProductCode),
   138  		UsageDimension: aws.String(report.Dimension),
   139  		UsageQuantity:  aws.Int64(report.Value),
   140  		Timestamp:      aws.Time(report.Timestamp),
   141  	}
   142  
   143  	resp, err := ams.AwsMeteringSvc.MeterUsage(params)
   144  	if err != nil {
   145  		return errors.Wrap(err, "Invalid metering service id.")
   146  	}
   147  	if resp.MeteringRecordId == nil {
   148  		return errors.Wrap(err, "Invalid metering service id.")
   149  	}
   150  
   151  	mlog.Debug("Sent record to AWS metering service", mlog.String("dimension", report.Dimension), mlog.Int64("value", report.Value), mlog.String("timestamp", report.Timestamp.String()))
   152  
   153  	return nil
   154  }