feat: allow specifying the graphdb database name (#235)

Author: rcrisanti

pkg/externalfunctions/dataextraction.go

@@ -881,39 +881,51 @@ func StoreElementsInVectorDatabase(elements []sharedtypes.CodeGenerationElement,
 	}
 }
 
+func graphDbNameOrDefault(dbname string) string {
+	if dbname == "" {
+		logging.Log.Debugf(&logging.ContextMap{}, "No database name provided for aali-graphdb, falling back to default 'aali'")
+		return "aali"
+	} else {
+		return dbname
+	}
+}
+
 // StoreElementsInGraphDatabase stores elements in the graph database.
 //
 // Tags:
 //   - @displayName: Store Elements in Graph Database
 //
 // Parameters:
+//   - dbname: the name of the graphdb to target. If not provided, defaults to "aali".
 //   - elements: code generation elements.
-func StoreElementsInGraphDatabase(elements []sharedtypes.CodeGenerationElement) {
+func StoreElementsInGraphDatabase(dbname string, elements []sharedtypes.CodeGenerationElement) {
 	ctx := &logging.ContextMap{}
 
+	dbname = graphDbNameOrDefault(dbname)
+
 	// Initialize the graph database.
-	err := graphdb.Initialize(config.GlobalConfig.GRAPHDB_ADDRESS)
+	err := graphdb.Initialize(config.GlobalConfig.GRAPHDB_ADDRESS, dbname)
 	if err != nil {
-		errMsg := fmt.Sprintf("error initializing graphdb: %v", err)
+		errMsg := fmt.Sprintf("error initializing graphdb %q: %v", dbname, err)
 		logging.Log.Error(ctx, errMsg)
 		panic(errMsg)
 	}
 
-	err = graphdb.GraphDbDriver.CreateSchema()
+	err = graphdb.GraphDbDriver.CreateSchema(dbname)
 	if err != nil {
 		logPanic(ctx, "error error creating aali schema: %v", err)
 	}
 
 	// Add the elements to the graph database.
-	err = graphdb.GraphDbDriver.AddCodeGenerationElementNodes(elements)
+	err = graphdb.GraphDbDriver.AddCodeGenerationElementNodes(dbname, elements)
 	if err != nil {
 		errMsg := fmt.Sprintf("error adding code gen element nodes to graphdb: %v", err)
 		logging.Log.Error(ctx, errMsg)
 		panic(errMsg)
 	}
 
 	// Add the dependencies to the graph database.
-	err = graphdb.GraphDbDriver.CreateCodeGenerationRelationships(elements)
+	err = graphdb.GraphDbDriver.CreateCodeGenerationRelationships(dbname, elements)
 	if err != nil {
 		errMsg := fmt.Sprintf("error adding code gen relationships to graphdb: %v", err)
 		logging.Log.Error(ctx, errMsg)
@@ -1317,33 +1329,36 @@ func StoreExamplesInVectorDatabase(examples []sharedtypes.CodeGenerationExample,
 //   - @displayName: Store Examples in Graph Database
 //
 // Parameters:
+//   - dbname: the name of the graphdb to target. If not provided, defaults to "aali".
 //   - examples: code generation examples.
-func StoreExamplesInGraphDatabase(examples []sharedtypes.CodeGenerationExample) {
+func StoreExamplesInGraphDatabase(dbname string, examples []sharedtypes.CodeGenerationExample) {
 	ctx := &logging.ContextMap{}
 
+	dbname = graphDbNameOrDefault(dbname)
+
 	// Initialize the graph database.
-	err := graphdb.Initialize(config.GlobalConfig.GRAPHDB_ADDRESS)
+	err := graphdb.Initialize(config.GlobalConfig.GRAPHDB_ADDRESS, dbname)
 	if err != nil {
-		errMsg := fmt.Sprintf("error initializing graphdb: %v", err)
+		errMsg := fmt.Sprintf("error initializing graphdb %q: %v", dbname, err)
 		logging.Log.Error(ctx, errMsg)
 		panic(errMsg)
 	}
 
-	err = graphdb.GraphDbDriver.CreateSchema()
+	err = graphdb.GraphDbDriver.CreateSchema(dbname)
 	if err != nil {
 		logPanic(ctx, "error error creating aali schema: %v", err)
 	}
 
 	// Add the elements to the graph database.
-	err = graphdb.GraphDbDriver.AddCodeGenerationExampleNodes(examples)
+	err = graphdb.GraphDbDriver.AddCodeGenerationExampleNodes(dbname, examples)
 	if err != nil {
 		errMsg := fmt.Sprintf("error adding code gen example nodes to graphdb: %v", err)
 		logging.Log.Error(ctx, errMsg)
 		panic(errMsg)
 	}
 
 	// Add the dependencies to the graph database.
-	err = graphdb.GraphDbDriver.CreateCodeGenerationExampleRelationships(examples)
+	err = graphdb.GraphDbDriver.CreateCodeGenerationExampleRelationships(dbname, examples)
 	if err != nil {
 		errMsg := fmt.Sprintf("error adding code gen example relationships to graphdb: %v", err)
 		logging.Log.Error(ctx, errMsg)
@@ -1593,34 +1608,36 @@ func StoreUserGuideSectionsInVectorDatabase(sections []sharedtypes.CodeGeneratio
 //   - @displayName: Store User Guide Sections in Graph Database
 //
 // Parameters:
-//   - elements: user guide sections.
-//   - label: label for the sections (UserGuide by default).
-func StoreUserGuideSectionsInGraphDatabase(sections []sharedtypes.CodeGenerationUserGuideSection) {
+//   - dbname: the name of the graphdb to target. If not provided, defaults to "aali".
+//   - sections: user guide sections.
+func StoreUserGuideSectionsInGraphDatabase(dbname string, sections []sharedtypes.CodeGenerationUserGuideSection) {
 	ctx := &logging.ContextMap{}
 
+	dbname = graphDbNameOrDefault(dbname)
+
 	// Initialize the graph database.
-	err := graphdb.Initialize(config.GlobalConfig.GRAPHDB_ADDRESS)
+	err := graphdb.Initialize(config.GlobalConfig.GRAPHDB_ADDRESS, dbname)
 	if err != nil {
 		errMsg := fmt.Sprintf("error initializing graphdb: %v", err)
 		logging.Log.Error(ctx, errMsg)
 		panic(errMsg)
 	}
 
-	err = graphdb.GraphDbDriver.CreateSchema()
+	err = graphdb.GraphDbDriver.CreateSchema(dbname)
 	if err != nil {
 		logPanic(ctx, "error error creating aali schema: %v", err)
 	}
 
 	// Add the elements to the graph database.
-	err = graphdb.GraphDbDriver.AddUserGuideSectionNodes(sections)
+	err = graphdb.GraphDbDriver.AddUserGuideSectionNodes(dbname, sections)
 	if err != nil {
 		errMsg := fmt.Sprintf("error adding user guide section nodes to graphdb: %v", err)
 		logging.Log.Error(ctx, errMsg)
 		panic(errMsg)
 	}
 
 	// Add the dependencies to the graph database.
-	err = graphdb.GraphDbDriver.CreateUserGuideSectionRelationships(sections)
+	err = graphdb.GraphDbDriver.CreateUserGuideSectionRelationships(dbname, sections)
 	if err != nil {
 		errMsg := fmt.Sprintf("error adding user guide section relationships to graphdb: %v", err)
 		logging.Log.Error(ctx, errMsg)

pkg/externalfunctions/dataextraction_test.go

@@ -761,7 +761,7 @@ func TestStoreElementsInGraphDatabase(t *testing.T) {
 	logging.InitLogger(&setup.config)
 
 	// do some initial checks
-	const DBNAME = "aali" // this is hardcoded in the graphdb driver for now
+	const DBNAME = "aali"
 	client, err := aali_graphdb.DefaultClient(setup.config.GRAPHDB_ADDRESS)
 	require.NoError(err)
 
@@ -797,7 +797,7 @@ func TestStoreElementsInGraphDatabase(t *testing.T) {
 		EnumValues: nil,
 	}
 
-	StoreElementsInGraphDatabase([]sharedtypes.CodeGenerationElement{element})
+	StoreElementsInGraphDatabase(DBNAME, []sharedtypes.CodeGenerationElement{element})
 
 	// query graphdb to make sure things are as they should be
 
@@ -862,7 +862,7 @@ func TestStoreExamplesInGraphDatabase(t *testing.T) {
 	logging.InitLogger(&setup.config)
 
 	// do some initial checks
-	const DBNAME = "aali" // this is hardcoded in the graphdb driver for now
+	const DBNAME = "aali"
 	client, err := aali_graphdb.DefaultClient(setup.config.GRAPHDB_ADDRESS)
 	require.NoError(err)
 
@@ -881,7 +881,7 @@ func TestStoreExamplesInGraphDatabase(t *testing.T) {
 		},
 	}
 
-	StoreExamplesInGraphDatabase(examples)
+	StoreExamplesInGraphDatabase(DBNAME, examples)
 
 	// query graphdb to make sure things are as they should be
 
@@ -933,7 +933,7 @@ func TestStoreUserGuideSectionsInGraphDatabase(t *testing.T) {
 	logging.InitLogger(&setup.config)
 
 	// do some initial checks
-	const DBNAME = "aali" // this is hardcoded in the graphdb driver for now
+	const DBNAME = "aali"
 	client, err := aali_graphdb.DefaultClient(setup.config.GRAPHDB_ADDRESS)
 	require.NoError(err)
 
@@ -993,7 +993,7 @@ func TestStoreUserGuideSectionsInGraphDatabase(t *testing.T) {
 		},
 	}
 
-	StoreUserGuideSectionsInGraphDatabase(sections)
+	StoreUserGuideSectionsInGraphDatabase(DBNAME, sections)
 
 	// query graphdb to make sure things are as they should be
 	type dbUserGuideSection struct {

pkg/externalfunctions/knowledgedb.go

@@ -193,6 +193,7 @@ func GetListCollections() (collectionsList []string) {
 //   - @displayName: Retrieve Dependencies
 //
 // Parameters:
+//   - dbname: the name of the graphdb to target. If not provided, defaults to "aali".
 //   - relationshipName: the name of the relationship to retrieve dependencies for.
 //   - relationshipDirection: the direction of the relationship to retrieve dependencies for.
 //   - sourceDocumentId: the document ID of the source node.
@@ -202,14 +203,17 @@ func GetListCollections() (collectionsList []string) {
 // Returns:
 //   - dependenciesIds: the list of dependencies
 func RetrieveDependencies(
+	dbname string,
 	relationshipName string,
 	relationshipDirection string,
 	sourceDocumentId string,
 	nodeTypesFilter sharedtypes.DbArrayFilter,
 	maxHopsNumber int) (dependenciesIds []string) {
 	ctx := &logging.ContextMap{}
+	dbname = graphDbNameOrDefault(dbname)
 	dependenciesIds, err := graphdb.GraphDbDriver.RetrieveDependencies(
 		ctx,
+		dbname,
 		relationshipName,
 		relationshipDirection,
 		sourceDocumentId,
@@ -254,18 +258,20 @@ func AddGraphDbParameter(parameters aali_graphdb.ParameterMap, name string, valu
 //   - @displayName: General Graph DB Query
 //
 // Parameters:
+//   - dbname: the name of the graphdb to target. If not provided, defaults to "aali".
 //   - query: the Cypher query to be executed.
 //   - parameters: parameters to pass to the query during execution
 //
 // Returns:
 //   - databaseResponse: the graph db response
-func GeneralGraphDbQuery(query string, parameters aali_graphdb.ParameterMap) []map[string]any {
+func GeneralGraphDbQuery(dbname string, query string, parameters aali_graphdb.ParameterMap) []map[string]any {
 	// Initialize the graph database.
-	err := graphdb.Initialize(config.GlobalConfig.GRAPHDB_ADDRESS)
+	dbname = graphDbNameOrDefault(dbname)
+	err := graphdb.Initialize(config.GlobalConfig.GRAPHDB_ADDRESS, dbname)
 	if err != nil {
 		logPanic(nil, "error initializing graphdb: %v", err)
 	}
-	res, err := graphdb.GraphDbDriver.WriteCypherQuery(query, parameters)
+	res, err := graphdb.GraphDbDriver.WriteCypherQuery(dbname, query, parameters)
 	if err != nil {
 		logPanic(nil, "error executing cypher query: %q", err)
 	}

pkg/externalfunctions/knowledgedb_test.go

@@ -540,6 +540,7 @@ func TestRetrieveDependencies(t *testing.T) {
 
 	require := require.New(t)
 	assert := assert.New(t)
+	const DBNAME = "aali-test"
 
 	setup := setupFlowkitTestContainers(
 		t,
@@ -554,8 +555,8 @@ func TestRetrieveDependencies(t *testing.T) {
 	config.GlobalConfig = &setup.config
 	logging.InitLogger(&setup.config)
 
-	require.NoError(graphdb.Initialize(setup.config.GRAPHDB_ADDRESS))
-	require.NoError(graphdb.GraphDbDriver.CreateSchema())
+	require.NoError(graphdb.Initialize(setup.config.GRAPHDB_ADDRESS, DBNAME))
+	require.NoError(graphdb.GraphDbDriver.CreateSchema(DBNAME))
 	driver := graphdb.GraphDbDriver
 
 	// first add data
@@ -575,19 +576,19 @@ func TestRetrieveDependencies(t *testing.T) {
 		{Name: "4", Level: 0},
 		{Name: "4a", IsFirstChild: true, Parent: "4", Level: 1},
 	}
-	require.NoError(driver.AddUserGuideSectionNodes(data))
-	require.NoError(driver.CreateUserGuideSectionRelationships(data))
+	require.NoError(driver.AddUserGuideSectionNodes(DBNAME, data))
+	require.NoError(driver.CreateUserGuideSectionRelationships(DBNAME, data))
 
 	// now test deps
 	// relationshipName := []string{"NextSibling", "NextParent", "HasFirstChild", "HasChild"}
 	// relationshipDirection := []string{"in", "out", "both"}
-	depIds := RetrieveDependencies("NextParent", "out", "3b", sharedtypes.DbArrayFilter{}, 2)
+	depIds := RetrieveDependencies(DBNAME, "NextParent", "out", "3b", sharedtypes.DbArrayFilter{}, 2)
 	assert.EqualValues([]string{"4"}, depIds)
 
-	depIds = RetrieveDependencies("NextSibling", "in", "3b", sharedtypes.DbArrayFilter{}, 2)
+	depIds = RetrieveDependencies(DBNAME, "NextSibling", "in", "3b", sharedtypes.DbArrayFilter{}, 2)
 	assert.EqualValues([]string{"3a"}, depIds)
 
-	depIds = RetrieveDependencies("HasChild", "both", "3b", sharedtypes.DbArrayFilter{}, 2)
+	depIds = RetrieveDependencies(DBNAME, "HasChild", "both", "3b", sharedtypes.DbArrayFilter{}, 2)
 	expected := []string{"3", "3b1", "3b2"}
 	assert.Len(depIds, len(expected))
 	for _, exp := range expected {
@@ -601,6 +602,7 @@ func TestGeneralGraphDbQuery(t *testing.T) {
 
 	require := require.New(t)
 	assert := assert.New(t)
+	const DBNAME = "test-db"
 
 	setup := setupFlowkitTestContainers(
 		t,
@@ -615,8 +617,8 @@ func TestGeneralGraphDbQuery(t *testing.T) {
 	config.GlobalConfig = &setup.config
 	logging.InitLogger(&setup.config)
 
-	require.NoError(graphdb.Initialize(setup.config.GRAPHDB_ADDRESS))
-	require.NoError(graphdb.GraphDbDriver.CreateSchema())
+	require.NoError(graphdb.Initialize(setup.config.GRAPHDB_ADDRESS, DBNAME))
+	require.NoError(graphdb.GraphDbDriver.CreateSchema(DBNAME))
 	driver := graphdb.GraphDbDriver
 
 	// first add data
@@ -636,11 +638,11 @@ func TestGeneralGraphDbQuery(t *testing.T) {
 		{Name: "4", Level: 0},
 		{Name: "4a", IsFirstChild: true, Parent: "4", Level: 1},
 	}
-	require.NoError(driver.AddUserGuideSectionNodes(data))
-	require.NoError(driver.CreateUserGuideSectionRelationships(data))
+	require.NoError(driver.AddUserGuideSectionNodes(DBNAME, data))
+	require.NoError(driver.CreateUserGuideSectionRelationships(DBNAME, data))
 
 	// now make a query
-	res := GeneralGraphDbQuery("MATCH (n {parent:'1'}) RETURN n.name AS name", nil)
+	res := GeneralGraphDbQuery(DBNAME, "MATCH (n {parent:'1'}) RETURN n.name AS name", nil)
 	expected := []string{"1a", "1b"}
 	require.Len(res, len(expected))
 	for _, r := range res {
@@ -655,6 +657,7 @@ func TestGeneralGraphDbQueryWithParameters(t *testing.T) {
 
 	require := require.New(t)
 	assert := assert.New(t)
+	const DBNAME = "test-db-params"
 
 	setup := setupFlowkitTestContainers(
 		t,
@@ -669,8 +672,8 @@ func TestGeneralGraphDbQueryWithParameters(t *testing.T) {
 	config.GlobalConfig = &setup.config
 	logging.InitLogger(&setup.config)
 
-	require.NoError(graphdb.Initialize(setup.config.GRAPHDB_ADDRESS))
-	require.NoError(graphdb.GraphDbDriver.CreateSchema())
+	require.NoError(graphdb.Initialize(setup.config.GRAPHDB_ADDRESS, DBNAME))
+	require.NoError(graphdb.GraphDbDriver.CreateSchema(DBNAME))
 	driver := graphdb.GraphDbDriver
 
 	// first add data
@@ -690,15 +693,15 @@ func TestGeneralGraphDbQueryWithParameters(t *testing.T) {
 		{Name: "4", Level: 0},
 		{Name: "4a", IsFirstChild: true, Parent: "4", Level: 1},
 	}
-	require.NoError(driver.AddUserGuideSectionNodes(data))
-	require.NoError(driver.CreateUserGuideSectionRelationships(data))
+	require.NoError(driver.AddUserGuideSectionNodes(DBNAME, data))
+	require.NoError(driver.CreateUserGuideSectionRelationships(DBNAME, data))
 
 	// make a parameter map
 	paramMap := aali_graphdb.ParameterMap{}
 	AddGraphDbParameter(paramMap, "parent", "1", "string")
 
 	// now make a query
-	res := GeneralGraphDbQuery("MATCH (n {parent:$parent}) RETURN n.name AS name", paramMap)
+	res := GeneralGraphDbQuery(DBNAME, "MATCH (n {parent:$parent}) RETURN n.name AS name", paramMap)
 	expected := []string{"1a", "1b"}
 	require.Len(res, len(expected))
 	for _, r := range res {