메인 콘텐츠로 이동하기
  1. Posts/

Golang context를 활용한 Custom Wrapper Function

·770 자

소개 #

저는 현재 Skrr라는 앱 서비스를 개발중인데요. 처음에는 Python(Flask)과 MongoDB를 사용하여 작성되었지만, 이번에 기술스택을 Golang과 PostgreSQL로 마이그레이션하고 있습니다. 이에 대한 여러 배경이 있었지만, 이 글에서는 해당 내용을 자세히는 다루지 않을 예정입니다.

해당 글에서는 데이터베이스 객체와 JWT에서 추출한 사용자 ID를 주입하여 각 API handler에서 사용할 수 있는 custom wrapper 함수를 만드는 방법을 소개드리고자 합니다. 이를 통해 코드 중복을 줄이고 코드의 가독성을 높일 수 있습니다.

Golang에는 여러 멋진 웹 프레임워크가 있습니다. 예를 들어, Gin, Fiber, Echo 등이 있습니다. 이러한 도구들은 쉽게 사용할 수 있는 미들웨어 기능을 제공합니다. 하지만 저는 이 코드에 별도의 프레임워크를 전혀 사용하지 않고, zero-dependency를 유지하고 싶었습니다. 그리고, Golang 자체가 네트워크 프로그래밍을 풍부하게 지원하기 때문에 별도의 백엔드 프레임워크의 필요성을 느끼지도 못했습니다. 결국, net/http 패키지와 gorilla/mux를 최소한의 의존성으로 사용하기로 결정했습니다. 그리고 각 API handler에서 데이터베이스 객체와 JWT에서 사용자 ID를 주입하고 사용할 수 있는 custom wrapper 함수를 만들어야 했습니다.

Wrapper Function #

wrapper 함수는 컴퓨터 프로그램을 작성할 때 세부 사항을 추상화함으로써 작업을 용이하게 만드는 데 사용됩니다. 그러나 이 블로그 글에서의 뜻하는 wrapper 함수의 의미는 일반적인 wrapper 함수의 의미와 사뭇 다릅니다. 이 글에서 wrapper 함수는 다른 함수를 감싸는 함수입니다. 즉, 다른 함수를 인수로 받고 함수를 반환하는 함수입니다. 이를 고차 함수라고도 합니다.

예제 코드 #

아래 코드는 제가 원하는 함수가 수행해야 할 작업의 간단한 예시입니다. 이전 프로젝트에서 사용한 Python(FastAPI) 코드입니다. 인수로 전달된 함수는 데이터베이스 연결을 반환하는 함수입니다. 또한, JWT에서 추출된 사용자 ID를 반환하는 함수입니다. 이러한 설계는 각 API handler에서 database connection과 user ID를 사용하기 훨씬 더 편리하게 만듭니다.


