snapshot of https://github.com/mattermost/mattermost.git (master, 2026-09-16; blobs > 131072B omitted)

server/platform/services/awsmeter/awsmeter.go

@@ -0,0 +1,156 @@
+// Copyright (c) 2015-present Mattermost, Inc. All Rights Reserved.
+// See LICENSE.txt for license information.
+
+package awsmeter
+
+import (
+	"context"
+	"encoding/json"
+	"os"
+	"time"
+
+	"github.com/aws/aws-sdk-go-v2/aws"
+	"github.com/aws/aws-sdk-go-v2/config"
+	"github.com/aws/aws-sdk-go-v2/credentials/ec2rolecreds"
+	"github.com/aws/aws-sdk-go-v2/feature/ec2/imds"
+	"github.com/aws/aws-sdk-go-v2/service/marketplacemetering"
+	"github.com/pkg/errors"
+
+	"github.com/mattermost/mattermost/server/public/model"
+	"github.com/mattermost/mattermost/server/public/shared/mlog"
+	"github.com/mattermost/mattermost/server/v8/channels/store"
+)
+
+type AwsMeter struct {
+	store   store.Store
+	service *AWSMeterService
+	config  *model.Config
+}
+
+type MeteringService interface {
+	MeterUsage(ctx context.Context, input *marketplacemetering.MeterUsageInput, optFns ...func(*marketplacemetering.Options)) (*marketplacemetering.MeterUsageOutput, error)
+}
+
+type AWSMeterService struct {
+	AwsDryRun      bool
+	AwsProductCode string
+	AwsMeteringSvc MeteringService
+}
+
+type AWSMeterReport struct {
+	Dimension string    `json:"dimension"`
+	Value     int32     `json:"value"`
+	Timestamp time.Time `json:"timestamp"`
+}
+
+func (o *AWSMeterReport) ToJSON() string {
+	b, _ := json.Marshal(o)
+	return string(b)
+}
+
+func New(ctx context.Context, store store.Store, config *model.Config) *AwsMeter {
+	svc := &AWSMeterService{
+		AwsDryRun:      false,
+		AwsProductCode: "12345", //TODO
+	}
+
+	service, err := newAWSMarketplaceMeteringService(ctx)
+	if err != nil {
+		mlog.Debug("Could not create AWS metering service", mlog.String("error", err.Error()))
+		return nil
+	}
+
+	svc.AwsMeteringSvc = service
+	return &AwsMeter{
+		store:   store,
+		service: svc,
+		config:  config,
+	}
+}
+
+func newAWSMarketplaceMeteringService(ctx context.Context) (*marketplacemetering.Client, error) {
+	region := os.Getenv("AWS_REGION")
+	imdsClient := imds.New(imds.Options{})
+	credsProvider := ec2rolecreds.New(func(o *ec2rolecreds.Options) { o.Client = imdsClient })
+
+	creds := aws.NewCredentialsCache(credsProvider)
+
+	_, err := creds.Retrieve(ctx)
+	if err != nil {
+		return nil, errors.Wrap(err, "cannot obtain credentials")
+	}
+
+	cfg, err := config.LoadDefaultConfig(ctx,
+		config.WithRegion(region),
+		config.WithCredentialsProvider(creds),
+	)
+	if err != nil {
+		return nil, errors.Wrap(err, "failed to load default config with EC2 credentials provider")
+	}
+
+	return marketplacemetering.NewFromConfig(cfg), nil
+}
+
+// a report entry is for all metrics
+func (awsm *AwsMeter) GetUserCategoryUsage(dimensions []string, startTime time.Time, endTime time.Time) []*AWSMeterReport {
+	reports := make([]*AWSMeterReport, 0)
+
+	for _, dimension := range dimensions {
+		var userCount int32
+		var err error
+
+		switch dimension {
+		case model.AwsMeteringDimensionUsageHrs:
+			userCount, err = awsm.store.User().AnalyticsActiveCountForPeriod(model.GetMillisForTime(startTime), model.GetMillisForTime(endTime), model.UserCountOptions{})
+			if err != nil {
+				mlog.Warn("Failed to obtain usage data", mlog.String("dimension", dimension), mlog.Time("start", startTime), mlog.Int("count", userCount), mlog.Err(err))
+				continue
+			}
+		default:
+			mlog.Debug("Dimension does not exist!", mlog.String("dimension", dimension))
+			continue
+		}
+
+		report := &AWSMeterReport{
+			Dimension: dimension,
+			Value:     userCount,
+			Timestamp: startTime,
+		}
+
+		reports = append(reports, report)
+	}
+
+	return reports
+}
+
+func (awsm *AwsMeter) ReportUserCategoryUsage(ctx context.Context, reports []*AWSMeterReport) error {
+	for _, report := range reports {
+		err := sendReportToMeteringService(ctx, awsm.service, report)
+		if err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+func sendReportToMeteringService(ctx context.Context, ams *AWSMeterService, report *AWSMeterReport) error {
+	params := &marketplacemetering.MeterUsageInput{
+		DryRun:         aws.Bool(ams.AwsDryRun),
+		ProductCode:    aws.String(ams.AwsProductCode),
+		UsageDimension: aws.String(report.Dimension),
+		UsageQuantity:  aws.Int32(report.Value),
+		Timestamp:      aws.Time(report.Timestamp),
+	}
+
+	resp, err := ams.AwsMeteringSvc.MeterUsage(ctx, params)
+	if err != nil {
+		return errors.Wrap(err, "Invalid metering service id.")
+	}
+	if resp.MeteringRecordId == nil {
+		return errors.Wrap(err, "Invalid metering service id.")
+	}
+
+	mlog.Debug("Sent record to AWS metering service", mlog.String("dimension", report.Dimension), mlog.Int("value", report.Value), mlog.String("timestamp", report.Timestamp.String()))
+
+	return nil
+}

server/platform/services/awsmeter/awsmeter_test.go

@@ -0,0 +1,146 @@
+// Copyright (c) 2015-present Mattermost, Inc. All Rights Reserved.
+// See LICENSE.txt for license information.
+
+package awsmeter
+
+import (
+	"context"
+	"errors"
+	"testing"
+	"time"
+
+	"github.com/aws/aws-sdk-go-v2/service/marketplacemetering"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/mattermost/mattermost/server/public/model"
+	"github.com/mattermost/mattermost/server/public/plugin/plugintest/mock"
+	"github.com/mattermost/mattermost/server/v8/channels/store/storetest/mocks"
+)
+
+type mockMarketplaceMeteringClient struct {
+	MeteringService
+}
+
+func (m *mockMarketplaceMeteringClient) MeterUsage(ctx context.Context, input *marketplacemetering.MeterUsageInput, optFns ...func(*marketplacemetering.Options)) (*marketplacemetering.MeterUsageOutput, error) {
+	return &marketplacemetering.MeterUsageOutput{
+		MeteringRecordId: new("1"),
+	}, nil
+}
+
+type mockMarketplaceMeteringClientWithError struct {
+	MeteringService
+}
+
+func (m *mockMarketplaceMeteringClientWithError) MeterUsage(ctx context.Context, input *marketplacemetering.MeterUsageInput, optFns ...func(*marketplacemetering.Options)) (*marketplacemetering.MeterUsageOutput, error) {
+	return nil, errors.New("error")
+}
+
+func TestAwsMeterUsage(t *testing.T) {
+	startTime := time.Now()
+	endTime := time.Now()
+	dimensions := []string{model.AwsMeteringDimensionUsageHrs}
+
+	userStoreMock := mocks.UserStore{}
+	userStoreMock.On("AnalyticsActiveCountForPeriod", model.GetMillisForTime(startTime), model.GetMillisForTime(endTime), mock.AnythingOfType("model.UserCountOptions")).Return(int32(2), nil)
+
+	storeMock := mocks.Store{}
+	storeMock.On("User").Return(&userStoreMock)
+
+	reports := make([]*AWSMeterReport, 1)
+	reports[0] = &AWSMeterReport{
+		Dimension: model.AwsMeteringDimensionUsageHrs,
+		Value:     2,
+		Timestamp: startTime,
+	}
+
+	// Define a mock struct to be used in your unit tests of myFunc.
+	svc := &AWSMeterService{
+		AwsDryRun:      false,
+		AwsProductCode: "12345",
+		AwsMeteringSvc: &mockMarketplaceMeteringClient{},
+	}
+
+	config := &model.Config{}
+	config.SetDefaults()
+
+	awsmeter := &AwsMeter{
+		store:   &storeMock,
+		service: svc,
+		config:  config,
+	}
+
+	t.Run("Send report for one usage category", func(t *testing.T) {
+		resultReports := awsmeter.GetUserCategoryUsage(dimensions, startTime, endTime)
+		require.NotNil(t, resultReports)
+		assert.Equal(t, 1, len(resultReports))
+		assert.Equal(t, reports[0].Dimension, resultReports[0].Dimension)
+		assert.Equal(t, reports[0].Value, resultReports[0].Value)
+		assert.Equal(t, reports[0].Timestamp, resultReports[0].Timestamp)
+
+		err := awsmeter.ReportUserCategoryUsage(context.Background(), resultReports)
+		require.NoError(t, err)
+	})
+
+	t.Run("Error in AWS service call", func(t *testing.T) {
+		awsmeter.service.AwsMeteringSvc = &mockMarketplaceMeteringClientWithError{}
+		resultReports := awsmeter.GetUserCategoryUsage(dimensions, startTime, endTime)
+		require.NotNil(t, resultReports)
+		assert.Equal(t, 1, len(resultReports))
+		err := awsmeter.ReportUserCategoryUsage(context.Background(), resultReports)
+		require.Error(t, err)
+	})
+
+	t.Run("Invalid dimension", func(t *testing.T) {
+		awsmeter.service.AwsMeteringSvc = &mockMarketplaceMeteringClient{}
+		dimensions = []string{"invalid dimension"}
+		resultReports := awsmeter.GetUserCategoryUsage(dimensions, startTime, endTime)
+		require.NotNil(t, resultReports)
+		assert.Equal(t, 0, len(resultReports))
+		err := awsmeter.ReportUserCategoryUsage(context.Background(), resultReports)
+		require.NoError(t, err)
+	})
+}
+
+func TestAwsMeterUsageWithDBError(t *testing.T) {
+	startTime := time.Now()
+	endTime := time.Now()
+	dimensions := []string{model.AwsMeteringDimensionUsageHrs}
+
+	userStoreMock := mocks.UserStore{}
+	userStoreMock.On("AnalyticsActiveCountForPeriod", model.GetMillisForTime(startTime), model.GetMillisForTime(endTime), mock.AnythingOfType("model.UserCountOptions")).Return(int32(0), errors.New("error"))
+
+	storeMock := mocks.Store{}
+	storeMock.On("User").Return(&userStoreMock)
+
+	reports := make([]*AWSMeterReport, 1)
+	reports[0] = &AWSMeterReport{
+		Dimension: model.AwsMeteringDimensionUsageHrs,
+		Value:     2,
+		Timestamp: startTime,
+	}
+
+	// Define a mock struct to be used in your unit tests of myFunc.
+	svc := &AWSMeterService{
+		AwsDryRun:      false,
+		AwsProductCode: "12345",
+		AwsMeteringSvc: &mockMarketplaceMeteringClient{},
+	}
+
+	config := &model.Config{}
+	config.SetDefaults()
+
+	awsmeter := &AwsMeter{
+		store:   &storeMock,
+		service: svc,
+		config:  config,
+	}
+
+	t.Run("Error in DB query", func(t *testing.T) {
+		resultReports := awsmeter.GetUserCategoryUsage(dimensions, startTime, endTime)
+		require.NotNil(t, resultReports)
+		assert.Equal(t, 0, len(resultReports))
+		err := awsmeter.ReportUserCategoryUsage(context.Background(), resultReports)
+		require.NoError(t, err)
+	})
+}