org.babyfish.jimmer.sql.fetcher.impl.FetcherTask Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of jimmer-sql Show documentation
A revolutionary ORM framework for both java and kotlin
The newest version!
package org.babyfish.jimmer.sql.fetcher.impl;

import org.babyfish.jimmer.meta.EmbeddedLevel;
import org.babyfish.jimmer.meta.ImmutableProp;
import org.babyfish.jimmer.meta.PropId;
import org.babyfish.jimmer.meta.TargetLevel;
import org.babyfish.jimmer.runtime.DraftContext;
import org.babyfish.jimmer.runtime.DraftSpi;
import org.babyfish.jimmer.runtime.ImmutableSpi;
import org.babyfish.jimmer.runtime.Internal;
import org.babyfish.jimmer.sql.fetcher.Fetcher;
import org.babyfish.jimmer.sql.fetcher.Field;
import org.babyfish.jimmer.sql.fetcher.RecursionStrategy;
import org.babyfish.jimmer.sql.runtime.JSqlClientImplementor;

import java.sql.Connection;
import java.util.*;
import java.util.stream.Collectors;

class FetcherTask {

    private final FetchingCache cache;

    private final JSqlClientImplementor sqlClient;

    private final Field field;

    private final int batchSize;

    private final DataLoader dataLoader;

    private Map pendingMap = new LinkedHashMap<>();

    public FetcherTask(
            FetchingCache cache,
            JSqlClientImplementor sqlClient,
            Connection con,
            FetchPath path,
            Field field
    ) {
        this.cache = cache;
        this.sqlClient = sqlClient;
        this.field = field;
        this.batchSize = determineBatchSize();
        this.dataLoader = new DataLoader(sqlClient, con, path, field);
    }

    public void add(DraftSpi draft) {
        add(draft, 1);
    }

    private void add(DraftSpi draft, int depth) {
        if (isLoaded(draft)) {
            return;
        }
        Object key = cache.createKey(field, draft);
        if (key == null) {
            return;
        }
        Object value = cache.get(field, key);
        if (value != null) {
            setDraftProp(draft, FetchingCache.unwrap(value), field);
            return;
        }
        pendingMap.computeIfAbsent(key, it -> new TaskData(key, depth)).getDrafts().add(draft);
    }

    public boolean execute() {
        if (pendingMap.isEmpty()) {
            return true;
        }
        Map handledMap;
        if (pendingMap.size() > batchSize) {
            Iterator> itr =
                    pendingMap.entrySet().iterator();
            if (batchSize == 1) {
                Map.Entry e = itr.next();
                handledMap = Collections.singletonMap(e.getKey(), e.getValue());
                itr.remove();
            } else {
                handledMap = new LinkedHashMap<>((batchSize * 4 + 2) / 3);
                for (int i = batchSize; i > 0; --i) {
                    Map.Entry e = itr.next();
                    handledMap.put(e.getKey(), e.getValue());
                    itr.remove();
                }
            }
        } else {
            handledMap = this.pendingMap;
            pendingMap = new LinkedHashMap<>();
        }
        Iterator> handledEntryItr =
                handledMap.entrySet().iterator();
        while (handledEntryItr.hasNext()) {
            Map.Entry e = handledEntryItr.next();
            Object key = e.getKey();
            Object value = cache.get(field, key);
            if (value != null) {
                value = FetchingCache.unwrap(value);
                TaskData taskData = e.getValue();
                afterLoad(taskData, value,false);
                handledEntryItr.remove();
            }
        }
        if (!handledMap.isEmpty()) {
            Map loadedMap = dataLoader.load(
                    handledMap
                            .values()
                            .stream()
                            .map(it -> it.getDrafts().get(0))
                            .collect(Collectors.toList())
            );
            for (Map.Entry e : handledMap.entrySet()) {
                TaskData taskData = e.getValue();
                Object value = loadedMap.get(taskData.getDrafts().get(0));
                afterLoad(taskData, value, true);
            }
        }
        return pendingMap.isEmpty();
    }

    private boolean isLoaded(DraftSpi draft) {
        if (!isLoaded(draft, field)) {
            return false;
        }
        Fetcher childFetcher = field.getChildFetcher(true);
        Object childValue = draft.__get(field.getProp().getId());
        if (childFetcher != null && childValue != null) {
            if (sqlClient.getFilters().getFilter(field.getProp().getTargetType()) != null) {
                return false;
            }
            for (Field childField : childFetcher.getFieldMap().values()) {
                if (!isLoaded(childValue, childField)) {
                    return false;
                }
            }
        }
        return true;
    }

    @SuppressWarnings("unchecked")
    private static boolean isLoaded(Object obj, Field field) {
        if (obj instanceof List) {
            List drafts = (List) obj;
            for (DraftSpi draft : drafts) {
                if (!isLoaded(draft, field)) {
                    return false;
                }
            }
            return true;
        }
        return ((DraftSpi) obj).__isLoaded(field.getProp().getId());
    }

    @SuppressWarnings("unchecked")
    private void afterLoad(TaskData taskData, Object value, boolean updateCache) {
        if (updateCache) {
            cache.put(field, taskData.getKey(), value);
        }
        for (DraftSpi draft : taskData.getDrafts()) {
            setDraftProp(draft, value, field);
        }
        if (!field.getProp().isAssociation(TargetLevel.ENTITY)) {
            return;
        }
        RecursionStrategy