From d14a5ad54a8098e6db80371d561606c2ed64ca94 Mon Sep 17 00:00:00 2001
From: oiooj <oiooj@qq.com>
Date: Mon, 24 Oct 2016 00:32:18 +0800
Subject: [PATCH 3/5] support unix socket connect for influx CLI

Signed-off-by: Guillem Jover <gjover@sipwise.com>
---
 client/influxdb.go    | 29 ++++++++++++++++++++++-------
 cmd/influx/cli/cli.go |  8 +++++---
 cmd/influx/main.go    |  3 +++
 3 files changed, 30 insertions(+), 10 deletions(-)

diff --git a/client/influxdb.go b/client/influxdb.go
index 623276dd5..e28ac70fe 100644
--- a/client/influxdb.go
+++ b/client/influxdb.go
@@ -87,13 +87,14 @@ func ParseConnectionString(path string, ssl bool) (url.URL, error) {
 // UserAgent: If not provided, will default "InfluxDBClient",
 // Timeout: If not provided, will default to 0 (no timeout)
 type Config struct {
-	URL       url.URL
-	Username  string
-	Password  string
-	UserAgent string
-	Timeout   time.Duration
-	Precision string
-	UnsafeSsl bool
+	URL        url.URL
+	UnixSocket string
+	Username   string
+	Password   string
+	UserAgent  string
+	Timeout    time.Duration
+	Precision  string
+	UnsafeSsl  bool
 }
 
 // NewConfig will create a config to be used in connecting to the client
@@ -106,6 +107,7 @@ func NewConfig() Config {
 // Client is used to make calls to the server.
 type Client struct {
 	url        url.URL
+	unixSocket string
 	username   string
 	password   string
 	httpClient *http.Client
@@ -137,8 +139,18 @@ func NewClient(c Config) (*Client, error) {
 		TLSClientConfig: tlsConfig,
 	}
 
+	if c.UnixSocket != "" {
+		// No need for compression in local communications.
+		tr.DisableCompression = true
+
+		tr.Dial = func(_, _ string) (net.Conn, error) {
+			return net.Dial("unix", c.UnixSocket)
+		}
+	}
+
 	client := Client{
 		url:        c.URL,
+		unixSocket: c.UnixSocket,
 		username:   c.Username,
 		password:   c.Password,
 		httpClient: &http.Client{Timeout: c.Timeout, Transport: tr},
@@ -751,6 +763,9 @@ func (bp *BatchPoints) UnmarshalJSON(b []byte) error {
 
 // Addr provides the current url as a string of the server the client is connected to.
 func (c *Client) Addr() string {
+	if c.unixSocket != "" {
+		return c.unixSocket
+	}
 	return c.url.String()
 }
 
diff --git a/cmd/influx/cli/cli.go b/cmd/influx/cli/cli.go
index ca94b00cc..5705067c8 100644
--- a/cmd/influx/cli/cli.go
+++ b/cmd/influx/cli/cli.go
@@ -41,6 +41,7 @@ type CommandLine struct {
 	Line             *liner.State
 	Host             string
 	Port             int
+	UnixSocket       string
 	Username         string
 	Password         string
 	Database         string
@@ -117,8 +118,8 @@ func (c *CommandLine) Run() error {
 
 	if err := c.Connect(""); err != nil {
 		return fmt.Errorf(
-			"Failed to connect to %s\nPlease check your connection settings and ensure 'influxd' is running.",
-			c.Client.Addr())
+			"Failed to connect to %s: %s\nPlease check your connection settings and ensure 'influxd' is running.",
+			c.Client.Addr(), err.Error())
 	}
 
 	// Modify precision.
@@ -302,6 +303,7 @@ func (c *CommandLine) Connect(cmd string) error {
 
 	config := client.NewConfig()
 	config.URL = u
+	config.UnixSocket = c.UnixSocket
 	config.Username = c.Username
 	config.Password = c.Password
 	config.UserAgent = "InfluxDBShell/" + c.ClientVersion
@@ -315,7 +317,7 @@ func (c *CommandLine) Connect(cmd string) error {
 
 	var v string
 	if _, v, e = c.Client.Ping(); e != nil {
-		return fmt.Errorf("Failed to connect to %s\n", c.Client.Addr())
+		return fmt.Errorf("Failed to connect to %s: %s\n", c.Client.Addr(), e.Error())
 	}
 	c.ServerVersion = v
 	// Update the command with the current connection information
diff --git a/cmd/influx/main.go b/cmd/influx/main.go
index d182f2bf9..2d5b0e0d7 100644
--- a/cmd/influx/main.go
+++ b/cmd/influx/main.go
@@ -39,6 +39,7 @@ func main() {
 	fs := flag.NewFlagSet("InfluxDB shell version "+version, flag.ExitOnError)
 	fs.StringVar(&c.Host, "host", client.DefaultHost, "Influxdb host to connect to.")
 	fs.IntVar(&c.Port, "port", client.DefaultPort, "Influxdb port to connect to.")
+	fs.StringVar(&c.UnixSocket, "socket", c.UnixSocket, "Influxdb unix socket to connect to.")
 	fs.StringVar(&c.Username, "username", c.Username, "Username to connect to the server.")
 	fs.StringVar(&c.Password, "password", c.Password, `Password to connect to the server.  Leaving blank will prompt for password (--password="").`)
 	fs.StringVar(&c.Database, "database", c.Database, "Database to connect to the server.")
@@ -64,6 +65,8 @@ func main() {
        Host to connect to.
   -port 'port #'
        Port to connect to.
+  -socket 'unix domain socket'
+       Unix socket to connect to.
   -database 'database name'
        Database to connect to the server.
   -password 'password'
-- 
2.18.0