github.com/pingcap/tiflow@v0.0.0-20240520035814-5bf52d54e205/pkg/sink/codec/avro/glue_client.go (about)

     1  // Copyright 2023 PingCAP, Inc.
     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  // See the License for the specific language governing permissions and
    12  // limitations under the License.
    13  
    14  package avro
    15  
    16  import (
    17  	"context"
    18  
    19  	"github.com/aws/aws-sdk-go-v2/service/glue"
    20  	"github.com/aws/aws-sdk-go-v2/service/glue/types"
    21  	"github.com/google/uuid"
    22  	"github.com/pingcap/errors"
    23  )
    24  
    25  // glueClient is a partial interface of glue client, used to mock glue client in unit test
    26  type glueClient interface {
    27  	GetRegistry(ctx context.Context,
    28  		params *glue.GetRegistryInput,
    29  		optFns ...func(*glue.Options)) (*glue.GetRegistryOutput, error)
    30  	CreateSchema(ctx context.Context,
    31  		params *glue.CreateSchemaInput,
    32  		optFns ...func(*glue.Options)) (*glue.CreateSchemaOutput, error)
    33  	RegisterSchemaVersion(ctx context.Context,
    34  		params *glue.RegisterSchemaVersionInput,
    35  		optFns ...func(*glue.Options)) (*glue.RegisterSchemaVersionOutput, error)
    36  	GetSchemaVersion(ctx context.Context,
    37  		params *glue.GetSchemaVersionInput,
    38  		optFns ...func(*glue.Options)) (*glue.GetSchemaVersionOutput, error)
    39  }
    40  
    41  type mockGlueClientImpl struct {
    42  	createSchemaInput           map[string]*glue.CreateSchemaInput          // schemaName -> schema
    43  	registerSchemaVersionsInput map[string]*glue.RegisterSchemaVersionInput // schemaName -> schema
    44  }
    45  
    46  func newMockGlueClientImpl() *mockGlueClientImpl {
    47  	return &mockGlueClientImpl{
    48  		createSchemaInput:           make(map[string]*glue.CreateSchemaInput),
    49  		registerSchemaVersionsInput: make(map[string]*glue.RegisterSchemaVersionInput),
    50  	}
    51  }
    52  
    53  func (m *mockGlueClientImpl) CreateSchema(
    54  	ctx context.Context,
    55  	params *glue.CreateSchemaInput,
    56  	optFns ...func(*glue.Options),
    57  ) (*glue.CreateSchemaOutput, error) {
    58  	id := uuid.New()
    59  	sid := id.String()
    60  	// here use Description to store schemaVersionID for test purpose
    61  	params.Description = &sid
    62  	m.createSchemaInput[*params.SchemaName] = params
    63  	m.registerSchemaVersionsInput[sid] = &glue.RegisterSchemaVersionInput{
    64  		SchemaId: &types.SchemaId{
    65  			SchemaName: params.SchemaName,
    66  			// here use schemaArn as schemaVersionId for test purpose
    67  			SchemaArn: &sid,
    68  		},
    69  		SchemaDefinition: params.SchemaDefinition,
    70  	}
    71  	return &glue.CreateSchemaOutput{
    72  		SchemaVersionId: &sid,
    73  	}, nil
    74  }
    75  
    76  func (m *mockGlueClientImpl) RegisterSchemaVersion(
    77  	ctx context.Context,
    78  	params *glue.RegisterSchemaVersionInput,
    79  	optFns ...func(*glue.Options),
    80  ) (*glue.RegisterSchemaVersionOutput, error) {
    81  	old, ok := m.createSchemaInput[*params.SchemaId.SchemaName]
    82  	if !ok {
    83  		return nil, errors.New("EntityNotFoundException")
    84  	}
    85  
    86  	// if schema definition is the same, don't create new schema version
    87  	if *old.SchemaDefinition == *params.SchemaDefinition {
    88  		return &glue.RegisterSchemaVersionOutput{
    89  			SchemaVersionId: old.Description,
    90  		}, nil
    91  	}
    92  
    93  	m.createSchemaInput[*params.SchemaId.SchemaName] = &glue.CreateSchemaInput{
    94  		SchemaDefinition: params.SchemaDefinition,
    95  	}
    96  	id := uuid.New()
    97  	sid := id.String()
    98  	params.SchemaId.SchemaArn = &sid
    99  	m.registerSchemaVersionsInput[*params.SchemaId.SchemaArn] = params
   100  	return &glue.RegisterSchemaVersionOutput{
   101  		SchemaVersionId: params.SchemaId.SchemaArn,
   102  	}, nil
   103  }
   104  
   105  func (m *mockGlueClientImpl) GetSchemaVersion(
   106  	ctx context.Context,
   107  	params *glue.GetSchemaVersionInput,
   108  	optFns ...func(*glue.Options),
   109  ) (*glue.GetSchemaVersionOutput, error) {
   110  	if params.SchemaVersionId == nil {
   111  		res, ok := m.createSchemaInput[*params.SchemaId.SchemaName]
   112  		if !ok {
   113  			return nil, errors.New("EntityNotFoundException")
   114  		}
   115  		return &glue.GetSchemaVersionOutput{SchemaDefinition: res.SchemaDefinition}, nil
   116  	}
   117  
   118  	res, ok := m.registerSchemaVersionsInput[*params.SchemaVersionId]
   119  	if !ok {
   120  		return nil, errors.New("EntityNotFoundException")
   121  	}
   122  	return &glue.GetSchemaVersionOutput{SchemaDefinition: res.SchemaDefinition}, nil
   123  }
   124  
   125  func (m *mockGlueClientImpl) GetRegistry(
   126  	ctx context.Context,
   127  	params *glue.GetRegistryInput,
   128  	optFns ...func(*glue.Options),
   129  ) (*glue.GetRegistryOutput, error) {
   130  	return nil, nil
   131  }