mirror of https://github.com/databricks/cli.git
100 lines
2.4 KiB
Go
100 lines
2.4 KiB
Go
package telemetry
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"os"
|
|
"time"
|
|
|
|
"github.com/databricks/cli/libs/telemetry/protos"
|
|
"github.com/databricks/databricks-sdk-go/client"
|
|
"github.com/databricks/databricks-sdk-go/config"
|
|
)
|
|
|
|
const (
|
|
// File containing debug logs from the upload process.
|
|
UploadLogsFileEnvVar = "DATABRICKS_CLI_TELEMETRY_UPLOAD_LOGS_FILE"
|
|
|
|
// File containing the PID of the telemetry upload process.
|
|
PidFileEnvVar = "DATABRICKS_CLI_TELEMETRY_PID_FILE"
|
|
)
|
|
|
|
type UploadConfig struct {
|
|
Logs []protos.FrontendLog `json:"logs"`
|
|
}
|
|
|
|
// Upload reads telemetry logs from stdin and uploads them to the telemetry endpoint.
|
|
// This function is always expected to be called in a separate child process from
|
|
// the main CLI process.
|
|
func Upload() (*ResponseBody, error) {
|
|
var err error
|
|
|
|
b, err := io.ReadAll(os.Stdin)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to read from stdin: %s\n", err)
|
|
}
|
|
|
|
in := UploadConfig{}
|
|
err = json.Unmarshal(b, &in)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to unmarshal input: %s\n", err)
|
|
}
|
|
|
|
if len(in.Logs) == 0 {
|
|
return nil, fmt.Errorf("No logs to upload: %s\n", err)
|
|
}
|
|
|
|
protoLogs := make([]string, len(in.Logs))
|
|
for i, log := range in.Logs {
|
|
b, err := json.Marshal(log)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to marshal log: %s\n", err)
|
|
}
|
|
protoLogs[i] = string(b)
|
|
}
|
|
|
|
// Parent process is responsible for setting environment variables to
|
|
// configure authentication.
|
|
apiClient, err := client.New(&config.Config{})
|
|
if err != nil {
|
|
return nil, fmt.Errorf("Failed to create API client: %s\n", err)
|
|
}
|
|
|
|
// Set a maximum total time to try telemetry uploads.
|
|
ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
|
|
defer cancel()
|
|
|
|
resp := &ResponseBody{}
|
|
for {
|
|
select {
|
|
case <-ctx.Done():
|
|
return nil, errors.New("Failed to flush telemetry log due to timeout")
|
|
|
|
default:
|
|
// Proceed
|
|
}
|
|
|
|
// Log the CLI telemetry events.
|
|
err := apiClient.Do(ctx, http.MethodPost, "/telemetry-ext", nil, nil, RequestBody{
|
|
UploadTime: time.Now().UnixMilli(),
|
|
Items: []string{},
|
|
ProtoLogs: protoLogs,
|
|
}, resp)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("Failed to upload telemetry logs: %s\n", err)
|
|
}
|
|
|
|
if len(resp.Errors) > 0 {
|
|
return nil, fmt.Errorf("Failed to upload telemetry logs: %s\n", resp.Errors)
|
|
}
|
|
|
|
if resp.NumProtoSuccess == int64(len(in.Logs)) {
|
|
return resp, nil
|
|
}
|
|
}
|
|
}
|