From e43771425c6178ca2b073daebf58eec4c66f819b Mon Sep 17 00:00:00 2001
From: jchappelow <140431406+jchappelow@users.noreply.github.com>
Date: Sun, 15 Oct 2023 20:52:06 -0500
Subject: [PATCH] log,grpc: simpler gRPC logger for server (#348)

* pkg/log,grpc: simpler gRPC logger for server

* include client address
---
 core/log/logger.go                       |  7 +++
 internal/services/grpc_server/logging.go | 77 +++++++++++++++---------
 internal/services/grpc_server/server.go  |  7 +--
 3 files changed, 57 insertions(+), 34 deletions(-)

diff --git a/core/log/logger.go b/core/log/logger.go
index c3ec104a7..4f97cf080 100644
--- a/core/log/logger.go
+++ b/core/log/logger.go
@@ -45,6 +45,13 @@ func (l *Logger) Fatal(msg string, fields ...Field) {
 	l.L.Fatal(msg, fields...)
 }
 
+func (l *Logger) Log(level Level, msg string, fields ...Field) {
+	if level >= ErrorLevel {
+		fields = append(fields, detailedFields...)
+	}
+	l.L.Log(level, msg, fields...)
+}
+
 func (l *Logger) Named(name string) *Logger {
 	if name == "" {
 		return l
diff --git a/internal/services/grpc_server/logging.go b/internal/services/grpc_server/logging.go
index 8a84aacbd..4cdcc905a 100644
--- a/internal/services/grpc_server/logging.go
+++ b/internal/services/grpc_server/logging.go
@@ -3,41 +3,62 @@ package server
 import (
 	"context"
 	"fmt"
+	"strings"
+	"time"
 
 	"github.com/kwilteam/kwil-db/core/log"
 
-	"github.com/grpc-ecosystem/go-grpc-middleware/v2/interceptors/logging"
 	"go.uber.org/zap"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/peer"
+	"google.golang.org/grpc/status"
 )
 
-// InterceptorLogger adapts zap logger to interceptor logger.
-// This code is copied from go-grpc-middleware/interceptors/logging/examples/zap/example_test.go
-func InterceptorLogger(l *log.Logger) logging.Logger {
+func codeToLevel(code codes.Code) log.Level {
+	switch code {
+	case codes.OK:
+		return log.InfoLevel // log.DebugLevel
+	case codes.NotFound, codes.Canceled, codes.AlreadyExists, codes.InvalidArgument, codes.Unauthenticated:
+		return log.InfoLevel
 
-	return logging.LoggerFunc(func(ctx context.Context, lvl logging.Level, msg string, fields ...any) {
-		var f []zap.Field
-		i := logging.Fields(fields).Iterator()
-		for i.Next() {
-			k, v := i.At()
-			f = append(f, zap.Any(k, v))
-		}
+	case codes.DeadlineExceeded, codes.PermissionDenied, codes.ResourceExhausted, codes.FailedPrecondition,
+		codes.Aborted, codes.OutOfRange, codes.Unavailable:
+		return log.WarnLevel
+
+	case codes.Unknown, codes.Unimplemented, codes.Internal, codes.DataLoss:
+		return log.ErrorLevel
+
+	default:
+		return log.WarnLevel
+	}
+}
 
-		// TODO: this is a hack to get rid of the extended fields every time we log
-		// log wrapper is not correctly cloned
-		// here use zap.Logger
-		lg := l.L.WithOptions(zap.AddCallerSkip(1)).With(f...)
-
-		switch lvl {
-		case logging.LevelDebug:
-			lg.Debug(msg)
-		case logging.LevelInfo:
-			lg.Info(msg)
-		case logging.LevelWarn:
-			lg.Warn(msg)
-		case logging.LevelError:
-			lg.Error(msg)
-		default:
-			panic(fmt.Sprintf("unknown level %v", lvl))
+// SimpleInterceptorLogger is a simplified gRPC server request logger. For an
+// alternative, see the example from
+// go-grpc-middleware/interceptors/logging/examples/zap/example_test.go
+func SimpleInterceptorLogger(l *log.Logger) grpc.UnaryServerInterceptor {
+	return func(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) {
+		tStart := time.Now()
+		resp, err := handler(ctx, req)
+		elapsedMs := float64(time.Since(tStart).Microseconds()) / 1e3
+		code := status.Code(err)
+		fields := []zap.Field{
+			zap.String("method", strings.Trim(info.FullMethod, "/")),
+			zap.String("elapsed", fmt.Sprintf("%.3fms", elapsedMs)),
+		}
+		if peer, ok := peer.FromContext(ctx); ok {
+			fields = append(fields, zap.String("addr", peer.Addr.String()))
+		}
+		fields = append(fields, zap.String("code", code.String()))
+		var msg string
+		if err != nil {
+			msg = "call failure"
+			fields = append(fields, zap.Error(err))
+		} else {
+			msg = "call success"
 		}
-	})
+		l.Log(codeToLevel(code), msg, fields...)
+		return resp, err
+	}
 }
diff --git a/internal/services/grpc_server/server.go b/internal/services/grpc_server/server.go
index a29cb5876..8a7322aca 100644
--- a/internal/services/grpc_server/server.go
+++ b/internal/services/grpc_server/server.go
@@ -5,7 +5,6 @@ import (
 
 	"github.com/kwilteam/kwil-db/core/log"
 
-	"github.com/grpc-ecosystem/go-grpc-middleware/v2/interceptors/logging"
 	"github.com/grpc-ecosystem/go-grpc-middleware/v2/interceptors/recovery"
 	"go.uber.org/zap"
 	"google.golang.org/grpc"
@@ -32,10 +31,6 @@ func New(logger log.Logger, lis net.Listener, opts ...Option) *Server {
 		recovery.WithRecoveryHandler(recoveryFunc),
 	}
 
-	loggingOpts := []logging.Option{
-		logging.WithLogOnEvents(logging.StartCall, logging.FinishCall),
-	}
-
 	s := &Server{
 		logger:   l,
 		listener: lis,
@@ -46,7 +41,7 @@ func New(logger log.Logger, lis net.Listener, opts ...Option) *Server {
 	}
 	srvOpts := append(s.srvOpts, grpc.ChainUnaryInterceptor(
 		recovery.UnaryServerInterceptor(recoveryOpts...),
-		logging.UnaryServerInterceptor(InterceptorLogger(&l), loggingOpts...),
+		SimpleInterceptorLogger(&l),
 	))
 
 	s.server = grpc.NewServer(srvOpts...)