zsjz-ai/src/main/java/com/qingjian/module/dm/clean/AbstractDataLoader.java

package com.qingjian.module.dm.clean;

import com.qingjian.common.base.Cleaned;
import com.qingjian.common.constants.StrConsts;
import lombok.extern.slf4j.Slf4j;
import org.duckdb.DuckDBAppender;
import org.duckdb.DuckDBConnection;
import org.noear.solon.Solon;
import org.noear.solon.data.dynamicds.DynamicDataSource;
import org.noear.solon.data.dynamicds.DynamicDsKey;

@Slf4j
public abstract class AbstractDataLoader<S, T extends Cleaned> implements DataLoader<S, T> {
    protected DuckDBConnection duckDBConnection;
    protected DuckDBAppender duckDBAppender;
    private final String tableName;

    public AbstractDataLoader(String tableName) {
        this.tableName = tableName;
        try {
            DynamicDataSource ds = Solon.context().getBean("db1");
            DynamicDsKey.with(StrConsts.DS_KEY_CASE, () -> {
                duckDBConnection = (DuckDBConnection) ds.getConnection();
                duckDBAppender = duckDBConnection.createAppender(DuckDBConnection.DEFAULT_SCHEMA, tableName);
            });
            log.info("初始化 DuckDB 数据加载器成功, tableName={}", tableName);
        } catch (Exception e) {
            log.error("初始化 DuckDB 数据加载器失败, tableName={}", tableName, e);
        }
    }

    @Override
    public void write(S record) {
        if (record == null) {
            log.info("跳过空记录写入, tableName={}", tableName);
            return;
        }
        T t = convert(record);
        long l = System.currentTimeMillis();
        fillData(t);
        log.info("写入 fillData 成功, tableName={}, cost={}", tableName, System.currentTimeMillis() - l);
        //doWrite(t);
    }

    @Override
    public T convert(S source) {
        return (T) source;
    }

    protected abstract void doWrite(T record);

    @Override
    public void fillData(T record) {

    }

    @Override
    public void close() {
        // 倒序关闭，且单独捕获异常，避免一个资源关闭失败影响另一个
        if (duckDBAppender != null) {
            try {
                duckDBAppender.close();
            } catch (Exception e) {
                // 添加日志记录 Appender 关闭失败
                log.error("DuckDB Appender 关闭失败, tableName={}", tableName, e);
            }
        }
        if (duckDBConnection != null) {
            try {
                duckDBConnection.close();
            } catch (Exception e) {
                log.error("DuckDB 连接关闭失败, tableName={}", tableName, e);
            }
        }
    }
}