Add metadata for functions and aggregates (apache#1204)

* Add metadata for functions and aggregates * Review comments * Style nits * Review comments
luis-sousa-pinto · Oct 19, 2018 · 6832a79 · 6832a79
1 parent fa8dbda
commit 6832a79
Show file tree

Hide file tree

Showing 5 changed files with 351 additions and 3 deletions.
diff --git a/cassandra_test.go b/cassandra_test.go
@@ -2184,6 +2184,126 @@ func TestGetColumnMetadata(t *testing.T) {
 	}
 }
 
+func TestAggregateMetadata(t *testing.T) {
+	session := createSession(t)
+	defer session.Close()
+	createAggregate(t, session)
+
+	aggregates, err := getAggregatesMetadata(session, "gocql_test")
+	if err != nil {
+		t.Fatalf("failed to query aggregate metadata with err: %v", err)
+	}
+	if aggregates == nil {
+		t.Fatal("failed to query aggregate metadata, nil returned")
+	}
+	if len(aggregates) != 1 {
+		t.Fatal("expected only a single aggregate")
+	}
+	aggregate := aggregates[0]
+
+	expectedAggregrate := AggregateMetadata{
+		Keyspace:      "gocql_test",
+		Name:          "average",
+		ArgumentTypes: []TypeInfo{NativeType{typ: TypeInt}},
+		InitCond:      "(0, 0)",
+		ReturnType:    NativeType{typ: TypeDouble},
+		StateType: TupleTypeInfo{
+			NativeType: NativeType{typ: TypeTuple},
+
+			Elems: []TypeInfo{
+				NativeType{typ: TypeInt},
+				NativeType{typ: TypeBigInt},
+			},
+		},
+		stateFunc: "avgstate",
+		finalFunc: "avgfinal",
+	}
+
+	// In this case cassandra is returning a blob
+	if flagCassVersion.Before(3, 0, 0) {
+		expectedAggregrate.InitCond = string([]byte{0, 0, 0, 4, 0, 0, 0, 0, 0, 0, 0, 8, 0, 0, 0, 0, 0, 0, 0, 0})
+	}
+
+	if !reflect.DeepEqual(aggregate, expectedAggregrate) {
+		t.Fatalf("aggregate is %+v, but expected %+v", aggregate, expectedAggregrate)
+	}
+}
+
+func TestFunctionMetadata(t *testing.T) {
+	session := createSession(t)
+	defer session.Close()
+	createFunctions(t, session)
+
+	functions, err := getFunctionsMetadata(session, "gocql_test")
+	if err != nil {
+		t.Fatalf("failed to query function metadata with err: %v", err)
+	}
+	if functions == nil {
+		t.Fatal("failed to query function metadata, nil returned")
+	}
+	if len(functions) != 2 {
+		t.Fatal("expected two functions")
+	}
+	avgState := functions[1]
+	avgFinal := functions[0]
+
+	avgStateBody := "if (val !=null) {state.setInt(0, state.getInt(0)+1); state.setLong(1, state.getLong(1)+val.intValue());}return state;"
+	expectedAvgState := FunctionMetadata{
+		Keyspace: "gocql_test",
+		Name:     "avgstate",
+		ArgumentTypes: []TypeInfo{
+			TupleTypeInfo{
+				NativeType: NativeType{typ: TypeTuple},
+
+				Elems: []TypeInfo{
+					NativeType{typ: TypeInt},
+					NativeType{typ: TypeBigInt},
+				},
+			},
+			NativeType{typ: TypeInt},
+		},
+		ArgumentNames: []string{"state", "val"},
+		ReturnType: TupleTypeInfo{
+			NativeType: NativeType{typ: TypeTuple},
+
+			Elems: []TypeInfo{
+				NativeType{typ: TypeInt},
+				NativeType{typ: TypeBigInt},
+			},
+		},
+		CalledOnNullInput: true,
+		Language:          "java",
+		Body:              avgStateBody,
+	}
+	if !reflect.DeepEqual(avgState, expectedAvgState) {
+		t.Fatalf("function is %+v, but expected %+v", avgState, expectedAvgState)
+	}
+
+	finalStateBody := "double r = 0; if (state.getInt(0) == 0) return null; r = state.getLong(1); r/= state.getInt(0); return Double.valueOf(r);"
+	expectedAvgFinal := FunctionMetadata{
+		Keyspace: "gocql_test",
+		Name:     "avgfinal",
+		ArgumentTypes: []TypeInfo{
+			TupleTypeInfo{
+				NativeType: NativeType{typ: TypeTuple},
+
+				Elems: []TypeInfo{
+					NativeType{typ: TypeInt},
+					NativeType{typ: TypeBigInt},
+				},
+			},
+		},
+		ArgumentNames:     []string{"state"},
+		ReturnType:        NativeType{typ: TypeDouble},
+		CalledOnNullInput: true,
+		Language:          "java",
+		Body:              finalStateBody,
+	}
+	if !reflect.DeepEqual(avgFinal, expectedAvgFinal) {
+		t.Fatalf("function is %+v, but expected %+v", avgFinal, expectedAvgFinal)
+	}
+}
+
 // Integration test of querying and composition the keyspace metadata
 func TestKeyspaceMetadata(t *testing.T) {
 	session := createSession(t)
@@ -2192,6 +2312,7 @@ func TestKeyspaceMetadata(t *testing.T) {
 	if err := createTable(session, "CREATE TABLE gocql_test.test_metadata (first_id int, second_id int, third_id int, PRIMARY KEY (first_id, second_id))"); err != nil {
 		t.Fatalf("failed to create table with error '%v'", err)
 	}
+	createAggregate(t, session)
 
 	if err := session.Query("CREATE INDEX index_metadata ON test_metadata ( third_id )").Exec(); err != nil {
 		t.Fatalf("failed to create index with err: %v", err)
@@ -2246,6 +2367,17 @@ func TestKeyspaceMetadata(t *testing.T) {
 		// TODO(zariel): scan index info from system_schema
 		t.Errorf("Expected column index named 'index_metadata' but was '%s'", thirdColumn.Index.Name)
 	}
+
+	aggregate, found := keyspaceMetadata.Aggregates["average"]
+	if !found {
+		t.Fatal("failed to find the aggreate in metadata")
+	}
+	if aggregate.FinalFunc.Name != "avgfinal" {
+		t.Fatalf("expected final function %s, but got %s", "avgFinal", aggregate.FinalFunc.Name)
+	}
+	if aggregate.StateFunc.Name != "avgstate" {
+		t.Fatalf("expected state function %s, but got %s", "avgstate", aggregate.StateFunc.Name)
+	}
 }
 
 // Integration test of the routing key calculation

diff --git a/common_test.go b/common_test.go
@@ -170,6 +170,39 @@ func createTestSession() *Session {
 	return session
 }
 
+func createFunctions(t *testing.T, session *Session) {
+	if err := session.Query(`
+		CREATE OR REPLACE FUNCTION gocql_test.avgState ( state tuple<int,bigint>, val int )
+		CALLED ON NULL INPUT
+		RETURNS tuple<int,bigint>
+		LANGUAGE java AS
+		$$if (val !=null) {state.setInt(0, state.getInt(0)+1); state.setLong(1, state.getLong(1)+val.intValue());}return state;$$;	`).Exec(); err != nil {
+		t.Fatalf("failed to create function with err: %v", err)
+	}
+	if err := session.Query(`
+		CREATE OR REPLACE FUNCTION gocql_test.avgFinal ( state tuple<int,bigint> )
+		CALLED ON NULL INPUT
+		RETURNS double
+		LANGUAGE java AS
+		$$double r = 0; if (state.getInt(0) == 0) return null; r = state.getLong(1); r/= state.getInt(0); return Double.valueOf(r);$$ 
+	`).Exec(); err != nil {
+		t.Fatalf("failed to create function with err: %v", err)
+	}
+}
+
+func createAggregate(t *testing.T, session *Session) {
+	createFunctions(t, session)
+	if err := session.Query(`
+		CREATE OR REPLACE AGGREGATE gocql_test.average(int)
+		SFUNC avgState
+		STYPE tuple<int,bigint>
+		FINALFUNC avgFinal
+		INITCOND (0,0);
+	`).Exec(); err != nil {
+		t.Fatalf("failed to create aggregate with err: %v", err)
+	}
+}
+
 func staticAddressTranslator(newAddr net.IP, newPort int) AddressTranslator {
 	return AddressTranslatorFunc(func(addr net.IP, port int) (net.IP, int) {
 		return newAddr, newPort

diff --git a/helpers.go b/helpers.go
@@ -191,6 +191,20 @@ func splitCompositeTypes(name string) []string {
 	return parts
 }
 
+func apacheToCassandraType(t string) string {
+	t = strings.Replace(t, apacheCassandraTypePrefix, "", -1)
+	t = strings.Replace(t, "(", "<", -1)
+	t = strings.Replace(t, ")", ">", -1)
+	types := strings.FieldsFunc(t, func(r rune) bool {
+		return r == '<' || r == '>' || r == ','
+	})
+	for _, typ := range types {
+		t = strings.Replace(t, typ, getApacheCassandraType(typ).String(), -1)
+	}
+	// This is done so it exactly matches what Cassandra returns
+	return strings.Replace(t, ",", ", ", -1)
+}
+
 func getApacheCassandraType(class string) Type {
 	switch strings.TrimPrefix(class, apacheCassandraTypePrefix) {
 	case "AsciiType":