Fallback in functions and aggregates metadata queries for yugabyte

Author: jameshartig

metadata.go

@@ -7,6 +7,7 @@ package gocql
 import (
 	"encoding/hex"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"strconv"
 	"strings"
@@ -1048,13 +1049,17 @@ func getFunctionsMetadata(session *Session, keyspaceName string) ([]FunctionMeta
 	if session.cfg.ProtoVersion == protoVersion1 || !session.hasAggregatesAndFunctions {
 		return nil, nil
 	}
+	var withoutBody bool
 	var tableName string
 	if session.useSystemSchema {
 		tableName = "system_schema.functions"
 	} else {
 		tableName = "system.schema_functions"
 	}
-	stmt := fmt.Sprintf(`
+FuncsStmt:
+	var stmt string
+	if !withoutBody {
+		stmt = fmt.Sprintf(`
 		SELECT
 			function_name,
 			argument_types,
@@ -1065,6 +1070,18 @@ func getFunctionsMetadata(session *Session, keyspaceName string) ([]FunctionMeta
 			return_type
 		FROM %s
 		WHERE keyspace_name = ?`, tableName)
+	} else {
+		stmt = fmt.Sprintf(`
+		SELECT
+			function_name,
+			argument_types,
+			argument_names,
+			called_on_null_input,
+			language,
+			return_type
+		FROM %s
+		WHERE keyspace_name = ?`, tableName)
+	}
 
 	var functions []FunctionMetadata
 
@@ -1073,14 +1090,25 @@ func getFunctionsMetadata(session *Session, keyspaceName string) ([]FunctionMeta
 		function := FunctionMetadata{Keyspace: keyspaceName}
 		var argumentTypes []string
 		var returnType string
-		err := rows.Scan(&function.Name,
-			&argumentTypes,
-			&function.ArgumentNames,
-			&function.Body,
-			&function.CalledOnNullInput,
-			&function.Language,
-			&returnType,
-		)
+		var err error
+		if !withoutBody {
+			err = rows.Scan(&function.Name,
+				&argumentTypes,
+				&function.ArgumentNames,
+				&function.Body,
+				&function.CalledOnNullInput,
+				&function.Language,
+				&returnType,
+			)
+		} else {
+			err = rows.Scan(&function.Name,
+				&argumentTypes,
+				&function.ArgumentNames,
+				&function.CalledOnNullInput,
+				&function.Language,
+				&returnType,
+			)
+		}
 		if err != nil {
 			return nil, err
 		}
@@ -1093,6 +1121,14 @@ func getFunctionsMetadata(session *Session, keyspaceName string) ([]FunctionMeta
 	}
 
 	if err := rows.Err(); err != nil {
+		if !withoutBody {
+			// Yugabyte doesn't support the body column in the functions table
+			var rerr RequestError
+			if errors.As(err, &rerr) && rerr.Code() == ErrCodeInvalid {
+				withoutBody = true
+				goto FuncsStmt
+			}
+		}
 		return nil, err
 	}
 
@@ -1103,14 +1139,17 @@ func getAggregatesMetadata(session *Session, keyspaceName string) ([]AggregateMe
 	if session.cfg.ProtoVersion == protoVersion1 || !session.hasAggregatesAndFunctions {
 		return nil, nil
 	}
+	var withoutReturnType bool
 	var tableName string
 	if session.useSystemSchema {
 		tableName = "system_schema.aggregates"
 	} else {
 		tableName = "system.schema_aggregates"
 	}
-
-	stmt := fmt.Sprintf(`
+AggsStmt:
+	var stmt string
+	if !withoutReturnType {
+		stmt = fmt.Sprintf(`
 		SELECT
 			aggregate_name,
 			argument_types,
@@ -1121,6 +1160,18 @@ func getAggregatesMetadata(session *Session, keyspaceName string) ([]AggregateMe
 			state_type
 		FROM %s
 		WHERE keyspace_name = ?`, tableName)
+	} else {
+		stmt = fmt.Sprintf(`
+		SELECT
+			aggregate_name,
+			argument_types,
+			final_func,
+			initcond,
+			state_func,
+			state_type
+		FROM %s
+		WHERE keyspace_name = ?`, tableName)
+	}
 
 	var aggregates []AggregateMetadata
 
@@ -1130,14 +1181,25 @@ func getAggregatesMetadata(session *Session, keyspaceName string) ([]AggregateMe
 		var argumentTypes []string
 		var returnType string
 		var stateType string
-		err := rows.Scan(&aggregate.Name,
-			&argumentTypes,
-			&aggregate.finalFunc,
-			&aggregate.InitCond,
-			&returnType,
-			&aggregate.stateFunc,
-			&stateType,
-		)
+		var err error
+		if !withoutReturnType {
+			err = rows.Scan(&aggregate.Name,
+				&argumentTypes,
+				&aggregate.finalFunc,
+				&aggregate.InitCond,
+				&returnType,
+				&aggregate.stateFunc,
+				&stateType,
+			)
+		} else {
+			err = rows.Scan(&aggregate.Name,
+				&argumentTypes,
+				&aggregate.finalFunc,
+				&aggregate.InitCond,
+				&aggregate.stateFunc,
+				&stateType,
+			)
+		}
 		if err != nil {
 			return nil, err
 		}
@@ -1151,6 +1213,14 @@ func getAggregatesMetadata(session *Session, keyspaceName string) ([]AggregateMe
 	}
 
 	if err := rows.Err(); err != nil {
+		if !withoutReturnType {
+			// Yugabyte doesn't support the return_type column in the aggregates table
+			var rerr RequestError
+			if errors.As(err, &rerr) && rerr.Code() == ErrCodeInvalid {
+				withoutReturnType = true
+				goto AggsStmt
+			}
+		}
 		return nil, err
 	}