added dbfs datasource and resource objects for uploading files

Author: stikkireddy

Makefile

@@ -29,11 +29,11 @@ python-setup:
 	@cd docs && python -m virtualenv venv
 	@cd docs && source venv/bin/activate && python -m pip install -r requirements.txt
 
-docs:
+docs: python-setup
 	@echo "==> Building Docs ..."
 	@cd docs && source venv/bin/activate && make clean && make html
 
-opendocs: docs
+opendocs: python-setup docs
 	@echo "==> Opening Docs ..."
 	@cd docs && open build/html/index.html
 

client/client.go

@@ -96,6 +96,14 @@ func (o *DBClientOption) getRequestURI(path string, apiVersion string) (string,
 	return requestURI, nil
 }
 
+func onlyNBytes(buf []byte, numBytes int64) []byte {
+	if len(buf) > int(numBytes) {
+		return buf[:numBytes]
+	} else {
+		return buf
+	}
+}
+
 // PerformQuery can be used in a client or directly
 func PerformQuery(option DBClientOption, method, path string, apiVersion string, headers map[string]string, marshalJson bool, useRawPath bool, data interface{}) ([]byte, error) {
 
@@ -133,7 +141,7 @@ func PerformQuery(option DBClientOption, method, path string, apiVersion string,
 			}
 
 			requestBody = bodyBytes
-			log.Println(string(requestBody))
+			log.Println(string(onlyNBytes(requestBody, 1e3)))
 		} else {
 			requestBody = []byte(data.(string))
 		}

client/model/dbfs.go

@@ -0,0 +1,7 @@
+package model
+
+type FileInfo struct {
+	Path     string `json:"path,omitempty"`
+	IsDir    bool   `json:"is_dir,omitempty"`
+	FileSize int64  `json:"file_size,omitempty"`
+}

client/service/api.go

@@ -65,6 +65,11 @@ func (c DBApiClient) Jobs() JobsAPI {
 	return jobsAPI.init(c)
 }
 
