diff --git a/keydb_extension.go b/keydb_extension.go
index c4afcb2..05ad779 100644
--- a/keydb_extension.go
+++ b/keydb_extension.go
@@ -20,7 +20,7 @@ func init() {
 type KeyDBHandler struct {
 	Address  string `json:"address"`
 	Password string `json:"password"`
-	DB       string `json:"db"`
+	DB       int    `json:"db"`
 	client   *redis.Client
 }
 
@@ -32,17 +32,13 @@ func (KeyDBHandler) CaddyModule() caddy.ModuleInfo {
 }
 
 func (h *KeyDBHandler) Provision(ctx caddy.Context) error {
-	db, err := strconv.Atoi(h.DB)
-	if err != nil {
-		return err
-	}
-	if db < 0 || db > 15 {
-		return fmt.Errorf("invalid db value: %d", db)
+	if h.DB < 0 || h.DB > 15 {
+		return fmt.Errorf("invalid db value: %d", h.DB)
 	}
 	h.client = redis.NewClient(&redis.Options{
 		Addr:     h.Address,
 		Password: h.Password,
-		DB:       db,
+		DB:       h.DB, // Directly use h.DB
 	})
 	return nil
 }
@@ -83,9 +79,15 @@ func (h *KeyDBHandler) UnmarshalCaddyfile(d *caddyfile.Dispenser) error {
 					return d.ArgErr()
 				}
 			case "db":
-				if !d.Args(&h.DB) {
+				var dbString string
+				if !d.Args(&dbString) {
 					return d.ArgErr()
 				}
+				db, err := strconv.Atoi(dbString)
+				if err != nil {
+					return err
+				}
+				h.DB = db
 			}
 		}
 	}