-
Notifications
You must be signed in to change notification settings - Fork 0
/
client.go
167 lines (147 loc) · 5.02 KB
/
client.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
package client
import (
"bytes"
"context"
"crypto/tls"
"encoding/json"
"fmt"
"sync"
"github.com/apache/arrow/go/v16/arrow"
"github.com/apache/arrow/go/v16/arrow/flight"
"github.com/apache/arrow/go/v16/arrow/ipc"
"github.com/cloudquery/plugin-sdk/v4/message"
"github.com/cloudquery/plugin-sdk/v4/plugin"
"github.com/cloudquery/plugin-sdk/v4/schema"
"github.com/rs/zerolog"
"google.golang.org/grpc"
"google.golang.org/grpc/credentials"
"google.golang.org/grpc/credentials/insecure"
)
// Client is the client for the Apache Arrow destination
type Client struct {
logger zerolog.Logger
spec Spec
done chan struct{}
flightClient flight.Client
flightDoPutClients []flight.FlightService_DoPutClient
lock sync.RWMutex
wg sync.WaitGroup
writers map[string]*flight.Writer
plugin.UnimplementedSource
}
var _ plugin.Client = (*Client)(nil)
// New creates a new Apache Arrow destination client
func New(ctx context.Context, logger zerolog.Logger, specBytes []byte, opts plugin.NewClientOptions) (plugin.Client, error) {
c := &Client{
logger: logger.With().Str("module", "arrowflight").Logger(),
done: make(chan struct{}),
writers: make(map[string]*flight.Writer),
}
if opts.NoConnection {
return c, nil
}
if err := json.Unmarshal(specBytes, &c.spec); err != nil {
return nil, fmt.Errorf("failed to unmarshal spec: %w", err)
}
c.spec.SetDefaults()
var transportCredentials credentials.TransportCredentials
if c.spec.TlsEnabled {
transportCredentials = credentials.NewTLS(&tls.Config{
ServerName: c.spec.TlsServerName,
InsecureSkipVerify: c.spec.TlsInsecureSkipVerify,
})
} else {
transportCredentials = insecure.NewCredentials()
}
var grpcDialOptions []grpc.DialOption
grpcDialOptions = append(grpcDialOptions, grpc.WithTransportCredentials(transportCredentials))
var callOptions []grpc.CallOption
if c.spec.MaxCallRecvMsgSize != nil {
callOptions = append(callOptions, grpc.MaxCallRecvMsgSize(*c.spec.MaxCallRecvMsgSize))
}
if c.spec.MaxCallSendMsgSize != nil {
callOptions = append(callOptions, grpc.MaxCallSendMsgSize(*c.spec.MaxCallSendMsgSize))
}
grpcDialOptions = append(grpcDialOptions, grpc.WithDefaultCallOptions(callOptions...))
var err error
if c.flightClient, err = flight.NewClientWithMiddlewareCtx(ctx, c.spec.Addr, NewAuthHandler(c.spec.Handshake, c.spec.Token), nil, grpcDialOptions...); err != nil {
return nil, fmt.Errorf("failed to create flight client: %w", err)
}
if c.spec.Handshake != "" {
if err = c.flightClient.Authenticate(ctx); err != nil {
return nil, fmt.Errorf("failed to authenticate flight client: %w", err)
}
}
return c, nil
}
// Read is called when a table is read
func (c *Client) Read(ctx context.Context, table *schema.Table, res chan<- arrow.Record) error {
c.logger.Debug().Str("table", table.Name).Msg("read")
flightInfo, err := c.getFlightInfo(ctx, table.Name)
if err != nil {
return fmt.Errorf("failed to get flight info: %w", err)
}
var reader *ipc.Reader
if reader, err = ipc.NewReader(bytes.NewReader(flightInfo.GetSchema())); err != nil {
return fmt.Errorf("failed to create reader: %w", err)
}
defer reader.Release()
for _, endpoint := range flightInfo.GetEndpoint() {
if err = c.doGet(ctx, endpoint, reader.Schema(), res); err != nil {
return err
}
}
return nil
}
// Write receives messages from the plugin and writes them to the destination
func (c *Client) Write(ctx context.Context, messages <-chan message.WriteMessage) error {
for msg := range messages {
switch v := msg.(type) {
case *message.WriteMigrateTable:
if err := c.MigrateTable(ctx, v); err != nil {
return fmt.Errorf("failed to migrate table: %w", err)
}
case *message.WriteInsert:
if err := c.Insert(ctx, v); err != nil {
return fmt.Errorf("failed to insert: %w", err)
}
case *message.WriteDeleteStale:
if err := c.DeleteStale(ctx, v); err != nil {
return fmt.Errorf("failed to delete stale: %w", err)
}
case *message.WriteDeleteRecord:
if err := c.DeleteRecord(ctx, v); err != nil {
return fmt.Errorf("failed to delete record: %w", err)
}
default:
return fmt.Errorf("unknown message type: %T", v)
}
}
return nil
}
// Close is called when the client is closed
func (c *Client) Close(context.Context) error {
if c.flightClient == nil {
return fmt.Errorf("client already closed or not initialized")
}
close(c.done)
c.logger.Debug().Msg("closing writers")
for _, writer := range c.writers {
if err := writer.Close(); err != nil {
c.logger.Error().Err(err).Msg("failed to close writer")
}
}
waitTimeout(&c.wg, c.spec.CloseTimeout.Duration())
c.logger.Debug().Msg("closing do put clients")
for _, flightDoPutClient := range c.flightDoPutClients {
if err := flightDoPutClient.CloseSend(); err != nil {
c.logger.Error().Err(err).Msg("failed to close send")
}
}
c.logger.Info().Msg("closing flight client")
if err := c.flightClient.Close(); err != nil {
c.logger.Error().Err(err).Msg("failed to close flight client")
}
c.flightClient = nil
return nil
}