go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/auth_service/internal/pubsub/client.go (about)

     1  // Copyright 2024 The LUCI Authors.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package pubsub
    16  
    17  import (
    18  	"context"
    19  
    20  	"cloud.google.com/go/iam"
    21  	"cloud.google.com/go/pubsub"
    22  	"google.golang.org/grpc/codes"
    23  	"google.golang.org/grpc/status"
    24  
    25  	"go.chromium.org/luci/common/errors"
    26  	"go.chromium.org/luci/common/logging"
    27  )
    28  
    29  // PubsubClient abstracts functionality to connect with Pubsub.
    30  //
    31  // Non-production implementations are used for unit testing.
    32  type PubsubClient interface {
    33  	// Close closes the connection to the Pubsub server.
    34  	Close() error
    35  
    36  	// GetIAMPolicy returns the IAM policy for the AuthDBChange topic.
    37  	GetIAMPolicy(ctx context.Context) (*iam.Policy, error)
    38  
    39  	// SetIAMPolicy sets the IAM policy for the AuthDBChange topic.
    40  	SetIAMPolicy(ctx context.Context, policy *iam.Policy) error
    41  
    42  	// Publish publishes the message to the AuthDBChange topic.
    43  	Publish(ctx context.Context, msg *pubsub.Message) error
    44  }
    45  
    46  type prodClient struct {
    47  	baseClient *pubsub.Client
    48  	projectID  string
    49  }
    50  
    51  // newProdClient creates a new production Pubsub client (not a mock).
    52  func newProdClient(ctx context.Context) (*prodClient, error) {
    53  	projectID := getProject(ctx)
    54  	client, err := pubsub.NewClient(ctx, projectID)
    55  	if err != nil {
    56  		return nil, errors.Annotate(err, "failed to create PubSub client for project %s", projectID).Err()
    57  	}
    58  
    59  	return &prodClient{
    60  		baseClient: client,
    61  		projectID:  projectID,
    62  	}, nil
    63  }
    64  
    65  func (c *prodClient) Close() error {
    66  	if c.baseClient != nil {
    67  		if err := c.baseClient.Close(); err != nil {
    68  			return errors.Annotate(err, "error closing PubSub client").Err()
    69  		}
    70  		c.baseClient = nil
    71  	}
    72  	return nil
    73  }
    74  
    75  func (c *prodClient) GetIAMPolicy(ctx context.Context) (*iam.Policy, error) {
    76  	if c.baseClient == nil {
    77  		return nil, status.Error(codes.Internal, "aborting - no PubSub client")
    78  	}
    79  
    80  	p, err := c.baseClient.Topic(AuthDBChangeTopicName).IAM().Policy(ctx)
    81  	if err != nil {
    82  		return nil, err
    83  	}
    84  
    85  	return p, nil
    86  }
    87  
    88  func (c *prodClient) SetIAMPolicy(ctx context.Context, policy *iam.Policy) error {
    89  	if c.baseClient == nil {
    90  		return status.Error(codes.Internal, "aborting - no PubSub client")
    91  	}
    92  
    93  	err := c.baseClient.Topic(AuthDBChangeTopicName).IAM().SetPolicy(ctx, policy)
    94  	if err != nil {
    95  		return err
    96  	}
    97  
    98  	return nil
    99  }
   100  
   101  func (c *prodClient) Publish(ctx context.Context, msg *pubsub.Message) (retErr error) {
   102  	if c.baseClient == nil {
   103  		return status.Error(codes.Internal, "aborting - no PubSub client")
   104  	}
   105  
   106  	topic := c.baseClient.Topic(AuthDBChangeTopicName)
   107  	ok, err := topic.Exists(ctx)
   108  	if err != nil {
   109  		return errors.Annotate(err, "error checking topic existence").Err()
   110  	}
   111  	if !ok {
   112  		// The topic doesn't exist; it must be created before we publish.
   113  		logging.Infof(ctx, "creating topic %s in project %s",
   114  			AuthDBChangeTopicName, c.projectID)
   115  		topic, err = c.baseClient.CreateTopic(ctx, AuthDBChangeTopicName)
   116  		if err != nil {
   117  			return errors.Annotate(err, "error creating topic %s in project %s",
   118  				AuthDBChangeTopicName, c.projectID).Err()
   119  		}
   120  	}
   121  
   122  	defer topic.Stop()
   123  	result := topic.Publish(ctx, msg)
   124  	if _, err := result.Get(ctx); err != nil {
   125  		switch status.Code(err) {
   126  		case codes.PermissionDenied:
   127  			return status.Errorf(codes.PermissionDenied,
   128  				"missing permission to publish PubSub message for project %s on topic %s",
   129  				c.projectID, AuthDBChangeTopicName)
   130  		default:
   131  			return status.Errorf(codes.Internal,
   132  				"error publishing Pubsub message: %s", err)
   133  		}
   134  	}
   135  
   136  	return nil
   137  }