Add timeout options for http and nebula graph connection

vesoft-inc · Apr 21, 2020 · 329c29e · 329c29e
1 parent 976b908
commit 329c29e
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
diff --git a/pkg/client/client.go b/pkg/client/client.go
@@ -1,11 +1,14 @@
 package client
 
 import (
+	"time"
+
 	nebula "github.com/vesoft-inc/nebula-go"
 )
 
 func NewNebulaConnection(addr, user, password string) (*nebula.GraphClient, error) {
-	client, err := nebula.NewClient(addr)
+	opts := nebula.WithTimeout(10 * time.Second)
+	client, err := nebula.NewClient(addr, opts)
 	if err != nil {
 		return nil, err
 	}

diff --git a/pkg/reader/reader.go b/pkg/reader/reader.go
@@ -9,6 +9,7 @@ import (
 	"os"
 	"path"
 	"strings"
+	"time"
 
 	"github.com/vesoft-inc/nebula-importer/pkg/base"
 	"github.com/vesoft-inc/nebula-importer/pkg/config"
@@ -92,7 +93,9 @@ func (r *FileReader) prepareDataFile() (*string, error) {
 	}
 	defer file.Close()
 
-	resp, err := http.Get(*r.File.Path)
+	client := http.Client{Timeout: 2 * time.Second}
+
+	resp, err := client.Get(*r.File.Path)
 	if err != nil {
 		return nil, err
 	}