-
Notifications
You must be signed in to change notification settings - Fork 26
feat: Support sources in SDK V3 #864
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Merged
Merged
Changes from all commits
Commits
Show all changes
14 commits
Select commit
Hold shift + click to select a range
e502681
feat: Add SDKV3 support for sources and update protocol
yevgenypats c9d4321
fmt
yevgenypats 84ea9d2
fix lint
yevgenypats 2734aef
minor fixes
yevgenypats 093f04f
rebase
yevgenypats 09f9478
fix lints
yevgenypats c479525
fix json on sdk side
yevgenypats a37d929
fix: source_name
yevgenypats 4e58515
fix more stuff
yevgenypats 0a83c35
fix fmt
yevgenypats 4f1e337
fix tests
yevgenypats 099728d
fix transformer
yevgenypats 9e1a409
fmt fix
yevgenypats 7a33710
comment some stuff out untill another iteration
yevgenypats File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,191 @@ | ||
| package destination | ||
|
|
||
| import ( | ||
| "bytes" | ||
| "context" | ||
| "encoding/json" | ||
| "fmt" | ||
| "io" | ||
|
|
||
| "github.com/apache/arrow/go/v13/arrow" | ||
| "github.com/apache/arrow/go/v13/arrow/ipc" | ||
| pb "github.com/cloudquery/plugin-pb-go/pb/destination/v1" | ||
| "github.com/cloudquery/plugin-pb-go/specs" | ||
| "github.com/cloudquery/plugin-sdk/v3/plugins/destination" | ||
| "github.com/cloudquery/plugin-sdk/v3/schema" | ||
| "github.com/rs/zerolog" | ||
| "golang.org/x/sync/errgroup" | ||
| "google.golang.org/grpc/codes" | ||
| "google.golang.org/grpc/status" | ||
| ) | ||
|
|
||
| type Server struct { | ||
| pb.UnimplementedDestinationServer | ||
| Plugin *destination.Plugin | ||
| Logger zerolog.Logger | ||
| spec specs.Destination | ||
| } | ||
|
|
||
| func (s *Server) Configure(ctx context.Context, req *pb.Configure_Request) (*pb.Configure_Response, error) { | ||
| var spec specs.Destination | ||
| if err := json.Unmarshal(req.Config, &spec); err != nil { | ||
| return nil, status.Errorf(codes.InvalidArgument, "failed to unmarshal spec: %v", err) | ||
| } | ||
| s.spec = spec | ||
| return &pb.Configure_Response{}, s.Plugin.Init(ctx, s.Logger, spec) | ||
| } | ||
|
|
||
| func (s *Server) GetName(context.Context, *pb.GetName_Request) (*pb.GetName_Response, error) { | ||
| return &pb.GetName_Response{ | ||
| Name: s.Plugin.Name(), | ||
| }, nil | ||
| } | ||
|
|
||
| func (s *Server) GetVersion(context.Context, *pb.GetVersion_Request) (*pb.GetVersion_Response, error) { | ||
| return &pb.GetVersion_Response{ | ||
| Version: s.Plugin.Version(), | ||
| }, nil | ||
| } | ||
|
|
||
| func (s *Server) Migrate(ctx context.Context, req *pb.Migrate_Request) (*pb.Migrate_Response, error) { | ||
| schemas, err := schema.NewSchemasFromBytes(req.Tables) | ||
| if err != nil { | ||
| return nil, status.Errorf(codes.InvalidArgument, "failed to create schemas: %v", err) | ||
| } | ||
| tables, err := schema.NewTablesFromArrowSchemas(schemas) | ||
| if err != nil { | ||
| return nil, status.Errorf(codes.InvalidArgument, "failed to create tables: %v", err) | ||
| } | ||
| s.setPKsForTables(tables) | ||
|
|
||
| return &pb.Migrate_Response{}, s.Plugin.Migrate(ctx, tables) | ||
| } | ||
|
|
||
| // Note the order of operations in this method is important! | ||
| // Trying to insert into the `resources` channel before starting the reader goroutine will cause a deadlock. | ||
| func (s *Server) Write(msg pb.Destination_WriteServer) error { | ||
| resources := make(chan arrow.Record) | ||
|
|
||
| r, err := msg.Recv() | ||
| if err != nil { | ||
| if err == io.EOF { | ||
| return msg.SendAndClose(&pb.Write_Response{}) | ||
| } | ||
| return status.Errorf(codes.Internal, "failed to receive msg: %v", err) | ||
| } | ||
|
|
||
| schemas, err := schema.NewSchemasFromBytes(r.Tables) | ||
| if err != nil { | ||
| return status.Errorf(codes.InvalidArgument, "failed to create schemas: %v", err) | ||
| } | ||
| tables, err := schema.NewTablesFromArrowSchemas(schemas) | ||
| if err != nil { | ||
| return status.Errorf(codes.InvalidArgument, "failed to create tables: %v", err) | ||
| } | ||
| var sourceSpec specs.Source | ||
| if r.SourceSpec == nil { | ||
| // this is for backward compatibility | ||
| sourceSpec = specs.Source{ | ||
| Name: r.Source, | ||
| } | ||
| } else { | ||
| if err := json.Unmarshal(r.SourceSpec, &sourceSpec); err != nil { | ||
| return status.Errorf(codes.InvalidArgument, "failed to unmarshal source spec: %v", err) | ||
| } | ||
| } | ||
| syncTime := r.Timestamp.AsTime() | ||
| s.setPKsForTables(tables) | ||
| eg, ctx := errgroup.WithContext(msg.Context()) | ||
| eg.Go(func() error { | ||
| return s.Plugin.Write(ctx, sourceSpec, tables, syncTime, resources) | ||
| }) | ||
|
|
||
| for { | ||
| r, err := msg.Recv() | ||
| if err == io.EOF { | ||
| close(resources) | ||
| if err := eg.Wait(); err != nil { | ||
| return status.Errorf(codes.Internal, "write failed: %v", err) | ||
| } | ||
| return msg.SendAndClose(&pb.Write_Response{}) | ||
| } | ||
| if err != nil { | ||
| close(resources) | ||
| if wgErr := eg.Wait(); wgErr != nil { | ||
| return status.Errorf(codes.Internal, "failed to receive msg: %v and write failed: %v", err, wgErr) | ||
| } | ||
| return status.Errorf(codes.Internal, "failed to receive msg: %v", err) | ||
| } | ||
| rdr, err := ipc.NewReader(bytes.NewReader(r.Resource)) | ||
| if err != nil { | ||
| close(resources) | ||
| if wgErr := eg.Wait(); wgErr != nil { | ||
| return status.Errorf(codes.InvalidArgument, "failed to create reader: %v and write failed: %v", err, wgErr) | ||
| } | ||
| return status.Errorf(codes.InvalidArgument, "failed to create reader: %v", err) | ||
| } | ||
| for rdr.Next() { | ||
| rec := rdr.Record() | ||
| rec.Retain() | ||
| select { | ||
| case resources <- rec: | ||
| case <-ctx.Done(): | ||
| close(resources) | ||
| if err := eg.Wait(); err != nil { | ||
| return status.Errorf(codes.Internal, "Context done: %v and failed to wait for plugin: %v", ctx.Err(), err) | ||
| } | ||
| return status.Errorf(codes.Internal, "Context done: %v", ctx.Err()) | ||
| } | ||
| } | ||
| if err := rdr.Err(); err != nil { | ||
| return status.Errorf(codes.InvalidArgument, "failed to read resource: %v", err) | ||
| } | ||
| } | ||
| } | ||
|
|
||
| func setCQIDAsPrimaryKeysForTables(tables schema.Tables) { | ||
| for _, table := range tables { | ||
| for i, col := range table.Columns { | ||
| table.Columns[i].PrimaryKey = col.Name == schema.CqIDColumn.Name | ||
| } | ||
| setCQIDAsPrimaryKeysForTables(table.Relations) | ||
| } | ||
| } | ||
|
|
||
| func (s *Server) GetMetrics(context.Context, *pb.GetDestinationMetrics_Request) (*pb.GetDestinationMetrics_Response, error) { | ||
| stats := s.Plugin.Metrics() | ||
| b, err := json.Marshal(stats) | ||
| if err != nil { | ||
| return nil, fmt.Errorf("failed to marshal stats: %w", err) | ||
| } | ||
| return &pb.GetDestinationMetrics_Response{ | ||
| Metrics: b, | ||
| }, nil | ||
| } | ||
|
|
||
| func (s *Server) DeleteStale(ctx context.Context, req *pb.DeleteStale_Request) (*pb.DeleteStale_Response, error) { | ||
| schemas, err := schema.NewSchemasFromBytes(req.Tables) | ||
| if err != nil { | ||
| return nil, status.Errorf(codes.InvalidArgument, "failed to create schemas: %v", err) | ||
| } | ||
| tables, err := schema.NewTablesFromArrowSchemas(schemas) | ||
| if err != nil { | ||
| return nil, status.Errorf(codes.InvalidArgument, "failed to create tables: %v", err) | ||
| } | ||
|
|
||
| if err := s.Plugin.DeleteStale(ctx, tables, req.Source, req.Timestamp.AsTime()); err != nil { | ||
| return nil, err | ||
| } | ||
|
|
||
| return &pb.DeleteStale_Response{}, nil | ||
| } | ||
|
|
||
| func (s *Server) setPKsForTables(tables schema.Tables) { | ||
| if s.spec.PKMode == specs.PKModeCQID { | ||
| setCQIDAsPrimaryKeysForTables(tables) | ||
| } | ||
| } | ||
|
|
||
| func (s *Server) Close(ctx context.Context, _ *pb.Close_Request) (*pb.Close_Response, error) { | ||
| return &pb.Close_Response{}, s.Plugin.Close(ctx) | ||
| } | ||
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This could be using
schema.NewTablesFromBytes(r.Tables)(which is a little bit shorter)