utils.go 14 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533
  1. package main
  2. import (
  3. "fmt"
  4. nebula "github.com/vesoft-inc/nebula-go/v3"
  5. "log"
  6. "strings"
  7. )
  8. // 表示企业关系结果
  9. //type AllRelationResult struct {
  10. // RelatedCompanies []string // 有关系的企业列表
  11. // Paths []string // 对应的路径
  12. //}
  13. //
  14. //// 批量获取企业的 VID
  15. //func getVidsByName(session *nebula.Session, names []string) (map[string]string, error) {
  16. // if len(names) == 0 {
  17. // return nil, nil
  18. // }
  19. // conditions := ""
  20. // for i, name := range names {
  21. // if i > 0 {
  22. // conditions += " OR "
  23. // }
  24. // conditions += fmt.Sprintf("Legal.name == \"%s\"", name)
  25. // }
  26. // query := fmt.Sprintf(`
  27. //USE %s;
  28. //LOOKUP ON Legal WHERE %s YIELD id(vertex) AS vid, properties(vertex).name AS name
  29. //`, Table_Space, conditions)
  30. // resp, err := session.Execute(query)
  31. // if err != nil {
  32. // return nil, err
  33. // }
  34. // nameToVid := make(map[string]string)
  35. // for _, row := range resp.GetRows() {
  36. // if len(row.Values) >= 2 {
  37. // if row.Values[0].SVal != nil && row.Values[1].SVal != nil {
  38. // nameToVid[string(row.Values[1].SVal)] = string(row.Values[0].SVal)
  39. // }
  40. // }
  41. // }
  42. // return nameToVid, nil
  43. //}
  44. //
  45. //// 获取 VID 对应的名称
  46. //func getVidName(session *nebula.Session, vid string) (string, error) {
  47. // query := fmt.Sprintf(`
  48. //USE %s;
  49. //FETCH PROP ON Legal "%s" YIELD properties(vertex).name AS name
  50. //`, Table_Space, vid)
  51. // resp, err := session.Execute(query)
  52. // if err != nil {
  53. // return "", err
  54. // }
  55. // names, err := getFirstColumnStrings(resp)
  56. // if err != nil || len(names) == 0 {
  57. // return "", fmt.Errorf("未找到 VID %s 的名称", vid)
  58. // }
  59. // return names[0], nil
  60. //}
  61. //
  62. //// 查找路径
  63. //func findPath(session *nebula.Session, fromVid, toVid string, maxStep int, pathCache map[string][]string) ([]string, error) {
  64. // key := fmt.Sprintf("%s->%s:%d", fromVid, toVid, maxStep)
  65. // if cachedPath, ok := pathCache[key]; ok {
  66. // return cachedPath, nil
  67. // }
  68. // query := fmt.Sprintf(`FIND ALL PATH FROM "%s" TO "%s" OVER Invest UPTO %d STEPS YIELD path as p`, fromVid, toVid, maxStep)
  69. // resp, err := session.Execute(query)
  70. // if err != nil {
  71. // return nil, err
  72. // }
  73. // path, err := getFirstColumnStrings(resp)
  74. // if err != nil {
  75. // return nil, err
  76. // }
  77. // pathCache[key] = path
  78. // return path, nil
  79. //}
  80. //
  81. //// 检查共同祖先
  82. //func checkCommonAncestor(session *nebula.Session, aVid, bVid string, deep int, pathCache map[string][]string) (bool, []string, string) {
  83. // key := fmt.Sprintf("%s&%s:%d", aVid, bVid, deep)
  84. // if cachedPath, ok := pathCache[key]; ok {
  85. // if len(cachedPath) > 0 {
  86. // return true, cachedPath, cachedPath[1]
  87. // }
  88. // return false, nil, ""
  89. // }
  90. // query := fmt.Sprintf(`
  91. // (
  92. // GO 1 TO %d STEPS FROM "%s" OVER Invest REVERSELY YIELD dst(edge) AS ancestor
  93. // )
  94. // INTERSECT
  95. // (
  96. // GO 1 TO %d STEPS FROM "%s" OVER Invest REVERSELY YIELD dst(edge) AS ancestor
  97. // );
  98. // `, deep, aVid, deep, bVid)
  99. // resp, err := session.Execute(query)
  100. // if err != nil {
  101. // return false, nil, ""
  102. // }
  103. // ancestors, err := getFirstColumnStrings(resp)
  104. // if err != nil || len(ancestors) == 0 {
  105. // pathCache[key] = nil
  106. // return false, nil, ""
  107. // }
  108. // pathA, _ := findPath(session, aVid, ancestors[0], deep, pathCache)
  109. // pathB, _ := findPath(session, bVid, ancestors[0], deep, pathCache)
  110. // var path []string
  111. // if len(pathB) > 1 {
  112. // path = append(pathA, pathB[1:]...)
  113. // } else {
  114. // path = append(pathA, pathB...)
  115. // }
  116. // pathCache[key] = path
  117. // return true, path, ancestors[0]
  118. //}
  119. //
  120. //// 将 VID 路径转换为名称路径
  121. //func convertVidPathToNamePath(session *nebula.Session, vidPath []string) (string, error) {
  122. // namePath := ""
  123. // for i, vid := range vidPath {
  124. // name, err := getVidName(session, vid)
  125. // if err != nil {
  126. // return "", err
  127. // }
  128. // if i > 0 {
  129. // namePath += "->"
  130. // }
  131. // namePath += name
  132. // }
  133. // return namePath, nil
  134. //}
  135. //
  136. //// 检查企业关系
  137. //func CheckLegalRelations(session *nebula.Session, names []string, deep int) (AllRelationResult, error) {
  138. // result := AllRelationResult{}
  139. // checked := make(map[string]bool)
  140. // nameToVid, err := getVidsByName(session, names)
  141. // if err != nil {
  142. // return result, err
  143. // }
  144. // pathCache := make(map[string][]string)
  145. // relatedCompaniesSet := make(map[string]bool)
  146. // var paths []string
  147. //
  148. // for i := 0; i < len(names); i++ {
  149. // for j := i + 1; j < len(names); j++ {
  150. // a, b := names[i], names[j]
  151. // vidA, okA := nameToVid[a]
  152. // vidB, okB := nameToVid[b]
  153. // if !okA || !okB {
  154. // continue
  155. // }
  156. //
  157. // key := vidA + "|" + vidB
  158. // if checked[key] {
  159. // continue
  160. // }
  161. // checked[key] = true
  162. //
  163. // // 1. a -> b
  164. // pathAB, err := findPath(session, vidA, vidB, deep, pathCache)
  165. // if err != nil {
  166. // log.Printf("查找 %s 到 %s 的路径失败: %v", a, b, err)
  167. // continue
  168. // }
  169. // if len(pathAB) > 0 {
  170. // pathStr, err := convertVidPathToNamePath(session, pathAB)
  171. // if err != nil {
  172. // log.Printf("转换 %s 到 %s 的路径失败: %v", a, b, err)
  173. // continue
  174. // }
  175. // paths = append(paths, pathStr)
  176. // for _, vid := range pathAB {
  177. // name, err := getVidName(session, vid)
  178. // if err != nil {
  179. // log.Printf("获取 VID %s 对应的名称失败: %v", vid, err)
  180. // continue
  181. // }
  182. // relatedCompaniesSet[name] = true
  183. // }
  184. // continue
  185. // }
  186. //
  187. // // 2. b -> a
  188. // pathBA, err := findPath(session, vidB, vidA, deep, pathCache)
  189. // if err != nil {
  190. // log.Printf("查找 %s 到 %s 的路径失败: %v", b, a, err)
  191. // continue
  192. // }
  193. // if len(pathBA) > 0 {
  194. // pathStr, err := convertVidPathToNamePath(session, pathBA)
  195. // if err != nil {
  196. // log.Printf("转换 %s 到 %s 的路径失败: %v", b, a, err)
  197. // continue
  198. // }
  199. // paths = append(paths, pathStr)
  200. // for _, vid := range pathBA {
  201. // name, err := getVidName(session, vid)
  202. // if err != nil {
  203. // log.Printf("获取 VID %s 对应的名称失败: %v", vid, err)
  204. // continue
  205. // }
  206. // relatedCompaniesSet[name] = true
  207. // }
  208. // continue
  209. // }
  210. //
  211. // // 3. common ancestor
  212. // common, path, _ := checkCommonAncestor(session, vidA, vidB, deep, pathCache)
  213. // if common {
  214. // pathStr, err := convertVidPathToNamePath(session, path)
  215. // if err != nil {
  216. // log.Printf("转换 %s 和 %s 到共同祖先的路径失败: %v", a, b, err)
  217. // continue
  218. // }
  219. // paths = append(paths, pathStr)
  220. // for _, vid := range path {
  221. // name, err := getVidName(session, vid)
  222. // if err != nil {
  223. // log.Printf("获取 VID %s 对应的名称失败: %v", vid, err)
  224. // continue
  225. // }
  226. // relatedCompaniesSet[name] = true
  227. // }
  228. // }
  229. // }
  230. // }
  231. //
  232. // for company := range relatedCompaniesSet {
  233. // result.RelatedCompanies = append(result.RelatedCompanies, company)
  234. // }
  235. // result.Paths = paths
  236. //
  237. // return result, nil
  238. //}
  239. //
  240. //// getFirstColumnStrings 适配 nebula-go v3 取出字符串类型列
  241. //func getFirstColumnStrings(resp *nebula.ResultSet) ([]string, error) {
  242. // if resp == nil {
  243. // return nil, fmt.Errorf("result set is nil")
  244. // }
  245. //
  246. // var values []string
  247. // for _, row := range resp.GetRows() {
  248. // if len(row.Values) == 0 {
  249. // continue
  250. // }
  251. // val := row.Values[0]
  252. // switch {
  253. // case val.SVal != nil:
  254. // values = append(values, string(val.SVal))
  255. // case val.IVal != nil:
  256. // values = append(values, fmt.Sprintf("%d", *val.IVal))
  257. // case val.BVal != nil:
  258. // values = append(values, fmt.Sprintf("%v", *val.BVal))
  259. // default:
  260. // log.Printf("未知类型值: %+v", val)
  261. // }
  262. // }
  263. // return values, nil
  264. //}
  265. func CheckLegalRelations(session *nebula.Session, names []string, deep int) ([]RelationResult, error) {
  266. results := []RelationResult{}
  267. checked := make(map[string]bool)
  268. nameToVid, err := getAllVids(session, names)
  269. if err != nil {
  270. return nil, err
  271. }
  272. vidToName := reverseMap(nameToVid)
  273. for i := 0; i < len(names); i++ {
  274. for j := i + 1; j < len(names); j++ {
  275. a, b := names[i], names[j]
  276. vidA, okA := nameToVid[a]
  277. vidB, okB := nameToVid[b]
  278. if !okA || !okB {
  279. continue
  280. }
  281. key := vidA + "|" + vidB
  282. if checked[key] {
  283. continue
  284. }
  285. checked[key] = true
  286. // 1. a -> b
  287. pathAB, err := findPath(session, vidA, vidB, deep)
  288. if err != nil {
  289. return nil, err
  290. }
  291. if len(pathAB) > 0 {
  292. readablePath := convertPathToNames(pathAB, vidToName)
  293. results = append(results, RelationResult{A: a, B: b, RelationType: "direct_or_indirect", Path: readablePath})
  294. continue
  295. }
  296. // 2. b -> a
  297. pathBA, err := findPath(session, vidB, vidA, deep)
  298. if err != nil {
  299. return nil, err
  300. }
  301. if len(pathBA) > 0 {
  302. readablePath := convertPathToNames(pathBA, vidToName)
  303. results = append(results, RelationResult{A: b, B: a, RelationType: "direct_or_indirect", Path: readablePath})
  304. continue
  305. }
  306. // 3. common ancestor
  307. common, ancestorVid, err := checkCommonAncestor(session, vidA, vidB, deep)
  308. if err != nil {
  309. return nil, err
  310. }
  311. if common {
  312. ancestorName := getAncestorName(session, ancestorVid, vidToName)
  313. aName := vidToName[vidA]
  314. bName := vidToName[vidB]
  315. if ancestorName != "" && aName != "" && bName != "" {
  316. readablePath := []string{
  317. fmt.Sprintf("%s -> %s", aName, ancestorName),
  318. fmt.Sprintf("%s -> %s", bName, ancestorName),
  319. }
  320. results = append(results, RelationResult{A: a, B: b, RelationType: "common_ancestor", Path: readablePath})
  321. }
  322. }
  323. }
  324. }
  325. return results, nil
  326. }
  327. func getAllVids(session *nebula.Session, names []string) (map[string]string, error) {
  328. nameToVid := make(map[string]string)
  329. for _, name := range names {
  330. vid, err := getVidByName(session, name)
  331. if err != nil {
  332. log.Printf("获取 %s 的 VID 失败: %v", name, err)
  333. continue
  334. }
  335. nameToVid[name] = vid
  336. }
  337. return nameToVid, nil
  338. }
  339. func checkCommonAncestor(session *nebula.Session, aVid, bVid string, deep int) (bool, string, error) {
  340. query := fmt.Sprintf(`
  341. (
  342. GO 1 TO %d STEPS FROM "%s" OVER Invest REVERSELY YIELD dst(edge) AS ancestor
  343. )
  344. INTERSECT
  345. (
  346. GO 1 TO %d STEPS FROM "%s" OVER Invest REVERSELY YIELD dst(edge) AS ancestor
  347. );
  348. `, deep, aVid, deep, bVid)
  349. resp, err := session.Execute(query)
  350. if err != nil {
  351. return false, "", err
  352. }
  353. ancestors, err := getFirstColumnStrings(resp)
  354. if err != nil || len(ancestors) == 0 {
  355. return false, "", nil
  356. }
  357. return true, ancestors[0], nil
  358. }
  359. func findPath(session *nebula.Session, fromVid, toVid string, maxStep int) ([]string, error) {
  360. query := fmt.Sprintf(`FIND ALL PATH FROM "%s" TO "%s" OVER Invest UPTO %d STEPS YIELD path as p`, fromVid, toVid, maxStep)
  361. resp, err := session.Execute(query)
  362. if err != nil {
  363. return nil, err
  364. }
  365. return getFirstColumnStrings(resp)
  366. }
  367. func getVidByName(session *nebula.Session, name string) (string, error) {
  368. query := fmt.Sprintf(`
  369. USE `+Table_Space+`;
  370. LOOKUP ON Legal WHERE Legal.name == "%s" YIELD id(vertex)`, name)
  371. resp, err := session.Execute(query)
  372. if err != nil {
  373. return "", err
  374. }
  375. values, err := getFirstColumnStrings(resp)
  376. if err != nil || len(values) == 0 {
  377. return "", fmt.Errorf("未找到公司: %s", name)
  378. }
  379. return values[0], nil
  380. }
  381. type RelationResult struct {
  382. A, B string // 公司名
  383. RelationType string // 关系类型:"direct_or_indirect", "common_ancestor"
  384. Path []string // 路径中的公司名称,以更直观的形式展示
  385. }
  386. // getFirstColumnStrings 适配 nebula-go v3 取出字符串类型列
  387. func getFirstColumnStrings(resp *nebula.ResultSet) ([]string, error) {
  388. if resp == nil {
  389. return nil, fmt.Errorf("result set is nil")
  390. }
  391. var values []string
  392. for _, row := range resp.GetRows() {
  393. if len(row.Values) == 0 {
  394. continue
  395. }
  396. val := row.Values[0]
  397. switch {
  398. case val.SVal != nil:
  399. values = append(values, string(val.SVal))
  400. case val.IVal != nil:
  401. values = append(values, fmt.Sprintf("%d", *val.IVal))
  402. case val.BVal != nil:
  403. values = append(values, fmt.Sprintf("%v", *val.BVal))
  404. case val.PVal != nil:
  405. // 处理点类型
  406. src := val.PVal.GetSrc()
  407. log.Println(src)
  408. //if src.GetId != nil {
  409. // values = append(values, string(*src.SVal))
  410. //} else if src.IVal != nil {
  411. // values = append(values, fmt.Sprintf("%d", *src.IVal))
  412. //} else {
  413. // log.Printf("未知的点源 ID 类型: %+v", src)
  414. //}
  415. default:
  416. log.Printf("未知类型值: %+v", val)
  417. }
  418. }
  419. return values, nil
  420. }
  421. func reverseMap(m map[string]string) map[string]string {
  422. result := make(map[string]string)
  423. for k, v := range m {
  424. result[v] = k
  425. }
  426. return result
  427. }
  428. func convertPathToNames(path []string, vidToName map[string]string) []string {
  429. readablePath := make([]string, 0, len(path))
  430. for i := 0; i < len(path)-1; i++ {
  431. fromName, okFrom := vidToName[path[i]]
  432. toName, okTo := vidToName[path[i+1]]
  433. if okFrom && okTo && fromName != toName {
  434. readablePath = append(readablePath, fmt.Sprintf("%s -> %s", fromName, toName))
  435. }
  436. }
  437. return readablePath
  438. }
  439. func getAncestorName(session *nebula.Session, ancestorVid string, vidToName map[string]string) string {
  440. if name, ok := vidToName[ancestorVid]; ok {
  441. return name
  442. }
  443. query := fmt.Sprintf(`
  444. USE `+Table_Space+`;
  445. FETCH PROP ON Legal "%s" YIELD Legal.name;
  446. `, ancestorVid)
  447. resp, err := session.Execute(query)
  448. if err != nil {
  449. log.Printf("获取祖先公司名称失败: %v", err)
  450. return ""
  451. }
  452. names, err := getFirstColumnStrings(resp)
  453. if err != nil || len(names) == 0 {
  454. return ""
  455. }
  456. return names[0]
  457. }
  458. //2222222222222222//
  459. // CheckLegalRelationships 检查给定企业之间是否存在关系
  460. func CheckLegalRelationships(session *nebula.Session, names []string, deep, stype int) (bool, string, error) {
  461. if len(names) < 2 {
  462. return false, "", fmt.Errorf("企业数量不足,至少需要两个")
  463. }
  464. found := false // 是否找到任意关系
  465. for i := 0; i < len(names); i++ {
  466. start := names[i]
  467. // 构造剩下的企业列表作为目标
  468. targets := []string{}
  469. for j := 0; j < len(names); j++ {
  470. if i != j {
  471. targets = append(targets, fmt.Sprintf(`"%s"`, names[j]))
  472. }
  473. }
  474. targetList := strings.Join(targets, ", ")
  475. // 构造 nGQL 查询
  476. query := fmt.Sprintf(`
  477. USE %s;
  478. MATCH p=(a:Legal{name:"%s"})-[*1..%d]-(b:Legal)
  479. WHERE b.Legal.name IN [%s]
  480. RETURN p LIMIT 1
  481. `, Table_Space, start, deep, targetList)
  482. // 执行查询
  483. resp, err := session.Execute(query)
  484. if err != nil {
  485. return false, "", fmt.Errorf("查询失败: %w", err)
  486. }
  487. if !resp.IsSucceed() {
  488. return false, "", fmt.Errorf("查询执行失败: %s", resp.GetErrorMsg())
  489. }
  490. if resp.GetRowSize() > 0 {
  491. found = true
  492. // stype == 0:只要找到一个关系就返回
  493. if stype == 0 {
  494. return true, "企业之间存在关系", nil
  495. }
  496. }
  497. }
  498. // 所有组合查完,判断是否找到关系
  499. if found {
  500. return true, "企业之间存在关系", nil
  501. }
  502. return false, "企业之间无关联", nil
  503. }