From f6592d9bbd0dd6bbf0ef88090dcfb21c6951c3bb Mon Sep 17 00:00:00 2001
From: Shreyas Goenka <shreyas.goenka@databricks.com>
Date: Wed, 5 Mar 2025 14:54:30 +0100
Subject: [PATCH] Add synchronous logger for telemetry

---
 libs/telemetry/context.go | 25 ++++++++++++++++++++
 libs/telemetry/logger.go  | 50 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 75 insertions(+)
 create mode 100644 libs/telemetry/context.go
 create mode 100644 libs/telemetry/logger.go

diff --git a/libs/telemetry/context.go b/libs/telemetry/context.go
new file mode 100644
index 000000000..e556e462c
--- /dev/null
+++ b/libs/telemetry/context.go
@@ -0,0 +1,25 @@
+package telemetry
+
+import (
+	"context"
+	"errors"
+)
+
+// Private type to store the telemetry logger in the context
+type telemetryLogger int
+
+// Key to store the telemetry logger in the context
+var telemetryLoggerKey telemetryLogger
+
+func WithNewLogger(ctx context.Context) context.Context {
+	return context.WithValue(ctx, telemetryLoggerKey, &logger{})
+}
+
+func fromContext(ctx context.Context) *logger {
+	v := ctx.Value(telemetryLoggerKey)
+	if v == nil {
+		panic(errors.New("telemetry logger not found in the context"))
+	}
+
+	return v.(*logger)
+}
diff --git a/libs/telemetry/logger.go b/libs/telemetry/logger.go
new file mode 100644
index 000000000..4fbce2f52
--- /dev/null
+++ b/libs/telemetry/logger.go
@@ -0,0 +1,50 @@
+package telemetry
+
+import (
+	"context"
+
+	"github.com/databricks/cli/libs/telemetry/protos"
+	"github.com/google/uuid"
+)
+
+func Log(ctx context.Context, event protos.DatabricksCliLog) {
+	fromContext(ctx).log(event)
+}
+
+func GetLogs(ctx context.Context) []protos.FrontendLog {
+	return fromContext(ctx).getLogs()
+}
+
+func HasLogs(ctx context.Context) bool {
+	return len(fromContext(ctx).getLogs()) > 0
+}
+
+func SetExecutionContext(ctx context.Context, ec protos.ExecutionContext) {
+	fromContext(ctx).setExecutionContext(ec)
+}
+
+type logger struct {
+	logs []protos.FrontendLog
+}
+
+func (l *logger) log(event protos.DatabricksCliLog) {
+	if l.logs == nil {
+		l.logs = make([]protos.FrontendLog, 0)
+	}
+	l.logs = append(l.logs, protos.FrontendLog{
+		FrontendLogEventID: uuid.New().String(),
+		Entry: protos.FrontendLogEntry{
+			DatabricksCliLog: event,
+		},
+	})
+}
+
+func (l *logger) getLogs() []protos.FrontendLog {
+	return l.logs
+}
+
+func (l *logger) setExecutionContext(ec protos.ExecutionContext) {
+	for i := range l.logs {
+		l.logs[i].Entry.DatabricksCliLog.ExecutionContext = &ec
+	}
+}