Merge pull request #12 from ColdWaterLW/support-1746

Support 1746

ast/context.go

@@ -1,9 +1,12 @@
 package ast
 
+import "fmt"
+
 type Context struct {
-	QueryType string // select, insert, update, delete
-	Variable  map[string]string
-	Sqls      map[string]*SqlNode
+	QueryType        string // select, insert, update, delete
+	Variable         map[string]string
+	Sqls             map[string]*SqlNode
+	DefaultNamespace string // namespace of current mapper
 }
 
 func NewContext() *Context {
@@ -24,5 +27,10 @@ func (c *Context) SetVariable(k, v string) {
 
 func (c *Context) GetSql(k string) (*SqlNode, bool) {
 	sql, ok := c.Sqls[k]
+	if ok {
+		return sql, true
+	}
+	// 当存在跨namespace引用时，需要通过namespace区分引用的SQL id
+	sql, ok = c.Sqls[fmt.Sprintf("%v.%v", c.DefaultNamespace, k)]
 	return sql, ok
 }

ast/mapper.go

@@ -70,7 +70,10 @@ func (m *Mapper) GetStmt(ctx *Context) (string, error) {
 
 func (m *Mapper) GetStmts(ctx *Context, skipErrorQuery bool) ([]string, error) {
 	var stmts []string
-	ctx.Sqls = m.SqlNodes
+	if len(ctx.Sqls) == 0 {
+		ctx.Sqls = m.SqlNodes
+	}
+	ctx.DefaultNamespace = m.NameSpace
 	for _, a := range m.QueryNodes {
 		data, err := a.GetStmt(ctx)
 		if err == nil {

ast/mappers.go

@@ -0,0 +1,44 @@
+package ast
+
+import (
+	"errors"
+	"fmt"
+)
+
+type Mappers struct {
+	mappers []*Mapper
+}
+
+func NewMappers() *Mappers {
+	return &Mappers{}
+}
+
+func (s *Mappers) AddMapper(ms ...*Mapper) error {
+	for _, m := range ms {
+		if m == nil {
+			return errors.New("can not add null mapper to mappers")
+		}
+		s.mappers = append(s.mappers, m)
+	}
+	return nil
+}
+
+func (s *Mappers) GetStmts(skipErrorQuery bool) ([]string, error) {
+	ctx := NewContext()
+	stmts := []string{}
+	for _, m := range s.mappers {
+		for id, node := range m.SqlNodes {
+			ctx.Sqls[fmt.Sprintf("%v.%v", m.NameSpace, id)] = node
+		}
+	}
+
+	for _, m := range s.mappers {
+		ctx.DefaultNamespace = m.NameSpace
+		stmt, err := m.GetStmts(ctx, skipErrorQuery)
+		if err != nil {
+			return nil, fmt.Errorf("get sqls from mapper failed, namespace: %v, err: %v", m.NameSpace, err)
+		}
+		stmts = append(stmts, stmt...)
+	}
+	return stmts, nil
+}

go.mod

@@ -5,4 +5,5 @@ go 1.14
 require (
 	github.com/pingcap/parser v3.0.12+incompatible
 	github.com/pingcap/tidb v0.0.0-20200312110807-8c4696b3f340 // v3.0.12
+	github.com/stretchr/testify v1.3.0
 )

parser.go

@@ -2,6 +2,7 @@ package parser
 
 import (
 	"encoding/xml"
+	"errors"
 	"fmt"
 	"io"
 	"strings"
@@ -27,6 +28,47 @@ func ParseXML(data string) (string, error) {
 	return stmt, nil
 }
 
+// ParseXMLs is a parser for parse all query in several XML files to []string one by one;
+// you can set `skipErrorQuery` true to ignore invalid query.
+func ParseXMLs(data []string, skipErrorQuery bool) ([]string, error) {
+	ms := ast.NewMappers()
+	for i := range data {
+		r := strings.NewReader(data[i])
+		d := xml.NewDecoder(r)
+		n, err := parse(d)
+		if err != nil {
+			if skipErrorQuery {
+				continue
+			} else {
+				return nil, err
+			}
+		}
+
+		if n == nil {
+			continue
+		}
+
+		m, ok := n.(*ast.Mapper)
+		if !ok {
+			if skipErrorQuery {
+				continue
+			} else {
+				return nil, errors.New("the mapper is not found")
+			}
+		}
+		err = ms.AddMapper(m)
+		if err != nil && !skipErrorQuery {
+			return nil, fmt.Errorf("add mapper failed: %v", err)
+		}
+	}
+	stmts, err := ms.GetStmts(skipErrorQuery)
+	if err != nil {
+		return nil, err
+	}
+
+	return stmts, nil
+}
+
 // ParseXMLQuery is a parser for parse all query in XML to []string one by one;
 // you can set `skipErrorQuery` true to ignore invalid query.
 func ParseXMLQuery(data string, skipErrorQuery bool) ([]string, error) {

parser_test.go

@@ -2,6 +2,8 @@ package parser
 
 import (
 	"testing"
+
+	"github.com/stretchr/testify/assert"
 )
 
 func testParser(t *testing.T, xmlData, expect string) {
@@ -1046,3 +1048,62 @@ func TestOtherwise_issue1193(t *testing.T) {
         `, "SELECT * FROM `fruits` WHERE `name`=? AND `price`=? AND `category`=?;",
 	)
 }
+
+func TestParseXMLs(t *testing.T) {
+	xmlCommonData := `
+<?xml version="1.0" encoding="UTF-8"?><!--Converted at: Mon Jun 07 09:48:24 CST 2021-->
+<!DOCTYPE mapper PUBLIC "-//mybatis.org//DTD Mapper 3.0//EN"
+		"http://mybatis.org/dtd/mybatis-3-mapper.dtd">
+<mapper namespace="test.common">
+	<sql id="prefix">
+		SELECT * FROM (
+	</sql>
+
+	<sql id="suffix">
+		WHERE a=1 )
+	</sql>
+
+	<select id="sql1" parameterType="customer" resultMap="custResultMap">
+		<include refid="prefix"/>
+		SELECT a,b FROM tb1
+		<include refid="suffix"/>
+	</select>
+</mapper>
+`
+	xmlData := `
+<?xml version="1.0" encoding="UTF-8"?><!--Converted at: Tue May 10 15:50:21 CST 2022-->
+<!DOCTYPE mapper PUBLIC "-//mybatis.org//DTD Mapper 3.0//EN"
+"http://mybatis.org/dtd/mybatis-3-mapper.dtd">
+<mapper namespace="employee">
+  <select id="queryEmpHireSepList" parameterType="employee" resultType="employeeResult">
+	    <include refid="test.common.prefix"/>
+	   	SELECT a,b FROM tb1 
+		<include refid="test.common.suffix"/>
+  </select>
+</mapper>
+`
+
+	sqls, err := ParseXMLs([]string{xmlCommonData, xmlData}, false)
+	if err != nil {
+		if !assert.NoError(t, err) {
+			t.Fatal(err)
+		}
+	}
+	assert.Equal(t, 2, len(sqls))
+	assert.Equal(t, `
+		SELECT * FROM (
+	
+		SELECT a,b FROM tb1
+		
+		WHERE a=1 )
+	`, sqls[0])
+
+	assert.Equal(t, `
+		SELECT * FROM (
+	
+	   	SELECT a,b FROM tb1 
+		
+		WHERE a=1 )
+	`, sqls[1])
+
+}