FragmentGenerator.java

/*
 *    Copyright 2006-2026 the original author or authors.
 *
 *    Licensed under the Apache License, Version 2.0 (the "License");
 *    you may not use this file except in compliance with the License.
 *    You may obtain a copy of the License at
 *
 *       https://www.apache.org/licenses/LICENSE-2.0
 *
 *    Unless required by applicable law or agreed to in writing, software
 *    distributed under the License is distributed on an "AS IS" BASIS,
 *    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *    See the License for the specific language governing permissions and
 *    limitations under the License.
 */
package org.mybatis.generator.runtime.dynamicsql.java.elements;

import static org.mybatis.generator.api.dom.OutputUtilities.javaIndent;

import java.util.ArrayList;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Objects;
import java.util.Set;
import java.util.stream.Collectors;

import org.jspecify.annotations.Nullable;
import org.mybatis.generator.api.CommentGenerator;
import org.mybatis.generator.api.IntrospectedColumn;
import org.mybatis.generator.api.IntrospectedTable;
import org.mybatis.generator.api.dom.java.FullyQualifiedJavaType;
import org.mybatis.generator.api.dom.java.Method;
import org.mybatis.generator.api.dom.java.Parameter;
import org.mybatis.generator.internal.util.JavaBeansUtil;
import org.mybatis.generator.internal.util.StringUtility;
import org.mybatis.generator.runtime.JavaMethodAndImports;
import org.mybatis.generator.runtime.JavaMethodParts;
import org.mybatis.generator.runtime.mybatis3.ListUtilities;

public class FragmentGenerator {

    private final IntrospectedTable introspectedTable;
    private final String resultMapId;
    private final String tableFieldName;
    protected final boolean useSnakeCase;
    private final FullyQualifiedJavaType recordType;
    private final CommentGenerator commentGenerator;

    private FragmentGenerator(Builder builder) {
        this.introspectedTable = Objects.requireNonNull(builder.introspectedTable);
        this.resultMapId = Objects.requireNonNull(builder.resultMapId);
        tableFieldName = Objects.requireNonNull(builder.tableFieldName);
        useSnakeCase = builder.useSnakeCase;
        recordType = Objects.requireNonNull(builder.recordType);
        commentGenerator = Objects.requireNonNull(builder.commentGenerator);
    }

    public String calculateFieldName(String tableFieldName, IntrospectedColumn column) {
        String fieldName = column.getJavaProperty();
        if (useSnakeCase) {
            fieldName = StringUtility.convertCamelCaseToSnakeCase(fieldName);
        }

        if (fieldName.equals(tableFieldName)) {
            // name collision, no shortcut generated
            fieldName = tableFieldName + "." + fieldName; //$NON-NLS-1$
        }
        return fieldName;
    }

    public String getSelectList() {
        return introspectedTable.getAllColumns().stream()
                .map(c -> calculateFieldName(tableFieldName, c))
                .collect(Collectors.joining(", ")); //$NON-NLS-1$
    }

    public JavaMethodParts getPrimaryKeyWhereClauseAndParameters() {
        JavaMethodParts.Builder builder = new JavaMethodParts.Builder();

        boolean first = true;
        for (IntrospectedColumn column : introspectedTable.getPrimaryKeyColumns()) {
            String parameterName = column.getJavaProperty();
            if (!useSnakeCase) {
                parameterName += "_"; //$NON-NLS-1$
            }

            String fieldName = calculateFieldName(tableFieldName, column);
            builder.withImport(column.getFullyQualifiedJavaType());
            builder.withParameter(new Parameter(column.getFullyQualifiedJavaType(), parameterName));
            if (first) {
                builder.withBodyLine(javaIndent(1) + "c.where(" + fieldName //$NON-NLS-1$
                        + ", isEqualTo(" + parameterName //$NON-NLS-1$
                        + "))"); //$NON-NLS-1$
                first = false;
            } else {
                builder.withBodyLine(javaIndent(1) + ".and(" + fieldName //$NON-NLS-1$
                        + ", isEqualTo(" + parameterName //$NON-NLS-1$
                        + "))"); //$NON-NLS-1$
            }
        }
        builder.withBodyLine(");"); //$NON-NLS-1$

        return builder.build();
    }

    public List<String> getPrimaryKeyWhereClauseForUpdate(String prefix) {
        List<String> lines = new ArrayList<>();

        boolean first = true;
        for (IntrospectedColumn column : introspectedTable.getPrimaryKeyColumns()) {
            String fieldName = calculateFieldName(tableFieldName, column);
            String methodName = JavaBeansUtil.getCallingGetterMethodName(column);
            if (first) {
                lines.add(prefix + ".where(" + fieldName //$NON-NLS-1$
                        + ", isEqualTo(row::" + methodName //$NON-NLS-1$
                        + "))"); //$NON-NLS-1$
                first = false;
            } else {
                lines.add(prefix + ".and(" + fieldName //$NON-NLS-1$
                        + ", isEqualTo(row::" + methodName //$NON-NLS-1$
                        + "))"); //$NON-NLS-1$
            }
        }

        return lines;
    }

