302 lines
7.0 KiB
Go
302 lines
7.0 KiB
Go
|
package services
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"fmt"
|
||
|
|
|
||
|
|
"github.com/neo4j/neo4j-go-driver/v5/neo4j"
|
||
|
|
|
||
|
|
"knowledge-graph-backend/models"
|
||
|
|
)
|
||
|
|
|
||
|
|
type Neo4jService struct {
|
||
|
|
driver neo4j.DriverWithContext
|
||
|
|
}
|
||
|
|
|
||
|
|
func NewNeo4jService(driver neo4j.DriverWithContext) *Neo4jService {
|
||
|
|
return &Neo4jService{driver: driver}
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *Neo4jService) GetGraphData() models.GraphData {
|
||
|
|
ctx := context.Background()
|
||
|
|
var nodes []models.Node
|
||
|
|
var edges []models.Edge
|
||
|
|
|
||
|
|
nodeResult, err := neo4j.ExecuteQuery(ctx, s.driver,
|
||
|
|
"MATCH (n) RETURN n",
|
||
|
|
map[string]any{},
|
||
|
|
neo4j.EagerResultTransformer,
|
||
|
|
)
|
||
|
|
if err != nil {
|
||
|
|
fmt.Printf("Error querying nodes: %v\n", err)
|
||
|
|
return models.GraphData{}
|
||
|
|
}
|
||
|
|
for _, record := range nodeResult.Records {
|
||
|
|
if v, ok := record.Get("n"); ok {
|
||
|
|
if n, ok := v.(neo4j.Node); ok {
|
||
|
|
nodes = append(nodes, neo4jNodeToModel(n))
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
edgeResult, err := neo4j.ExecuteQuery(ctx, s.driver,
|
||
|
|
"MATCH (a)-[r]->(b) RETURN a.id AS source, b.id AS target, r",
|
||
|
|
map[string]any{},
|
||
|
|
neo4j.EagerResultTransformer,
|
||
|
|
)
|
||
|
|
if err != nil {
|
||
|
|
fmt.Printf("Error querying edges: %v\n", err)
|
||
|
|
return models.GraphData{Nodes: nodes}
|
||
|
|
}
|
||
|
|
for _, record := range edgeResult.Records {
|
||
|
|
source, _ := record.Get("source")
|
||
|
|
target, _ := record.Get("target")
|
||
|
|
v, _ := record.Get("r")
|
||
|
|
if rel, ok := v.(neo4j.Relationship); ok {
|
||
|
|
edge := neo4jRelToModel(rel)
|
||
|
|
if s, ok := source.(string); ok {
|
||
|
|
edge.Source = s
|
||
|
|
}
|
||
|
|
if t, ok := target.(string); ok {
|
||
|
|
edge.Target = t
|
||
|
|
}
|
||
|
|
edges = append(edges, edge)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
return models.GraphData{Nodes: nodes, Edges: edges}
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *Neo4jService) GetNodeByID(id string) (models.Node, bool) {
|
||
|
|
ctx := context.Background()
|
||
|
|
result, err := neo4j.ExecuteQuery(ctx, s.driver,
|
||
|
|
"MATCH (n {id: $id}) RETURN n",
|
||
|
|
map[string]any{"id": id},
|
||
|
|
neo4j.EagerResultTransformer,
|
||
|
|
)
|
||
|
|
if err != nil || len(result.Records) == 0 {
|
||
|
|
return models.Node{}, false
|
||
|
|
}
|
||
|
|
v, _ := result.Records[0].Get("n")
|
||
|
|
if n, ok := v.(neo4j.Node); ok {
|
||
|
|
return neo4jNodeToModel(n), true
|
||
|
|
}
|
||
|
|
return models.Node{}, false
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *Neo4jService) SearchNodes(query string) []models.Node {
|
||
|
|
if query == "" {
|
||
|
|
return []models.Node{}
|
||
|
|
}
|
||
|
|
ctx := context.Background()
|
||
|
|
result, err := neo4j.ExecuteQuery(ctx, s.driver,
|
||
|
|
`MATCH (n)
|
||
|
|
WHERE toLower(n.label) CONTAINS toLower($query)
|
||
|
|
OR toLower(n.id) CONTAINS toLower($query)
|
||
|
|
OR (n.type IS NOT NULL AND toLower(n.type) CONTAINS toLower($query))
|
||
|
|
RETURN n`,
|
||
|
|
map[string]any{"query": query},
|
||
|
|
neo4j.EagerResultTransformer,
|
||
|
|
)
|
||
|
|
if err != nil {
|
||
|
|
return []models.Node{}
|
||
|
|
}
|
||
|
|
var nodes []models.Node
|
||
|
|
for _, record := range result.Records {
|
||
|
|
v, _ := record.Get("n")
|
||
|
|
if n, ok := v.(neo4j.Node); ok {
|
||
|
|
nodes = append(nodes, neo4jNodeToModel(n))
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return nodes
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *Neo4jService) GetNeighbors(nodeID string) (models.NeighborResponse, bool) {
|
||
|
|
ctx := context.Background()
|
||
|
|
|
||
|
|
_, exists := s.GetNodeByID(nodeID)
|
||
|
|
if !exists {
|
||
|
|
return models.NeighborResponse{}, false
|
||
|
|
}
|
||
|
|
|
||
|
|
result, err := neo4j.ExecuteQuery(ctx, s.driver,
|
||
|
|
`MATCH ({id: $id})-[r]-(m)
|
||
|
|
RETURN m, r, startNode(r).id AS source, endNode(r).id AS target`,
|
||
|
|
map[string]any{"id": nodeID},
|
||
|
|
neo4j.EagerResultTransformer,
|
||
|
|
)
|
||
|
|
if err != nil {
|
||
|
|
return models.NeighborResponse{Nodes: []models.Node{}, Edges: []models.Edge{}}, true
|
||
|
|
}
|
||
|
|
|
||
|
|
neighborMap := make(map[string]models.Node)
|
||
|
|
edgeMap := make(map[string]models.Edge)
|
||
|
|
|
||
|
|
for _, record := range result.Records {
|
||
|
|
if v, ok := record.Get("m"); ok {
|
||
|
|
if n, ok := v.(neo4j.Node); ok {
|
||
|
|
modelNode := neo4jNodeToModel(n)
|
||
|
|
neighborMap[modelNode.ID] = modelNode
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
if v, ok := record.Get("r"); ok {
|
||
|
|
if rel, ok := v.(neo4j.Relationship); ok {
|
||
|
|
edge := neo4jRelToModel(rel)
|
||
|
|
if source, ok := record.Get("source"); ok {
|
||
|
|
if s, ok := source.(string); ok {
|
||
|
|
edge.Source = s
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if target, ok := record.Get("target"); ok {
|
||
|
|
if t, ok := target.(string); ok {
|
||
|
|
edge.Target = t
|
||
|
|
}
|
||
|
|
}
|
||
|
|
edgeMap[edge.ID] = edge
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
var neighborNodes []models.Node
|
||
|
|
for _, n := range neighborMap {
|
||
|
|
neighborNodes = append(neighborNodes, n)
|
||
|
|
}
|
||
|
|
|
||
|
|
var relatedEdges []models.Edge
|
||
|
|
for _, e := range edgeMap {
|
||
|
|
relatedEdges = append(relatedEdges, e)
|
||
|
|
}
|
||
|
|
|
||
|
|
return models.NeighborResponse{
|
||
|
|
Nodes: neighborNodes,
|
||
|
|
Edges: relatedEdges,
|
||
|
|
}, true
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *Neo4jService) GetStats() map[string]int {
|
||
|
|
ctx := context.Background()
|
||
|
|
result, err := neo4j.ExecuteQuery(ctx, s.driver,
|
||
|
|
`MATCH (n)
|
||
|
|
RETURN count(n) AS totalNodes,
|
||
|
|
sum(CASE WHEN n.type = '概念' THEN 1 ELSE 0 END) AS conceptNodes,
|
||
|
|
sum(CASE WHEN n.type = '工具' THEN 1 ELSE 0 END) AS toolNodes,
|
||
|
|
sum(CASE WHEN n.type = '应用' THEN 1 ELSE 0 END) AS applicationNodes`,
|
||
|
|
map[string]any{},
|
||
|
|
neo4j.EagerResultTransformer,
|
||
|
|
)
|
||
|
|
if err != nil || len(result.Records) == 0 {
|
||
|
|
return map[string]int{}
|
||
|
|
}
|
||
|
|
|
||
|
|
record := result.Records[0]
|
||
|
|
totalNodes := getInt(record, "totalNodes")
|
||
|
|
conceptNodes := getInt(record, "conceptNodes")
|
||
|
|
toolNodes := getInt(record, "toolNodes")
|
||
|
|
applicationNodes := getInt(record, "applicationNodes")
|
||
|
|
|
||
|
|
totalEdges := 0
|
||
|
|
edgeResult, err := neo4j.ExecuteQuery(ctx, s.driver,
|
||
|
|
"MATCH ()-[r]->() RETURN count(r) AS totalEdges",
|
||
|
|
map[string]any{},
|
||
|
|
neo4j.EagerResultTransformer,
|
||
|
|
)
|
||
|
|
if err == nil && len(edgeResult.Records) > 0 {
|
||
|
|
totalEdges = getInt(edgeResult.Records[0], "totalEdges")
|
||
|
|
}
|
||
|
|
|
||
|
|
return map[string]int{
|
||
|
|
"totalNodes": totalNodes,
|
||
|
|
"totalEdges": totalEdges,
|
||
|
|
"conceptNodes": conceptNodes,
|
||
|
|
"toolNodes": toolNodes,
|
||
|
|
"applicationNodes": applicationNodes,
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func neo4jNodeToModel(n neo4j.Node) models.Node {
|
||
|
|
props := n.Props
|
||
|
|
node := models.Node{
|
||
|
|
ID: getStr(props, "id"),
|
||
|
|
Label: getStr(props, "label"),
|
||
|
|
Type: getStr(props, "type"),
|
||
|
|
Properties: make(map[string]interface{}),
|
||
|
|
}
|
||
|
|
if v, ok := props["x"]; ok && v != nil {
|
||
|
|
node.X = getFloat64(v)
|
||
|
|
}
|
||
|
|
if v, ok := props["y"]; ok && v != nil {
|
||
|
|
node.Y = getFloat64(v)
|
||
|
|
}
|
||
|
|
if v, ok := props["style"]; ok && v != nil {
|
||
|
|
if s, ok := v.(map[string]interface{}); ok {
|
||
|
|
node.Style = s
|
||
|
|
}
|
||
|
|
}
|
||
|
|
for k, v := range props {
|
||
|
|
switch k {
|
||
|
|
case "id", "label", "type", "x", "y", "style":
|
||
|
|
default:
|
||
|
|
node.Properties[k] = v
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return node
|
||
|
|
}
|
||
|
|
|
||
|
|
func neo4jRelToModel(r neo4j.Relationship) models.Edge {
|
||
|
|
props := r.Props
|
||
|
|
edge := models.Edge{
|
||
|
|
ID: getStr(props, "id"),
|
||
|
|
Label: getStr(props, "label"),
|
||
|
|
Type: r.Type,
|
||
|
|
Properties: make(map[string]interface{}),
|
||
|
|
}
|
||
|
|
if v, ok := props["style"]; ok && v != nil {
|
||
|
|
if s, ok := v.(map[string]interface{}); ok {
|
||
|
|
edge.Style = s
|
||
|
|
}
|
||
|
|
}
|
||
|
|
for k, v := range props {
|
||
|
|
switch k {
|
||
|
|
case "id", "label", "style":
|
||
|
|
default:
|
||
|
|
edge.Properties[k] = v
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return edge
|
||
|
|
}
|
||
|
|
|
||
|
|
func getStr(props map[string]any, key string) string {
|
||
|
|
if v, ok := props[key]; ok && v != nil {
|
||
|
|
if s, ok := v.(string); ok {
|
||
|
|
return s
|
||
|
|
}
|
||
|
|
}
|
||
|
|
return ""
|
||
|
|
}
|
||
|
|
|
||
|
|
func getFloat64(v any) float64 {
|
||
|
|
switch val := v.(type) {
|
||
|
|
case float64:
|
||
|
|
return val
|
||
|
|
case int64:
|
||
|
|
return float64(val)
|
||
|
|
case int:
|
||
|
|
return float64(val)
|
||
|
|
}
|
||
|
|
return 0
|
||
|
|
}
|
||
|
|
|
||
|
|
func getInt(record *neo4j.Record, key string) int {
|
||
|
|
v, _ := record.Get(key)
|
||
|
|
switch val := v.(type) {
|
||
|
|
case int64:
|
||
|
|
return int(val)
|
||
|
|
case int:
|
||
|
|
return val
|
||
|
|
case float64:
|
||
|
|
return int(val)
|
||
|
|
}
|
||
|
|
return 0
|
||
|
|
}
|