Skip to content
Closed
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 18 additions & 18 deletions job.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,16 +2,16 @@ package main

import (
"context"
"crypto/rsa"
"crypto/x509"
"encoding/pem"
"fmt"
"net/url"
"os"
"regexp"
"strconv"
"strings"
"time"
"crypto/rsa"
"crypto/x509"
"encoding/pem"

_ "github.com/ClickHouse/clickhouse-go/v2" // register the ClickHouse driver
"github.com/cenkalti/backoff"
Expand Down Expand Up @@ -401,32 +401,32 @@ func (j *Job) updateConnections() {
level.Error(j.log).Log("msg", "Failed to parse Snowflake URL", "url", conn, "err", err)
continue
}

queryParams := u.Query()
privateKeyPath := os.ExpandEnv(queryParams.Get("private_key_file"))

cfg := &gosnowflake.Config{
Account: u.Host,
User: u.User.Username(),
Role: queryParams.Get("role"),
Database: queryParams.Get("database"),
Schema: queryParams.Get("schema"),
}

if privateKeyPath != "" {
// RSA key auth
keyBytes, err := os.ReadFile(privateKeyPath)
if err != nil {
level.Error(j.log).Log("msg", "Failed to read private key file", "path", privateKeyPath, "err", err)
continue
}

keyBlock, _ := pem.Decode(keyBytes)
if keyBlock == nil {
level.Error(j.log).Log("msg", "Failed to decode PEM block", "path", privateKeyPath)
continue
}

var privateKey *rsa.PrivateKey
if parsedKey, err := x509.ParsePKCS8PrivateKey(keyBlock.Bytes); err == nil {
privateKey, _ = parsedKey.(*rsa.PrivateKey)
Expand All @@ -436,16 +436,16 @@ func (j *Job) updateConnections() {
level.Error(j.log).Log("msg", "Failed to parse private key", "err", err)
continue
}

cfg.Authenticator = gosnowflake.AuthTypeJwt
cfg.PrivateKey = privateKey

dsn, err := gosnowflake.DSN(cfg)
if err != nil {
level.Error(j.log).Log("msg", "Failed to create Snowflake DSN with RSA", "err", err)
continue
}

newConn.snowflakeConfig = cfg
newConn.snowflakeDSN = dsn
newConn.host = u.Host
Expand All @@ -460,20 +460,20 @@ func (j *Job) updateConnections() {
cfg.Port = port
}
}

dsn, err := gosnowflake.DSN(cfg)
if err != nil {
level.Error(j.log).Log("msg", "Failed to create Snowflake DSN with password", "err", err)
continue
}

newConn.conn, err = sqlx.Open("snowflake", dsn)
if err != nil {
level.Error(j.log).Log("msg", "Failed to open Snowflake connection", "err", err)
continue
}
}

j.conns = append(j.conns, newConn)
continue
}
Expand Down Expand Up @@ -651,21 +651,21 @@ func (c *connection) connect(job *Job) error {
}
c.tokenExpirationTime = time.Now().Add(time.Hour)
}

db, err := sqlx.Open("snowflake", c.snowflakeDSN)
if err != nil {
return fmt.Errorf("failed to open Snowflake connection: %w (host: %s)", err, c.host)
}

db.SetMaxOpenConns(1)
db.SetMaxIdleConns(0)
db.SetConnMaxLifetime(30 * time.Minute)

if err := db.Ping(); err != nil {
db.Close()
return fmt.Errorf("failed to ping Snowflake: %w (host: %s)", err, c.host)
}

c.conn = db
return nil
}
Expand Down