+func (c DBApiClient) DBFS() DBFSAPI {
+	var DBFSAPI DBFSAPI
+	return DBFSAPI.init(c)
+}
+
 func (c DBApiClient) InstancePools() InstancePoolsAPI {
 	var instancePoolsAPI InstancePoolsAPI
 	return instancePoolsAPI.init(c)

client/service/dbfs.go

@@ -0,0 +1,223 @@
+package service
+
+import (
+	"encoding/base64"
+	"encoding/json"
+	"github.com/databrickslabs/databricks-terraform/client/model"
+	"log"
+	"net/http"
+)
+
+// TokensAPI exposes the Secrets API
+type DBFSAPI struct {
+	Client DBApiClient
+}
+
+func (a DBFSAPI) init(client DBApiClient) DBFSAPI {
+	a.Client = client
+	return a
+}
+
+func (a DBFSAPI) Create(path string, overwrite bool, data string) error {
+	byteArr, err := base64.StdEncoding.DecodeString(data)
+	if err != nil {
+		return err
+	}
+	byteChunks := split(byteArr, 1e6)
+	handle, err := a.createHandle(path, overwrite)
+	if err != nil {
+		return err
+	}
+	for _, byteChunk := range byteChunks {
+		b64Data := base64.StdEncoding.EncodeToString(byteChunk)
+		err := a.addBlock(b64Data, handle)
+		if err != nil {
+			return err
+		}
+	}
+	err = a.closeHandle(handle)
+	return err
+}
+
+func (a DBFSAPI) Read(path string) (string, error) {
+	var bytesFetched []byte
+	fetchLoop := true
+	offSet := int64(0)
+	length := int64(1e6)
+	for fetchLoop == true {
+		bytesRead, bytes, err := a.read(path, offSet, length)
+		if err != nil {
+			return "", err
+		}
+		log.Println(bytesRead)
+		if bytesRead == 0 || bytesRead < length {
+			fetchLoop = false
+		}
+
+		bytesFetched = append(bytesFetched, bytes...)
+		offSet = offSet + length
+	}
+	resp := base64.StdEncoding.EncodeToString(bytesFetched)
+	return resp, nil
+}
+
+func (a DBFSAPI) Delete(path string, recursive bool) error {
+	deleteRequest := struct {
+		Path      string `json:"path,omitempty" url:"path,omitempty"`
+		Recursive bool   `json:"recursive,omitempty" url:"recursive,omitempty"`
+	}{
+		Path:      path,
+		Recursive: recursive,
+	}
+	_, err := a.Client.performQuery(http.MethodPost, "/dbfs/delete", "2.0", nil, deleteRequest)
+
+	return err
+}
+
+func (a DBFSAPI) read(path string, offset, length int64) (int64, []byte, error) {
+	var readBytes struct {
+		BytesRead int64  `json:"bytes_read,omitempty" url:"bytes_read,omitempty"`
+		Data      string `json:"data,omitempty" url:"data,omitempty"`
+	}
+	readRequest := struct {
+		Path   string `json:"path,omitempty" url:"path,omitempty"`
+		Offset int64  `json:"offset,omitempty" url:"offset,omitempty"`
+		Length int64  `json:"length,omitempty" url:"length,omitempty"`
+	}{
+		Path:   path,
+		Offset: offset,
+		Length: length,
+	}
+	resp, err := a.Client.performQuery(http.MethodGet, "/dbfs/read", "2.0", nil, readRequest)
+	if err != nil {
+		return readBytes.BytesRead, []byte{}, err
+	}
+	err = json.Unmarshal(resp, &readBytes)
+	if err != nil {
+		return 0, []byte{}, err
+	}
+	dataBytes, err := base64.StdEncoding.DecodeString(readBytes.Data)
+	return readBytes.BytesRead, dataBytes, err
+}
+
+func (a DBFSAPI) Status(path string) (model.FileInfo, error) {
+	var fileInfo model.FileInfo
+	statusRequest := struct {
+		Path string `json:"path,omitempty" url:"path,omitempty"`
+	}{
+		Path: path,
+	}
+	resp, err := a.Client.performQuery(http.MethodGet, "/dbfs/get-status", "2.0", nil, statusRequest)
+	if err != nil {
+		return fileInfo, err
+	}
+	err = json.Unmarshal(resp, &fileInfo)
+	return fileInfo, err
+}
+
+func (a DBFSAPI) List(path string, recursive bool) ([]model.FileInfo, error) {
+	if recursive == true {
+		var paths []model.FileInfo
+		a.recursiveAddPaths(path, &paths)
+		return paths, nil
+	} else {
+		return a.list(path)
+	}
+}
+
+func (a DBFSAPI) recursiveAddPaths(path string, pathList *[]model.FileInfo) {
+	fileInfoList, _ := a.list(path)
+	for _, v := range fileInfoList {
+		if v.IsDir == false {
+			*pathList = append(*pathList, v)
+		} else if v.IsDir == true {
+			a.recursiveAddPaths(v.Path, pathList)
+		}
+	}
+}
+
+func (a DBFSAPI) list(path string) ([]model.FileInfo, error) {
+	var dbfsList struct {
+		Files []model.FileInfo `json:"files,omitempty" url:"files,omitempty"`
+	}
+	listRequest := struct {
+		Path string `json:"path,omitempty" url:"path,omitempty"`
+	}{}
+	listRequest.Path = path
+
+	resp, err := a.Client.performQuery(http.MethodGet, "/dbfs/list", "2.0", nil, listRequest)
+	if err != nil {
+		return dbfsList.Files, err
+	}
+
+	err = json.Unmarshal(resp, &dbfsList)
+	return dbfsList.Files, err
+}
+
+func (a DBFSAPI) Mkdirs(path string) error {
+	mkDirsRequest := struct {
+		Path string `json:"path,omitempty" url:"path,omitempty"`
+	}{}
+	mkDirsRequest.Path = path
+
+	_, err := a.Client.performQuery(http.MethodPost, "/dbfs/mkdirs", "2.0", nil, mkDirsRequest)
+
+	return err
+}
+
+func (a DBFSAPI) createHandle(path string, overwrite bool) (int64, error) {
+	var handle struct {
+		Handle int64 `json:"handle,omitempty" url:"handle,omitempty"`
+	}
+	createDBFSHandleRequest := struct {
+		Path      string `json:"path,omitempty" url:"path,omitempty"`
+		Overwrite bool   `json:"overwrite,omitempty" url:"overwrite,omitempty"`
+	}{
+		Path:      path,
+		Overwrite: overwrite,
+	}
+
+	resp, err := a.Client.performQuery(http.MethodPost, "/dbfs/create", "2.0", nil, createDBFSHandleRequest)
+	if err != nil {
+		return handle.Handle, err
+	}
+
+	err = json.Unmarshal(resp, &handle)
+	return handle.Handle, err
+}
+
+func (a DBFSAPI) addBlock(data string, handle int64) error {
+	var addDBFSBlockRequest = struct {
+		Data   string `json:"data,omitempty" url:"data,omitempty"`
+		Handle int64  `json:"handle,omitempty" url:"handle,omitempty"`
+	}{
+		Data:   data,
+		Handle: handle,
+	}
+	_, err := a.Client.performQuery(http.MethodPost, "/dbfs/add-block", "2.0", nil, addDBFSBlockRequest)
+	return err
+}
+
+func (a DBFSAPI) closeHandle(handle int64) error {
+	closeHandleRequest := struct {
+		Handle int64 `json:"handle,omitempty" url:"handle,omitempty"`
+	}{
+		Handle: handle,
+	}
+
+	_, err := a.Client.performQuery(http.MethodPost, "/dbfs/close", "2.0", nil, closeHandleRequest)
+	return err
+}
+
+func split(buf []byte, lim int) [][]byte {
+	var chunk []byte
+	chunks := make([][]byte, 0, len(buf)/lim+1)
+	for len(buf) >= lim {
+		chunk, buf = buf[:lim], buf[lim:]
+		chunks = append(chunks, chunk)
+	}
+	if len(buf) > 0 {
+		chunks = append(chunks, buf[:len(buf)])
+	}
+	return chunks
+}

client/service/dbfs_integration_test.go

@@ -0,0 +1,82 @@
+package service
+
+import (
+	"bytes"
+	"encoding/base64"
+	"github.com/stretchr/testify/assert"
+	"io/ioutil"
+	"strings"
+	"testing"
+)
+
+func TestRead(t *testing.T) {
+	dat, _ := ioutil.ReadFile("/Users/Sri.Tikkireddy/Downloads/DatabricksConcepts&BestPractices.pdf")
+	t.Log(len(dat))
+}
+
+func TestSplit(t *testing.T) {
+	dat, _ := ioutil.ReadFile("/Users/Sri.Tikkireddy/Downloads/DatabricksConcepts&BestPractices.pdf")
+	t.Log(len(dat))
+	t.Log(len(split(dat, 1e6)))
+}
+
+func GenString(times int) []byte {
+	var buf bytes.Buffer
+	for i := 0; i < times; i++ {
+		buf.WriteString("Hello world how are you doing?\n")
+	}
+	return buf.Bytes()
+}
+
+func TestCreateFile(t *testing.T) {
+	if testing.Short() {
+		t.Skip("skipping integration test in short mode.")
+	}
+
+	path := "/sri/randomfile"
+
+	randomStr := GenString(500000)
+	t.Log(len(randomStr))
+	t.Log(len(base64.StdEncoding.EncodeToString(randomStr)))
+
+	client := GetIntegrationDBAPIClient()
+	err := client.DBFS().Create(path, true, base64.StdEncoding.EncodeToString(randomStr))
+	assert.NoError(t, err, err)
+
+	err = client.DBFS().Delete(path, false)
+	assert.NoError(t, err, err)
+}
+
+func TestReadFile(t *testing.T) {
+	if testing.Short() {
+		t.Skip("skipping integration test in short mode.")
+	}
+
+	path := "/sri/randomfile"
+	client := GetIntegrationDBAPIClient()
+	data, err := client.DBFS().Read(path)
+	assert.NoError(t, err, err)
+
+	byteArr, err := base64.StdEncoding.DecodeString(data)
+	assert.NoError(t, err, err)
+
+	t.Log(len(strings.Split(string(byteArr), "\n")))
+	t.Log(strings.Split(string(byteArr), "\n")[0])
+	t.Log(strings.Split(string(byteArr), "\n")[1])
+	t.Log(strings.Split(string(byteArr), "\n")[2])
+	t.Log(strings.Split(string(byteArr), "\n")[499999])
+	t.Log(strings.Split(string(byteArr), "\n")[500000])
+}
+
+func TestListRecursive(t *testing.T) {
+	if testing.Short() {
+		t.Skip("skipping integration test in short mode.")
+	}
+
+	path := "/andre_mesarovic/mlflow"
+	client := GetIntegrationDBAPIClient()
+	data, err := client.DBFS().List(path, true)
+	assert.NoError(t, err, err)
+
+	t.Log(data)
+}