utils.go 1.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081
  1. package dolphin
  2. import (
  3. pcast "github.com/pingcap/parser/ast"
  4. "github.com/kyleconroy/sqlc/internal/sql/ast"
  5. )
  6. type nodeSearch struct {
  7. list []pcast.Node
  8. check func(pcast.Node) bool
  9. }
  10. func (s *nodeSearch) Enter(n pcast.Node) (pcast.Node, bool) {
  11. if s.check(n) {
  12. s.list = append(s.list, n)
  13. }
  14. return n, false // skipChildren
  15. }
  16. func (s *nodeSearch) Leave(n pcast.Node) (pcast.Node, bool) {
  17. return n, true // ok
  18. }
  19. func collect(root pcast.Node, f func(pcast.Node) bool) []pcast.Node {
  20. if root == nil {
  21. return nil
  22. }
  23. ns := &nodeSearch{check: f}
  24. root.Accept(ns)
  25. return ns.list
  26. }
  27. type nodeVisit struct {
  28. fn func(pcast.Node)
  29. }
  30. func (s *nodeVisit) Enter(n pcast.Node) (pcast.Node, bool) {
  31. s.fn(n)
  32. return n, false // skipChildren
  33. }
  34. func (s *nodeVisit) Leave(n pcast.Node) (pcast.Node, bool) {
  35. return n, true // ok
  36. }
  37. func visit(root pcast.Node, f func(pcast.Node)) {
  38. if root == nil {
  39. return
  40. }
  41. ns := &nodeVisit{fn: f}
  42. root.Accept(ns)
  43. }
  44. // Maybe not useful?
  45. func text(nodes []pcast.Node) []string {
  46. str := make([]string, len(nodes))
  47. for i := range nodes {
  48. if nodes[i] == nil {
  49. continue
  50. }
  51. str[i] = nodes[i].Text()
  52. }
  53. return str
  54. }
  55. func parseTableName(n *pcast.TableName) *ast.TableName {
  56. return &ast.TableName{
  57. Schema: n.Schema.String(),
  58. Name: n.Name.String(),
  59. }
  60. }
  61. func isNotNull(n *pcast.ColumnDef) bool {
  62. for i := range n.Options {
  63. if n.Options[i].Tp == pcast.ColumnOptionNotNull {
  64. return true
  65. }
  66. }
  67. return false
  68. }