Repository: googleapis/genai-toolbox
File: internal/sources/trino/trino.go (Go)

Code before:

func (s *Source) SourceKind() string {
	return SourceKind
}

func (s *Source) ToConfig() sources.SourceConfig {
	return s.Config
}

func (s *Source) TrinoDB() *sql.DB {
	return s.Pool
}

func initTrinoConnectionPool(ctx context.Context, tracer trace.Tracer, name, host, port, user, password, catalog, schema, queryTimeout, accessToken string, kerberosEnabled, sslEnabled bool, sslCertPath, sslCert string, disableSslVerification bool) (*sql.DB, error) {
	//nolint:all // Reassigned ctx
	ctx, span := sources.InitConnectionSpan(ctx, tracer, SourceKind, name)
	defer span.End()

	// Build Trino DSN
	dsn, err := buildTrinoDSN(host, port, user, password, catalog, schema, queryTimeout, accessToken, kerberosEnabled, sslEnabled, sslCertPath, sslCert)
	if err != nil {
		return nil, fmt.Errorf("failed to build DSN: %w", err)
	}

	if disableSslVerification {
		fmt.Printf("WARNING: SSL verification is disabled for trino source %q. This is an insecure setting and should not be used in production.", name)
		tr := &http.Transport{
			TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
		}
		client := &http.Client{Transport: tr}
		clientName := fmt.Sprintf("insecure_trino_client_%s", name)
		if err := trinogo.RegisterCustomClient(clientName, client); err != nil {
			return nil, fmt.Errorf("failed to register custom client: %w", err)
		}
		dsn = fmt.Sprintf("%s&custom_client=%s", dsn, clientName)
	}

	db, err := sql.Open("trino", dsn)
	if err != nil {
		return nil, fmt.Errorf("failed to open connection: %w", err)
	}

	// Configure connection pool
	db.SetMaxOpenConns(10)
	db.SetMaxIdleConns(5)
	db.SetConnMaxLifetime(time.Hour)

	return db, nil
}

func buildTrinoDSN(host, port, user, password, catalog, schema, queryTimeout, accessToken string, kerberosEnabled, sslEnabled bool, sslCertPath, sslCert string) (string, error) {

Review comment (suggestion):
Instead of using `fmt.Printf`, let's grab the logger from ctx and use the `WarnContext()` to log a warn. (example: https://github.com/googleapis/genai-toolbox/blob/0f27f956c7d85e6f8a3f4c4589cf8020bc1d3d3b/internal/sources/http/http.go#L72)

Code after:
}

func (s *Source) SourceKind() string {
	return SourceKind
}

func (s *Source) ToConfig() sources.SourceConfig {
	return s.Config
}

func (s *Source) TrinoDB() *sql.DB {
	return s.Pool
}

func (s *Source) RunSQL(ctx context.Context, statement string, params []any) (any, error) {
	results, err := s.TrinoDB().QueryContext(ctx, statement, params...)
	if err != nil {
		return nil, fmt.Errorf("unable to execute query: %w", err)
	}
	defer results.Close()

	cols, err := results.Columns()
	if err != nil {
		return nil, fmt.Errorf("unable to retrieve column names: %w", err)
	}

	// create an array of values for each column, which can be re-used to scan each row
	rawValues := make([]any, len(cols))
	values := make([]any, len(cols))
	for i := range rawValues {
		values[i] = &rawValues[i]
	}

	var out []any
	for results.Next() {
		err := results.Scan(values...)
		if err != nil {
			return nil, fmt.Errorf("unable to parse row: %w", err)
		}
		vMap := make(map[string]any)
		for i, name := range cols {
			val := rawValues[i]
			if val == nil {
				vMap[name] = nil
				continue
			}

			// Convert byte arrays to strings for text fields
			if b, ok := val.([]byte); ok {
				vMap[name] = string(b)
			} else {