❄️

SPCSのServiceにgosnowflakeを使ってprogrammaticallyにアクセスする

に公開

ドキュメントにはPythonの例は載っているがgoを使ったアクセスは載っていないので備忘録として書いておく。
ref. https://docs.snowflake.com/en/developer-guide/snowpark-container-services/working-with-services#ingress-authentication

とりあえずPythonと同じことを実装しようと思って書いてるので、スーパー邪悪な実装になっている。

package main

import (
	"context"
	"database/sql"
	"fmt"
	"io"
	"net/http"
	"reflect"
	"time"
	"unsafe"

	"github.com/snowflakedb/gosnowflake"
)

var (
	account             = "<org>-<account>"
	user                = "<user>"
	authenticator       = gosnowflake.AuthTypeExternalBrowser
	database            = "<db>"
	schema              = "<schema>"
	role                = "<role>"
	spcsServiceEndpoint = "https://xxxxxx-<org>-<account>.snowflakecomputing.app/path/to/hello"
)

// Super-Darkness Method
func getSessionToken(driverConn any) (string, error) {
	v := reflect.ValueOf(driverConn)
	if v.Kind() != reflect.Ptr {
		return "", fmt.Errorf("driverConn is not a pointer")
	}
	v = v.Elem()
	if v.Kind() != reflect.Struct {
		return "", fmt.Errorf("driverConn elem is not a struct")
	}

	restField := v.FieldByName("rest")
	if !restField.IsValid() {
		return "", fmt.Errorf("rest field not found (driver changed?)")
	}
	if !restField.CanAddr() {
		return "", fmt.Errorf("rest field not addressable")
	}
	rest := reflect.NewAt(restField.Type(), unsafe.Pointer(restField.UnsafeAddr())).Elem()
	if rest.IsNil() {
		return "", fmt.Errorf("rest is nil")
	}
	restElem := rest
	if restElem.Kind() == reflect.Ptr {
		if restElem.IsNil() {
			return "", fmt.Errorf("rest pointer is nil")
		}
		restElem = restElem.Elem()
	}
	if restElem.Kind() != reflect.Struct {
		return "", fmt.Errorf("rest elem is not struct")
	}

	tokenAccessorField := restElem.FieldByName("TokenAccessor")
	if !tokenAccessorField.IsValid() {
		return "", fmt.Errorf("rest.TokenAccessor not found (driver changed?)")
	}
	ta, ok := tokenAccessorField.Interface().(gosnowflake.TokenAccessor)
	if !ok {
		return "", fmt.Errorf("rest.TokenAccessor is not gosnowflake.TokenAccessor")
	}

	token, _, _ := ta.GetTokens()
	if token == "" {
		return "", fmt.Errorf("empty session token")
	}
	return token, nil
}

func ptr(s string) *string { return &s }

func main() {
	// Build DSN with External Browser auth. Adjust as needed.
	cfg := &gosnowflake.Config{
		Account:       account,
		User:          user,
		Authenticator: authenticator,
		Database:      database,
		Schema:        schema,
		Role:          role,
		Params: map[string]*string{
			"GO_QUERY_RESULT_FORMAT": ptr("json"),
			"QUERY_TAG":              ptr("gosnowflake_ingress_demo"),
		},
	}
	dsn, err := gosnowflake.DSN(cfg)
	if err != nil {
		panic(err)
	}

	db, err := sql.Open("snowflake", dsn)
	if err != nil {
		panic(err)
	}
	defer db.Close()

	ctx, cancel := context.WithTimeout(context.Background(), 15*time.Minute)
	defer cancel()

	conn, err := db.Conn(ctx)
	if err != nil {
		panic(err)
	}
	defer conn.Close()

	var token string
	err = conn.Raw(func(dc any) error {
		tok, e := getSessionToken(dc)
		if e != nil {
			return fmt.Errorf("failed to extract token: %w", e)
		}
		token = tok
		return nil
	})
	if err != nil {
		panic(err)
	}

	req, err := http.NewRequestWithContext(ctx, http.MethodGet, spcsServiceEndpoint, nil)
	if err != nil {
		panic(err)
	}

	authz := fmt.Sprintf(`Snowflake Token="%s"`, token)
	req.Header.Set("Authorization", authz)

	httpClient := &http.Client{Timeout: 30 * time.Second}
	res, err := httpClient.Do(req)
	if err != nil {
		panic(err)
	}
	defer res.Body.Close()

	b, _ := io.ReadAll(res.Body)
	fmt.Printf("status=%s\nbody=%s\n", res.Status, string(b))
}

Discussion