inintal package

Author: mashiike

client.go

@@ -0,0 +1,33 @@
+package redshiftdatasqldriver
+
+import (
+	"context"
+
+	"github.com/aws/aws-sdk-go-v2/config"
+	"github.com/aws/aws-sdk-go-v2/service/redshiftdata"
+)
+
+type RedshiftDataClient interface {
+	ExecuteStatement(ctx context.Context, params *redshiftdata.ExecuteStatementInput, optFns ...func(*redshiftdata.Options)) (*redshiftdata.ExecuteStatementOutput, error)
+	DescribeStatement(ctx context.Context, params *redshiftdata.DescribeStatementInput, optFns ...func(*redshiftdata.Options)) (*redshiftdata.DescribeStatementOutput, error)
+	CancelStatement(ctx context.Context, params *redshiftdata.CancelStatementInput, optFns ...func(*redshiftdata.Options)) (*redshiftdata.CancelStatementOutput, error)
+	redshiftdata.GetStatementResultAPIClient
+}
+
+var RedshiftDataClientConstructor func(ctx context.Context, cfg *RedshiftDataConfig) (RedshiftDataClient, error)
+
+func newRedshiftDataClient(ctx context.Context, cfg *RedshiftDataConfig) (RedshiftDataClient, error) {
+	if RedshiftDataClientConstructor != nil {
+		return RedshiftDataClientConstructor(ctx, cfg)
+	}
+	return DefaultRedshiftDataClientConstructor(ctx, cfg)
+}
+
+func DefaultRedshiftDataClientConstructor(ctx context.Context, cfg *RedshiftDataConfig) (RedshiftDataClient, error) {
+	awsCfg, err := config.LoadDefaultConfig(ctx)
+	if err != nil {
+		return nil, err
+	}
+	client := redshiftdata.NewFromConfig(awsCfg, cfg.RedshiftDataOptFns...)
+	return client, nil
+}

conn.go