    public JavaMethodParts getAnnotatedConstructorArgs() {
        JavaMethodParts.Builder builder = new JavaMethodParts.Builder();

        builder.withImport(new FullyQualifiedJavaType("org.apache.ibatis.type.JdbcType")); //$NON-NLS-1$
        builder.withImport(new FullyQualifiedJavaType("org.apache.ibatis.annotations.Arg")); //$NON-NLS-1$
        builder.withImport(new FullyQualifiedJavaType("org.apache.ibatis.annotations.Results")); //$NON-NLS-1$

        builder.withAnnotation("@Results(id=\"" + resultMapId + "\")"); //$NON-NLS-1$ //$NON-NLS-2$

        Set<FullyQualifiedJavaType> imports = new HashSet<>();
        for (IntrospectedColumn introspectedColumn : introspectedTable.getPrimaryKeyColumns()) {
            builder.withAnnotation(getArgAnnotation(imports, introspectedColumn, true));
        }

        for (IntrospectedColumn introspectedColumn : introspectedTable.getNonPrimaryKeyColumns()) {
            builder.withAnnotation(getArgAnnotation(imports, introspectedColumn, false));
        }

        return builder.withImports(imports).build();
    }

    public JavaMethodParts getAnnotatedResults() {
        JavaMethodParts.Builder builder = new JavaMethodParts.Builder();

        builder.withImport(new FullyQualifiedJavaType("org.apache.ibatis.type.JdbcType")); //$NON-NLS-1$
        builder.withImport(new FullyQualifiedJavaType("org.apache.ibatis.annotations.Result")); //$NON-NLS-1$
        builder.withImport(new FullyQualifiedJavaType("org.apache.ibatis.annotations.Results")); //$NON-NLS-1$

        builder.withAnnotation("@Results(id=\"" + resultMapId + "\", value = {"); //$NON-NLS-1$ //$NON-NLS-2$

        StringBuilder sb = new StringBuilder();

        Set<FullyQualifiedJavaType> imports = new HashSet<>();
        Iterator<IntrospectedColumn> iterPk = introspectedTable.getPrimaryKeyColumns().iterator();
        Iterator<IntrospectedColumn> iterNonPk = introspectedTable.getNonPrimaryKeyColumns().iterator();
        while (iterPk.hasNext()) {
            IntrospectedColumn introspectedColumn = iterPk.next();
            sb.setLength(0);
            javaIndent(sb, 1);
            sb.append(getResultAnnotation(imports, introspectedColumn, true));

            if (iterPk.hasNext() || iterNonPk.hasNext()) {
                sb.append(',');
            }

            builder.withAnnotation(sb.toString());
        }

        while (iterNonPk.hasNext()) {
            IntrospectedColumn introspectedColumn = iterNonPk.next();
            sb.setLength(0);
            javaIndent(sb, 1);
            sb.append(getResultAnnotation(imports, introspectedColumn, false));

            if (iterNonPk.hasNext()) {
                sb.append(',');
            }

            builder.withAnnotation(sb.toString());
        }

        builder.withAnnotation("})") //$NON-NLS-1$
                .withImports(imports);

        return builder.build();
    }

    private String getArgAnnotation(Set<FullyQualifiedJavaType> imports, IntrospectedColumn introspectedColumn,
                                    boolean idColumn) {
        imports.add(introspectedColumn.getFullyQualifiedJavaType());

        return "@Arg(column=\"" //$NON-NLS-1$
                + introspectedColumn.getActualColumnName()
                + "\", javaType=" //$NON-NLS-1$
                + introspectedColumn.getFullyQualifiedJavaType().getShortName()
                + ".class" //$NON-NLS-1$
                + generateAdditionalItems(imports, introspectedColumn, idColumn)
                + ')';//$NON-NLS-1$
    }

    private String getResultAnnotation(Set<FullyQualifiedJavaType> imports, IntrospectedColumn introspectedColumn,
            boolean idColumn) {
        return "@Result(column=\"" //$NON-NLS-1$
                + introspectedColumn.getActualColumnName()
                + "\", property=\"" //$NON-NLS-1$
                + introspectedColumn.getJavaProperty()
                + '\"'
                + generateAdditionalItems(imports, introspectedColumn, idColumn)
                + ')'; //$NON-NLS-1$
    }

    private String generateAdditionalItems(Set<FullyQualifiedJavaType> imports, IntrospectedColumn introspectedColumn,
                                           boolean idColumn) {
        StringBuilder sb = new StringBuilder();

        introspectedColumn.getTypeHandler().map(FullyQualifiedJavaType::new).ifPresent(th -> {
            imports.add(th);
            sb.append(", typeHandler="); //$NON-NLS-1$
            sb.append(th.getShortName());
            sb.append(".class"); //$NON-NLS-1$
        });

        sb.append(", jdbcType=JdbcType."); //$NON-NLS-1$
        sb.append(introspectedColumn.getJdbcTypeName());
        if (idColumn) {
            sb.append(", id=true"); //$NON-NLS-1$
        }

        return sb.toString();
    }

