chore: fix export memos

This commit is contained in:
Steven
2024-02-04 20:20:14 +08:00
parent b8a9783db5
commit 1e07b70d23
10 changed files with 224 additions and 328 deletions

View File

@@ -27,40 +27,6 @@ const (
usernameContextKey ContextKey = iota
)
// Used to set modified context of ServerStream.
type WrappedStream struct {
ctx context.Context
stream grpc.ServerStream
}
func (w *WrappedStream) RecvMsg(m any) error {
return w.stream.RecvMsg(m)
}
func (w *WrappedStream) SendMsg(m any) error {
return w.stream.SendMsg(m)
}
func (w *WrappedStream) SendHeader(md metadata.MD) error {
return w.stream.SendHeader(md)
}
func (w *WrappedStream) SetHeader(md metadata.MD) error {
return w.stream.SetHeader(md)
}
func (w *WrappedStream) SetTrailer(md metadata.MD) {
w.stream.SetTrailer(md)
}
func (w *WrappedStream) Context() context.Context {
return w.ctx
}
func newWrappedStream(ctx context.Context, stream grpc.ServerStream) grpc.ServerStream {
return &WrappedStream{ctx, stream}
}
// GRPCAuthInterceptor is the auth interceptor for gRPC server.
type GRPCAuthInterceptor struct {
Store *store.Store
@@ -114,45 +80,6 @@ func (in *GRPCAuthInterceptor) AuthenticationInterceptor(ctx context.Context, re
return handler(childCtx, request)
}
func (in *GRPCAuthInterceptor) StreamAuthenticationInterceptor(srv any, stream grpc.ServerStream, serverInfo *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
md, ok := metadata.FromIncomingContext(stream.Context())
if !ok {
return status.Errorf(codes.Unauthenticated, "failed to parse metadata from incoming context")
}
accessToken, err := getTokenFromMetadata(md)
if err != nil {
return status.Errorf(codes.Unauthenticated, err.Error())
}
username, err := in.authenticate(stream.Context(), accessToken)
if err != nil {
if isUnauthorizeAllowedMethod(serverInfo.FullMethod) {
return handler(stream.Context(), stream)
}
return err
}
user, err := in.Store.GetUser(stream.Context(), &store.FindUser{
Username: &username,
})
if err != nil {
return errors.Wrap(err, "failed to get user")
}
if user == nil {
return errors.Errorf("user %q not exists", username)
}
if user.RowStatus == store.Archived {
return errors.Errorf("user %q is archived", username)
}
if isOnlyForAdminAllowedMethod(serverInfo.FullMethod) && user.Role != store.RoleHost && user.Role != store.RoleAdmin {
return errors.Errorf("user %q is not admin", username)
}
// Stores userID into context.
childCtx := context.WithValue(stream.Context(), usernameContextKey, username)
return handler(srv, newWrappedStream(childCtx, stream))
}
func (in *GRPCAuthInterceptor) authenticate(ctx context.Context, accessToken string) (string, error) {
if accessToken == "" {
return "", status.Errorf(codes.Unauthenticated, "access token not found")

View File

@@ -481,15 +481,9 @@ paths:
operationId: MemoService_ExportMemos
responses:
"200":
description: A successful response.(streaming responses)
description: A successful response.
schema:
type: object
properties:
result:
$ref: '#/definitions/v2ExportMemosResponse'
error:
$ref: '#/definitions/googlerpcStatus'
title: Stream result of v2ExportMemosResponse
$ref: '#/definitions/v2ExportMemosResponse'
default:
description: An unexpected error response.
schema:
@@ -1520,7 +1514,7 @@ definitions:
v2ExportMemosResponse:
type: object
properties:
file:
content:
type: string
format: byte
v2GetActivityResponse:

View File

@@ -508,17 +508,15 @@ func (s *APIV2Service) GetUserMemosStats(ctx context.Context, request *apiv2pb.G
return response, nil
}
func (s *APIV2Service) ExportMemos(request *apiv2pb.ExportMemosRequest, srv apiv2pb.MemoService_ExportMemosServer) error {
ctx := srv.Context()
fmt.Printf("%+v\n", ctx)
func (s *APIV2Service) ExportMemos(ctx context.Context, request *apiv2pb.ExportMemosRequest) (*apiv2pb.ExportMemosResponse, error) {
memoFind, err := s.buildFindMemosWithFilter(ctx, request.Filter, true)
if err != nil {
return err
return nil, status.Errorf(codes.Internal, "failed to build find memos with filter")
}
memos, err := s.Store.ListMemos(ctx, memoFind)
if err != nil {
return err
return nil, status.Errorf(codes.Internal, "failed to list memos")
}
buf := new(bytes.Buffer)
@@ -526,41 +524,27 @@ func (s *APIV2Service) ExportMemos(request *apiv2pb.ExportMemosRequest, srv apiv
for _, memo := range memos {
memoMessage, err := s.convertMemoFromStore(ctx, memo)
log.Info(memoMessage.Content)
if err != nil {
return errors.Wrap(err, "failed to convert memo")
return nil, errors.Wrap(err, "failed to convert memo")
}
file, err := writer.Create(time.Unix(memo.CreatedTs, 0).Format(time.RFC3339) + ".md")
if err != nil {
return status.Errorf(codes.Internal, "Failed to create memo file")
return nil, status.Errorf(codes.Internal, "Failed to create memo file")
}
_, err = file.Write([]byte(memoMessage.Content))
if err != nil {
return status.Errorf(codes.Internal, "Failed to write to memo file")
return nil, status.Errorf(codes.Internal, "Failed to write to memo file")
}
}
err = writer.Close()
if err != nil {
return status.Errorf(codes.Internal, "Failed to close zip file writer")
return nil, status.Errorf(codes.Internal, "Failed to close zip file writer")
}
exportChunk := &apiv2pb.ExportMemosResponse{}
sizeOfFile := len(buf.Bytes())
for currentByte := 0; currentByte < sizeOfFile; currentByte += ChunkSize {
if currentByte+ChunkSize > sizeOfFile {
exportChunk.File = buf.Bytes()[currentByte:sizeOfFile]
} else {
exportChunk.File = buf.Bytes()[currentByte : currentByte+ChunkSize]
}
err := srv.Send(exportChunk)
if err != nil {
return status.Error(codes.Internal, "Unable to stream ExportMemosResponse chunk")
}
}
return nil
return &apiv2pb.ExportMemosResponse{
Content: buf.Bytes(),
}, nil
}
func (s *APIV2Service) convertMemoFromStore(ctx context.Context, memo *store.Memo) (*apiv2pb.Memo, error) {

View File

@@ -46,9 +46,6 @@ func NewAPIV2Service(secret string, profile *profile.Profile, store *store.Store
grpc.ChainUnaryInterceptor(
authProvider.AuthenticationInterceptor,
),
grpc.ChainStreamInterceptor(
authProvider.StreamAuthenticationInterceptor,
),
)
apiv2Service := &APIV2Service{
Secret: secret,