@@ -0,0 +1,197 @@
+package redshiftdatasqldriver
+
+import (
+	"context"
+	"database/sql/driver"
+	"fmt"
+	"time"
+
+	"github.com/aws/aws-sdk-go-v2/aws"
+	"github.com/aws/aws-sdk-go-v2/service/redshiftdata"
+	"github.com/aws/aws-sdk-go-v2/service/redshiftdata/types"
+)
+
+type redshiftDataConn struct {
+	client   RedshiftDataClient
+	cfg      *RedshiftDataConfig
+	aliveCh  chan struct{}
+	isClosed bool
+}
+
+func newConn(client RedshiftDataClient, cfg *RedshiftDataConfig) *redshiftDataConn {
+	return &redshiftDataConn{
+		client:  client,
+		cfg:     cfg,
+		aliveCh: make(chan struct{}),
+	}
+}
+
+func (conn *redshiftDataConn) PrepareContext(ctx context.Context, query string) (driver.Stmt, error) {
+	return nil, fmt.Errorf("prepared statment %w", ErrNotSupported)
+}
+
+func (conn *redshiftDataConn) Prepare(query string) (driver.Stmt, error) {
+	return conn.PrepareContext(context.Background(), query)
+}
+
+func (conn *redshiftDataConn) Close() error {
+	if conn.isClosed {
+		return nil
+	}
+	conn.isClosed = true
+	close(conn.aliveCh)
+	return nil
+}
+
+func (conn *redshiftDataConn) BeginTx(ctx context.Context, opts driver.TxOptions) (driver.Tx, error) {
+	return nil, fmt.Errorf("transaction %w", ErrNotSupported)
+}
+
+func (conn *redshiftDataConn) Begin() (driver.Tx, error) {
+	return conn.BeginTx(context.Background(), driver.TxOptions{})
+}
+
+func (conn *redshiftDataConn) QueryContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Rows, error) {
+	params := &redshiftdata.ExecuteStatementInput{
+		Sql: nullif(query),
+	}
+	if len(args) > 0 {
+		params.Parameters = make([]types.SqlParameter, 0, len(args))
+		for _, arg := range args {
+			params.Parameters = append(params.Parameters, types.SqlParameter{
+				Name:  aws.String(arg.Name),
+				Value: aws.String(fmt.Sprintf("%v", arg.Value)),
+			})
+		}
+	}
+	p, output, err := conn.executeStatement(ctx, params)
+	if err != nil {
+		return nil, err
+	}
+	rows := newRows(coalesce(output.Id), p)
+	return rows, nil
+}
+
+func (conn *redshiftDataConn) ExecContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Result, error) {
+	params := &redshiftdata.ExecuteStatementInput{
+		Sql: nullif(query),
+	}
+	if len(args) > 0 {
+		params.Parameters = make([]types.SqlParameter, 0, len(args))
+		for _, arg := range args {
+			params.Parameters = append(params.Parameters, types.SqlParameter{
+				Name:  aws.String(arg.Name),
+				Value: aws.String(fmt.Sprintf("%v", arg.Value)),
+			})
+		}
+	}
+	_, output, err := conn.executeStatement(ctx, params)
+	if err != nil {
+		return nil, err
+	}
+	return newResult(output), nil
+}
+
+func (conn *redshiftDataConn) executeStatement(ctx context.Context, params *redshiftdata.ExecuteStatementInput) (*redshiftdata.GetStatementResultPaginator, *redshiftdata.DescribeStatementOutput, error) {
+	debugLogger.Printf("query: %s", coalesce(params.Sql))
+	params.ClusterIdentifier = conn.cfg.ClusterIdentifier
+	params.Database = conn.cfg.Database
+	params.DbUser = conn.cfg.DbUser
+	params.WorkgroupName = conn.cfg.WorkgroupName
+	params.SecretArn = conn.cfg.SecretsARN
+
+	if conn.cfg.Timeout == 0 {
+		conn.cfg.Timeout = 15 * time.Minute
+	}
+	if conn.cfg.Polling == 0 {
+		conn.cfg.Polling = 10 * time.Millisecond
+	}
+	ectx, cancel := context.WithTimeout(ctx, conn.cfg.Timeout)
+	defer cancel()
+	executeOutput, err := conn.client.ExecuteStatement(ectx, params)
+	if err != nil {
+		return nil, nil, fmt.Errorf("execute statement:%w", err)
+	}
+	queryStart := time.Now()
+	debugLogger.Printf("[%s] sucess execute statement: %s", *executeOutput.Id, coalesce(params.Sql))
+	describeOutput, err := conn.client.DescribeStatement(ectx, &redshiftdata.DescribeStatementInput{
+		Id: executeOutput.Id,
+	})
+	if err != nil {
+		return nil, nil, fmt.Errorf("describe statement:%w", err)
+	}
+	debugLogger.Printf("[%s] describe statement: status=%s pid=%d query_id=%d", *executeOutput.Id, describeOutput.Status, describeOutput.RedshiftPid, describeOutput.RedshiftQueryId)
+
+	var isFinished bool
+	defer func() {
+		if !isFinished {
+			describeOutput, err := conn.client.DescribeStatement(ctx, &redshiftdata.DescribeStatementInput{
+				Id: executeOutput.Id,
+			})
+			if err != nil {
+				errLogger.Printf("[%s] failed describe statement: %v", *executeOutput.Id, err)
+				return
+			}
+			if describeOutput.Status == types.StatusStringFinished ||
+				describeOutput.Status == types.StatusStringFailed ||
+				describeOutput.Status == types.StatusStringAborted {
+				return
+			}
+			debugLogger.Printf("[%s] try cancel statement", *executeOutput.Id)
+			output, err := conn.client.CancelStatement(ctx, &redshiftdata.CancelStatementInput{
+				Id: executeOutput.Id,
+			})
+			if err != nil {
+
+				errLogger.Printf("[%s] failed cancel statement: %v", *executeOutput.Id, err)
+				return
+			}
+			if !*output.Status {
+				debugLogger.Printf("[%s] cancel statement status is false", *executeOutput.Id)
+			}
+		}
+	}()
+	delay := time.NewTimer(conn.cfg.Polling)
+	for {
+		if describeOutput.Status == types.StatusStringAborted {
+			return nil, nil, fmt.Errorf("query aborted: %s", *describeOutput.Error)
+		}
+		if describeOutput.Status == types.StatusStringFailed {
+			return nil, nil, fmt.Errorf("query failed: %s", *describeOutput.Error)
+		}
+		if describeOutput.Status == types.StatusStringFinished {
+			break
+		}
+		debugLogger.Printf("[%s] wating finsih query: elapsed_time=%s", *executeOutput.Id, time.Since(queryStart))
+		delay.Reset(conn.cfg.Polling)
+		select {
+		case <-ectx.Done():
+			if !delay.Stop() {
+				<-delay.C
+			}
+			return nil, nil, ectx.Err()
+		case <-delay.C:
+		case <-conn.aliveCh:
+			if !delay.Stop() {
+				<-delay.C
+			}
+			return nil, nil, ErrConnClosed
+		}
+		describeOutput, err = conn.client.DescribeStatement(ctx, &redshiftdata.DescribeStatementInput{
+			Id: executeOutput.Id,
+		})
+		if err != nil {
+			return nil, nil, fmt.Errorf("describe statement:%w", err)
+		}
+	}
+	isFinished = true
+	debugLogger.Printf("[%s] success query: elapsed_time=%s", *executeOutput.Id, time.Since(queryStart))
+	if !*describeOutput.HasResultSet {
+		return nil, describeOutput, nil
+	}
+	debugLogger.Printf("[%s] query has result set: result_rows=%d", *executeOutput.Id, describeOutput.ResultRows)
+	p := redshiftdata.NewGetStatementResultPaginator(conn.client, &redshiftdata.GetStatementResultInput{
+		Id: executeOutput.Id,
+	})
+	return p, describeOutput, nil
+}

connector.go

@@ -0,0 +1,23 @@
+package redshiftdatasqldriver
+
+import (
+	"context"
+	"database/sql/driver"
+)
+
+type redshiftDataConnector struct {
+	d   *redshiftDataDriver
+	cfg *RedshiftDataConfig
+}
+
+func (c *redshiftDataConnector) Connect(ctx context.Context) (driver.Conn, error) {
+	client, err := newRedshiftDataClient(ctx, c.cfg)
+	if err != nil {
+		return nil, err
+	}
+	return newConn(client, c.cfg), nil
+}
+
+func (c *redshiftDataConnector) Driver() driver.Driver {
+	return c.d
+}

driver.go

@@ -0,0 +1,32 @@
+package redshiftdatasqldriver
+
+import (
+	"context"
+	"database/sql"
+	"database/sql/driver"
+)
+
+func init() {
+	sql.Register("redshift-data", &redshiftDataDriver{})
+}
+
+type redshiftDataDriver struct{}
+
+func (d *redshiftDataDriver) Open(dsn string) (driver.Conn, error) {
+	connector, err := d.OpenConnector(dsn)
+	if err != nil {
+		return nil, err
+	}
+	return connector.Connect(context.Background())
+}
+
+func (d *redshiftDataDriver) OpenConnector(dsn string) (driver.Connector, error) {
+	cfg, err := ParseDSN(dsn)
+	if err != nil {
+		return nil, err
+	}
+	return &redshiftDataConnector{
+		d:   d,
+		cfg: cfg,
+	}, nil
+}