    public List<String> getSetEqualLinesForUpdateStatement(List<IntrospectedColumn> columnList, String firstLinePrefix,
                                                           String subsequentLinePrefix, boolean terminate) {
        return getSetLinesForUpdateStatement(columnList, firstLinePrefix, subsequentLinePrefix, terminate,
                "equalTo"); //$NON-NLS-1$
    }

    public List<String> getSetEqualWhenPresentLinesForUpdateStatement(List<IntrospectedColumn> columnList,
                                                                      String firstLinePrefix,
                                                                      String subsequentLinePrefix, boolean terminate) {
        return getSetLinesForUpdateStatement(columnList, firstLinePrefix, subsequentLinePrefix, terminate,
                "equalToWhenPresent"); //$NON-NLS-1$
    }

    private List<String> getSetLinesForUpdateStatement(List<IntrospectedColumn> columnList, String firstLinePrefix,
                                                       String subsequentLinePrefix, boolean terminate,
                                                       String fragment) {
        List<String> lines = new ArrayList<>();
        List<IntrospectedColumn> columns = ListUtilities.filterColumnsForUpdate(columnList);
        Iterator<IntrospectedColumn> iter = columns.iterator();
        boolean first = true;
        while (iter.hasNext()) {
            IntrospectedColumn column = iter.next();
            String fieldName = calculateFieldName(tableFieldName, column);
            String methodName = JavaBeansUtil.getCallingGetterMethodName(column);

            String start;
            if (first) {
                start = firstLinePrefix;
                first = false;
            } else {
                start = subsequentLinePrefix;
            }

            String line = start
                    + ".set(" //$NON-NLS-1$
                    + fieldName
                    + ")." //$NON-NLS-1$
                    + fragment
                    + "(row::" //$NON-NLS-1$
                    + methodName
                    + ")"; //$NON-NLS-1$

            if (terminate && !iter.hasNext()) {
                line += ";"; //$NON-NLS-1$
            }

            lines.add(line);
        }

        return lines;
    }

    public JavaMethodAndImports generateGeneralSelectMethod(boolean isDistinct) {
        String methodName;
        String utilMethodName;
        if (isDistinct) {
            methodName = "selectDistinct"; //$NON-NLS-1$
            utilMethodName = "selectDistinct"; //$NON-NLS-1$
        } else {
            methodName = "select"; //$NON-NLS-1$
            utilMethodName = "selectList"; //$NON-NLS-1$
        }

        Set<FullyQualifiedJavaType> imports = new HashSet<>();

        FullyQualifiedJavaType parameterType = new FullyQualifiedJavaType(
                "org.mybatis.dynamic.sql.dsl.SelectDSLCompleter"); //$NON-NLS-1$

        imports.add(parameterType);
        imports.add(new FullyQualifiedJavaType("org.mybatis.dynamic.sql.util.mybatis3.MyBatis3Utils")); //$NON-NLS-1$

        FullyQualifiedJavaType returnType = FullyQualifiedJavaType.getNewListInstance();
        returnType.addTypeArgument(recordType);

        imports.add(returnType);

        Method method = new Method(methodName);
        method.setDefault(true);
        method.addParameter(new Parameter(parameterType, "completer")); //$NON-NLS-1$

        commentGenerator.addGeneralMethodAnnotation(method, introspectedTable, imports);

        method.setReturnType(returnType);
        String line = String.format(
                "return MyBatis3Utils.%s(this::selectMany, selectList, %s, completer);", //$NON-NLS-1$
                utilMethodName, tableFieldName);
        method.addBodyLine(line);

        return JavaMethodAndImports.withMethod(method)
                .withImports(imports)
                .build();
    }

    public static class Builder {
        private @Nullable IntrospectedTable introspectedTable;
        private @Nullable String resultMapId;
        private @Nullable String tableFieldName;
        private boolean useSnakeCase;
        private @Nullable FullyQualifiedJavaType recordType;
        private @Nullable CommentGenerator commentGenerator;

        public Builder withIntrospectedTable(IntrospectedTable introspectedTable) {
            this.introspectedTable = introspectedTable;
            return this;
        }

        public Builder withResultMapId(String resultMapId) {
            this.resultMapId = resultMapId;
            return this;
        }

        public Builder withTableFieldName(String tableFieldName) {
            this.tableFieldName = tableFieldName;
            return this;
        }

        public Builder useSnakeCase(boolean useSnakeCase) {
            this.useSnakeCase = useSnakeCase;
            return this;
        }

        public Builder withRecordType(FullyQualifiedJavaType recordType) {
            this.recordType = recordType;
            return this;
        }

        public Builder withCommentGenerator(CommentGenerator commentGenerator) {
            this.commentGenerator = commentGenerator;
            return this;
        }

        public FragmentGenerator build() {
            return new FragmentGenerator(this);
        }
    }
}