```python
@router.get("/{uuid}")
async def find_one_diary(
    uuid: str, user_uuid=Depends(valid_request), db=Depends(get_db)
):
    diary = await db.execute(select(Diary).where(Diary.uuid == uuid))
    diary = diary.scalar()
    if diary is None:
        return None
    if diary.user_uuid != user_uuid:
        raise HTTPException(status_code=403, detail="Not your diary")

    return diary

구현 #

기본 Golang 코드 #

package main

import "net/http"

func main() {
    http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
        w.Write([]byte("Hello World"))
    })
    http.ListenAndServe(":8080", r)
}

이것은 간단한 웹 서버를 만드는 기본 Golang 코드입니다. 우리가 주목해야 할 중요한 점은 r.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) 부분입니다. 이 부분을 우리의 custom wrapper 함수로 래핑할 것입니다. 따라서 wrapper 함수를 주입하면 다음과 같이 보일 것입니다.

	db := config.InitDatabase(os.Getenv("POSTGRES_URL"))
	defer db.Close()

    // ...

    schoolRouter := router.PathPrefix("/school").Subrouter()
	schoolRouter.Handle("/{id}", D(db, api.FindOneSchool)).Methods(http.MethodGet)
	schoolRouter.Handle("/search/name", D(db, api.SearchSchoolByName)).Methods(http.MethodGet)

	friendRouter := router.PathPrefix("/friend").Subrouter()
	friendRouter.Handle("", ValidateRequest(D(db, api.FindAllFriend))).Methods(http.MethodGet)

    // ...

데이터베이스 객체 주입 #

함수 D는 제가 구현한 wrapper 함수입니다. 첫 번째 인수는 데이터베이스 객체이고, 두 번째 인수는 우리가 감싸고자 하는 함수입니다. 반환 값은 handler로 사용할 함수입니다.

func D(db *sql.DB, fn func(http.ResponseWriter, *http.Request, *sql.DB)) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		fn(w, r, db)
	})
}

두 번째 인수로 http.Handler를 사용하는 Handle 함수에 대해 설명했는데, 우리는 함수를 http.Handler 유형으로 변환해야 합니다. 이를 위해 함수를 http.Handler 유형으로 변환하기 위해 http.HandlerFunc를 사용하는 것입니다.

func FindOneSchool(w http.ResponseWriter, r *http.Request, db *sql.DB) {
	id := mux.Vars(r)["id"]

	school, _ := storage.FindOneSchool(db, id)
	if school.ID == "" {
		utils.JsonResp(w, "no school found", http.StatusBadRequest)
		return
	}

	utils.JsonResp(w, school, http.StatusOK)
}

이제 handler 함수에서 데이터베이스 객체를 쉽게 사용할 수 있습니다. 이 경우에는 전달된 매개변수인 ID를 사용하여 학교를 찾습니다.

JWT에서 유저 객체 추출을 위한 미들웨어 #

다음 단계는 JWT에서 user ID를 주입하는 것입니다. 데이터베이스를 주입하는 것보다 한 단계 더 어렵습니다. 왜냐하면 JWT에서 user ID를 추출하고 handler 함수에 전달해야 하기 때문입니다. 또한, JWT가 유효하지 않을 때 함수 중간에 401_UNAUTHORIZED를 반환해야 합니다. 따라서 JWT에서 user ID를 추출하고 handler 함수에 전달하는 미들웨어 함수를 만들어야 합니다.

func ValidateRequest(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		tokenString := r.Header.Get("Authorization")
		if tokenString == "" {
			JsonResp(w, errors.New("authorization header missing"), http.StatusUnauthorized)
			return
		}

		// Split the "Bearer " prefix from the token string
		tokenParts := strings.Split(tokenString, " ")
		if len(tokenParts) != 2 || strings.ToLower(tokenParts[0]) != "bearer" {
			JsonResp(w, errors.New("invalid Authorization header format"), http.StatusUnauthorized)
			return
		}
		tokenString = tokenParts[1]

		claims, err := ValidateJwt(tokenString)
		if err != nil {
			JsonResp(w, err, http.StatusUnauthorized)
			return
		}

		ctx := context.WithValue(r.Context(), "user_id", userID)
		r = r.WithContext(ctx)

		next.ServeHTTP(w, r)
	})
}

이 함수의 핵심은 ctx := context.WithValue(r.Context(), "user_id", claims.UserID)r = r.WithContext(ctx) 입니다. 이것이 우리가 user ID를 handler 함수로 전달하는 방법입니다. 첫 번째 인수는 요청의 컨텍스트이고, 두 번째 인수는 value의 key이며, 세 번째 인수는 value입니다. 값이 포함된 새 컨텍스트를 생성한 후에는 http.Request 컨텍스트로 전달됩니다. 이것이 우리가 r = r.WithContext(ctx)를 사용하는 이유입니다.

JWT 검증만을 위한 경우에는 미들웨어 사용만으로 충분했습니다. 그러나 이 경우에는 다음과 같이 handler 함수에 user ID를 전달할 수 있습니다.

func FindAllFriend(w http.ResponseWriter, r *http.Request, db *sql.DB) {
	userID := r.Context().Value("user_id").(string)

	usecase.FindAllFriend(w, db, userID)
}

이것이 handler 함수에서 user ID를 사용하는 방법입니다. userID := r.Context().Value("user_id").(string) 부분은 컨텍스트에서 사용자 ID를 추출하는 방법입니다. 이제 handler 함수에서도 user ID를 쉽게 사용할 수 있습니다.

결론 #

이 글에서는 데이터베이스 객체와 JWT에서 user ID를 주입하고 각 API handler에서 사용할 수 있는 custom wrapper 함수를 만드는 방법을 보여주었습니다. 이를 통해 코드 중복을 줄이고 코드의 가독성을 높일 수 있습니다. 또한 Golang의 context를 사용할 좋은 기회였습니다. 이는 Golang에서 가장 중요한 기능 중 하나입니다. 그리고, 이 글이 여러분에게 도움이 되었기를 바랍니다.