Initial commit

Signed-off-by: nithinkdb <nithin.krishnamurthi@databricks.com>

Author: nithinkdb

connection.go

@@ -94,6 +94,13 @@ func (c *conn) IsValid() bool {
 // ExecContext honors the context timeout and return when it is canceled.
 // Statement ExecContext is the same as connection ExecContext
 func (c *conn) ExecContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Result, error) {
+	s, ok := ctx.(StagingCtx)
+	if ok {
+		if s.IsStagingOperation {
+			return c.ExecStagingOperation(s, query, args)
+		}
+	}
+
 	corrId := driverctx.CorrelationIdFromContext(ctx)
 	log := logger.WithContext(c.id, corrId, "")
 	msg, start := logger.Track("ExecContext")
@@ -133,6 +140,22 @@ func (c *conn) ExecContext(ctx context.Context, query string, args []driver.Name
 	return &res, nil
 }
 
+type StagingRow struct {
+	presignedUrl string
+	localFile    string
+	headers      string
+	operation    string
+}
+
+func (c *conn) ExecStagingOperation(ctx StagingCtx, query string, args []driver.NamedValue) (driver.Result, error) {
+	row, err := c.QueryContext(ctx, query, args)
+	if err != nil {
+		return nil, err
+	}
+
+	row.Next()
+}
+
 // QueryContext executes a query that may return rows, such as a
 // SELECT.
 //

connector.go

@@ -25,6 +25,36 @@ type connector struct {
 	client *http.Client
 }
 
+type StagingCtx struct {
+	IsStagingOperation      bool
+	StagingAllowedLocalPath string
+}
+
+func (ctx *StagingCtx) WithDefaults() {
+	ctx.IsStagingOperation = true
+	ctx.StagingAllowedLocalPath = "staging/"
+}
+
+func (StagingCtx) Deadline() (deadline time.Time, ok bool) {
+	return
+}
+
+func (StagingCtx) Done() <-chan struct{} {
+	return nil
+}
+
+func (StagingCtx) Err() error {
+	return nil
+}
+
+func (StagingCtx) Value(key any) any {
+	return nil
+}
+
+func (StagingCtx) String() string {
+	return "context.Background"
+}
+
 // Connect returns a connection to the Databricks database from a connection pool.
 func (c *connector) Connect(ctx context.Context) (driver.Conn, error) {
 	var catalogName *cli_service.TIdentifier

internal/config/config.go

@@ -82,24 +82,25 @@ func (c *Config) DeepCopy() *Config {
 
 // UserConfig is the set of configurations exposed to users
 type UserConfig struct {
-	Protocol          string
-	Host              string // from databricks UI
-	Port              int    // from databricks UI
-	HTTPPath          string // from databricks UI
-	Catalog           string
-	Schema            string
-	Authenticator     auth.Authenticator
-	AccessToken       string        // from databricks UI
-	MaxRows           int           // max rows per page
-	QueryTimeout      time.Duration // Timeout passed to server for query processing
-	UserAgentEntry    string
-	Location          *time.Location
-	SessionParams     map[string]string
-	RetryWaitMin      time.Duration
-	RetryWaitMax      time.Duration
-	RetryMax          int
-	Transport         http.RoundTripper
-	UseLz4Compression bool
+	Protocol                string
+	Host                    string // from databricks UI
+	Port                    int    // from databricks UI
+	HTTPPath                string // from databricks UI
+	Catalog                 string
+	Schema                  string
+	Authenticator           auth.Authenticator
+	AccessToken             string        // from databricks UI
+	MaxRows                 int           // max rows per page
+	QueryTimeout            time.Duration // Timeout passed to server for query processing
+	UserAgentEntry          string
+	Location                *time.Location
+	SessionParams           map[string]string
+	RetryWaitMin            time.Duration
+	RetryWaitMax            time.Duration
+	RetryMax                int
+	Transport               http.RoundTripper
+	UseLz4Compression       bool
+	StagingAllowedLocalPath string
 	CloudFetchConfig
 }
 
@@ -123,25 +124,26 @@ func (ucfg UserConfig) DeepCopy() UserConfig {
 	}
 
 	return UserConfig{
-		Protocol:          ucfg.Protocol,
-		Host:              ucfg.Host,
-		Port:              ucfg.Port,
-		HTTPPath:          ucfg.HTTPPath,
-		Catalog:           ucfg.Catalog,
-		Schema:            ucfg.Schema,
-		Authenticator:     ucfg.Authenticator,
-		AccessToken:       ucfg.AccessToken,
-		MaxRows:           ucfg.MaxRows,
-		QueryTimeout:      ucfg.QueryTimeout,
-		UserAgentEntry:    ucfg.UserAgentEntry,
-		Location:          loccp,
-		SessionParams:     sessionParams,
-		RetryWaitMin:      ucfg.RetryWaitMin,
-		RetryWaitMax:      ucfg.RetryWaitMax,
-		RetryMax:          ucfg.RetryMax,
-		Transport:         ucfg.Transport,
-		UseLz4Compression: ucfg.UseLz4Compression,
-		CloudFetchConfig:  ucfg.CloudFetchConfig,
+		Protocol:                ucfg.Protocol,
+		Host:                    ucfg.Host,
+		Port:                    ucfg.Port,
+		HTTPPath:                ucfg.HTTPPath,
+		Catalog:                 ucfg.Catalog,
+		Schema:                  ucfg.Schema,
+		Authenticator:           ucfg.Authenticator,
+		AccessToken:             ucfg.AccessToken,
+		MaxRows:                 ucfg.MaxRows,
+		QueryTimeout:            ucfg.QueryTimeout,
+		UserAgentEntry:          ucfg.UserAgentEntry,
+		Location:                loccp,
+		SessionParams:           sessionParams,
+		RetryWaitMin:            ucfg.RetryWaitMin,
+		RetryWaitMax:            ucfg.RetryWaitMax,
+		RetryMax:                ucfg.RetryMax,
+		Transport:               ucfg.Transport,
+		UseLz4Compression:       ucfg.UseLz4Compression,
+		StagingAllowedLocalPath: ucfg.StagingAllowedLocalPath,
+		CloudFetchConfig:        ucfg.CloudFetchConfig,
 	}
 }
 
@@ -176,7 +178,7 @@ func (ucfg UserConfig) WithDefaults() UserConfig {
 	}
 	ucfg.UseLz4Compression = false
 	ucfg.CloudFetchConfig = CloudFetchConfig{}.WithDefaults()
-
+	ucfg.StagingAllowedLocalPath = "staging/"
 	return ucfg
 }
 
@@ -279,6 +281,11 @@ func ParseDSN(dsn string) (UserConfig, error) {
 		ucfg.Location, err = time.LoadLocation(timezone)
 	}
 
+	if params.Has("stagingAllowedLocalPath") {
+		ucfg.StagingAllowedLocalPath = params.Get("stagingAllowedLocalPath")
+	}
+	params.Del("stagingAllowedLocalPath")
+
 	// any left over params are treated as session params
 	if len(params.Values) > 0 {
 		sessionParams := make(map[string]string)