send all headers within payload if they are correctly prefixed

trinodb · Dec 29, 2021 · 2e18f55 · 2e18f55
1 parent f2c444c
commit 2e18f55
Showing 1 changed file with 25 additions and 16 deletions.
diff --git a/trino/trino.go b/trino/trino.go
@@ -100,20 +100,23 @@ var (
 )
 
 const (
-	preparedStatementHeader    = "X-Trino-Prepared-Statement"
+	trinoHeaderPrefix = `X-Trino-`
+
+	preparedStatementHeader    = trinoHeaderPrefix+"Prepared-Statement"
 	preparedStatementName      = "_trino_go"
-	trinoUserHeader            = "X-Trino-User"
-	trinoSourceHeader          = "X-Trino-Source"
-	trinoCatalogHeader         = "X-Trino-Catalog"
-	trinoSchemaHeader          = "X-Trino-Schema"
-	trinoSessionHeader         = "X-Trino-Session"
-	trinoSetCatalogHeader      = "X-Trino-Set-Catalog"
-	trinoSetSchemaHeader       = "X-Trino-Set-Schema"
-	trinoSetPathHeader         = "X-Trino-Set-Path"
-	trinoSetSessionHeader      = "X-Trino-Set-Session"
-	trinoClearSessionHeader    = "X-Trino-Clear-Session"
-	trinoSetRoleHeader         = "X-Trino-Set-Role"
-	trinoExtraCredentialHeader = "X-Trino-Extra-Credential"
+
+	trinoUserHeader            = trinoHeaderPrefix+`User`
+	trinoSourceHeader          = trinoHeaderPrefix+`Source`
+	trinoCatalogHeader         = trinoHeaderPrefix+`Catalog`
+	trinoSchemaHeader          = trinoHeaderPrefix+`Schema`
+	trinoSessionHeader         = trinoHeaderPrefix+`Session`
+	trinoSetCatalogHeader      = trinoHeaderPrefix+`Set-Catalog`
+	trinoSetSchemaHeader       = trinoHeaderPrefix+`Set-Schema`
+	trinoSetPathHeader         = trinoHeaderPrefix+`Set-Path`
+	trinoSetSessionHeader      = trinoHeaderPrefix+`Set-Session`
+	trinoClearSessionHeader    = trinoHeaderPrefix+`Clear-Session`
+	trinoSetRoleHeader         = trinoHeaderPrefix+`Set-Role`
+	trinoExtraCredentialHeader = trinoHeaderPrefix+`Extra-Credential`
 
 	KerberosEnabledConfig    = "KerberosEnabled"
 	kerberosKeytabPathConfig = "KerberosKeytabPath"
@@ -654,9 +657,15 @@ func (st *driverStmt) exec(ctx context.Context, args []driver.NamedValue) (*stmt
 			if err != nil {
 				return nil, err
 			}
-			if arg.Name == trinoUserHeader {
-				st.user = arg.Value.(string)
-				hs.Add(trinoUserHeader, st.user)
+
+			if strings.HasPrefix(arg.Name, trinoHeaderPrefix) {
+				headerValue := arg.Value.(string)
+
+				if arg.Name == trinoUserHeader {
+					st.user = headerValue
+				}
+
+				hs.Add(arg.Name, headerValue)
 			} else {
 				if hs.Get(preparedStatementHeader) == "" {
 					hs.Add(preparedStatementHeader, preparedStatementName+"="+url.QueryEscape(st.query))