github.com/mendersoftware/go-lib-micro@v0.0.0-20240304135804-e8e39c59b148/store/v2/utils.go (about)

     1  // Copyright 2023 Northern.tech AS
     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 store
    16  
    17  import (
    18  	"context"
    19  	"strings"
    20  
    21  	"go.mongodb.org/mongo-driver/bson"
    22  
    23  	"github.com/mendersoftware/go-lib-micro/identity"
    24  	mdoc "github.com/mendersoftware/go-lib-micro/mongo/doc"
    25  	v1 "github.com/mendersoftware/go-lib-micro/store"
    26  )
    27  
    28  const FieldTenantID = "tenant_id"
    29  
    30  // WithTenantID adds the tenant_id field to a bson document using the value extracted
    31  // from the identity of the context
    32  func WithTenantID(ctx context.Context, doc interface{}) bson.D {
    33  	var (
    34  		tenantID string
    35  		res      bson.D
    36  	)
    37  
    38  	identity := identity.FromContext(ctx)
    39  	if identity != nil {
    40  		tenantID = identity.Tenant
    41  	}
    42  	tenantElem := bson.E{Key: FieldTenantID, Value: tenantID}
    43  
    44  	switch v := doc.(type) {
    45  	case map[string]interface{}:
    46  		res = make(bson.D, 0, len(v)+1)
    47  		for k, v := range v {
    48  			res = append(res, bson.E{Key: k, Value: v})
    49  		}
    50  	case bson.M:
    51  		res = make(bson.D, 0, len(v)+1)
    52  		for k, v := range v {
    53  			res = append(res, bson.E{Key: k, Value: v})
    54  		}
    55  	case bson.D:
    56  		res = make(bson.D, len(v), len(v)+1)
    57  		copy(res, v)
    58  
    59  	case bson.Marshaler:
    60  		b, err := v.MarshalBSON()
    61  		if err != nil {
    62  			return nil
    63  		}
    64  		err = bson.Unmarshal(b, &res)
    65  		if err != nil {
    66  			return nil
    67  		}
    68  	default:
    69  		return mdoc.DocumentFromStruct(v, tenantElem)
    70  	}
    71  	res = append(res, tenantElem)
    72  
    73  	return res
    74  }
    75  
    76  // ArrayWithTenantID adds the tenant_id field to an array of bson documents
    77  // using the value extracted from the identity of the context
    78  func ArrayWithTenantID(ctx context.Context, doc bson.A) bson.A {
    79  	res := bson.A{}
    80  	for _, item := range doc {
    81  		res = append(res, WithTenantID(ctx, item))
    82  	}
    83  	return res
    84  }
    85  
    86  // DbFromContext generates database name using tenant field from identity extracted
    87  // from context and original database name
    88  func DbFromContext(ctx context.Context, origDbName string) string {
    89  	return origDbName
    90  }
    91  
    92  // IsTenantDb returns a function of `TenantDbMatchFunc` that can be used for
    93  // checking if database has a tenant DB name format
    94  func IsTenantDb(baseDb string) v1.TenantDbMatchFunc {
    95  	prefix := baseDb + "-"
    96  	return func(name string) bool {
    97  		return strings.HasPrefix(name, prefix)
    98  	}
    99  }
   100  
   101  // TenantFromDbName attempts to extract tenant ID from provided tenant DB name.
   102  // Returns extracted tenant ID or an empty string.
   103  func TenantFromDbName(dbName string, baseDb string) string {
   104  	noBase := strings.TrimPrefix(dbName, baseDb+"-")
   105  	if noBase == dbName {
   106  		return ""
   107  	}
   108  	return noBase
   109  }
   110  
   111  // DbNameForTenant composes tenant's db name.
   112  func DbNameForTenant(tenantId string, baseDb string) string {
   113  	return baseDb
   114  }