package com.easyagents.federation.sql.adapter.jdbc;

import com.easyagents.federation.sql.api.FederationSqlEngine;
import com.easyagents.federation.sql.api.FederationSqlEngines;
import com.easyagents.federation.sql.api.FederationSqlErrorCode;
import com.easyagents.federation.sql.api.FederationSqlException;
import com.easyagents.federation.sql.api.SqlExecutionContext;
import com.easyagents.federation.sql.api.SqlQueryCommand;
import com.easyagents.federation.sql.compile.FederationSqlPlan;
import com.easyagents.federation.sql.compile.SqlCompileRequest;
import com.easyagents.federation.sql.compile.SqlExplainRequest;
import com.easyagents.federation.sql.execute.FederationResultCursor;
import com.easyagents.federation.sql.execute.FederationColumn;
import com.easyagents.federation.sql.execute.FederationQueryAdmissionController;
import com.easyagents.federation.sql.execute.FederationQueryMetricsSnapshot;
import com.easyagents.federation.sql.execute.FederationQueryPermit;
import com.easyagents.federation.sql.execute.LocalFederationQueryAdmissionController;
import com.easyagents.federation.sql.execute.QueryId;
import com.easyagents.federation.sql.execute.SqlExecutionOptions;
import com.easyagents.federation.sql.execute.SqlParameter;
import com.easyagents.federation.sql.execute.StatementLifecycle;
import com.easyagents.federation.sql.federation.FederationExecutionPolicy;
import com.easyagents.federation.sql.federation.FederationQueryScopeDefinition;
import com.easyagents.federation.sql.federation.FederationSourceBindingDefinition;
import com.easyagents.federation.sql.source.FederationDataSourceHandles;
import com.easyagents.federation.sql.source.FederationSourceDefinition;
import com.easyagents.federation.sql.source.RuntimeFingerprint;
import com.easyagents.federation.sql.source.SourceApplyOptions;
import com.easyagents.federation.sql.source.SourceId;
import java.io.PrintWriter;
import java.io.InputStream;
import java.io.Reader;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Proxy;
import java.math.BigDecimal;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.SQLTimeoutException;
import java.sql.Statement;
import java.sql.Types;
import java.time.Duration;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.BooleanSupplier;
import java.util.logging.Logger;
import javax.sql.DataSource;
import org.h2.jdbcx.JdbcDataSource;
import org.junit.After;
import org.junit.Assert;
import org.junit.Before;
import org.junit.Test;

/**
 * Calcite 编译、ServiceLoader Adapter 和直接 JDBC 流式执行集成测试。
 */
public class JdbcFederationSqlEngineTest {

    private static final SourceId SOURCE_ID = new SourceId("main");

    private TrackingDataSource dataSource;
    private FederationSqlEngine engine;

    /**
     * 创建隔离 H2 数据库与预热的数据源 Runtime。
     *
     * @throws SQLException 初始化失败
     */
    @Before
    public void setUp() throws SQLException {
        JdbcDataSource h2 = new JdbcDataSource();
        h2.setURL("jdbc:h2:mem:federation_" + System.nanoTime() + ";DB_CLOSE_DELAY=-1");
        dataSource = new TrackingDataSource(h2);
        try (Connection connection = dataSource.getConnection(); Statement statement = connection.createStatement()) {
            statement.execute("CREATE TABLE DEPARTMENT (ID INT PRIMARY KEY, NAME VARCHAR(64) NOT NULL)");
            statement.execute(
                "CREATE TABLE PERSON (ID INT PRIMARY KEY, NAME VARCHAR(64) NOT NULL, DEPARTMENT_ID INT NOT NULL)"
            );
            statement.execute("INSERT INTO DEPARTMENT VALUES (10, 'Engineering'), (20, 'Finance')");
            statement.execute(
                "INSERT INTO PERSON VALUES (1, 'Alice', 10), (2, 'Bob', 10), (3, 'Carol', 20)"
            );
        }

        RuntimeFingerprint fingerprint = new RuntimeFingerprint("H2", "2", "H2 JDBC Driver", "2", "1");
        engine = FederationSqlEngines.builder()
            .dataSourceResolver(definition -> FederationDataSourceHandles.shared(dataSource, fingerprint))
            .maximumPlanCacheEntries(32)
            .build();
        engine.sources().apply(definition(1), SourceApplyOptions.prewarmNow());
        Assert.assertEquals(0, dataSource.activeConnections());
    }

    /**
     * 关闭 Engine。
     */
    @After
    public void tearDown() {
        if (engine != null) {
            engine.close();
        }
    }

    /**
     * 验证参数、同源 Join、排序、分页、Explain 和流式资源释放。
     */
    @Test
    public void shouldCompileAndStreamSameSourceJoin() {
        String sql = "SELECT p.NAME, d.NAME AS DEPARTMENT_NAME "
            + "FROM PERSON p JOIN DEPARTMENT d ON p.DEPARTMENT_ID = d.ID "
            + "WHERE p.ID > ? ORDER BY p.ID FETCH NEXT 2 ROWS ONLY";
        SqlQueryCommand command = SqlQueryCommand.of(
            sql,
            SOURCE_ID,
            1,
            List.of(new SqlParameter(Types.INTEGER, 1))
        );

        try (FederationResultCursor cursor = engine.query(command)) {
            Assert.assertEquals(1, dataSource.activeConnections());
            Assert.assertEquals(2, cursor.columns().size());
            Assert.assertTrue(cursor.next());
            Assert.assertEquals(List.of("Bob", "Engineering"), cursor.row());
            Assert.assertTrue(cursor.next());
            Assert.assertEquals(List.of("Carol", "Finance"), cursor.row());
            Assert.assertFalse(cursor.next());
        }
        Assert.assertEquals(0, dataSource.activeConnections());

        FederationSqlPlan plan = engine.compile(new SqlCompileRequest(
            sql,
            SOURCE_ID,
            1,
            List.of(Types.INTEGER),
            "policy-v1"
        ));
        Assert.assertTrue(plan.executable());
        Assert.assertEquals(1, plan.parameterCount());
        Assert.assertEquals(List.of(0), plan.parameterMapping());
        Assert.assertTrue(plan.executableSql().contains("?"));
        Assert.assertTrue(plan.executableSql().toUpperCase().contains("PUBLIC"));
        Assert.assertFalse(plan.executableSql().toUpperCase().contains("APP"));
        Assert.assertFalse(plan.executableSql().contains("Bob"));
        Assert.assertFalse(
            engine.explain(new SqlExplainRequest(SqlCompileRequest.of(
                "SELECT COUNT(*) FROM PERSON",
                SOURCE_ID,
                1
            ))).relationalPlan().isBlank()
        );
    }

    /**
     * 验证缓存计划的每次签发保留各自的编译耗时与命中状态。
     */
    @Test
    public void shouldKeepPlanningMetricsOnEachIssuedPlan() {
        SqlCompileRequest request = SqlCompileRequest.of(
            "SELECT ID FROM PERSON WHERE ID >= 1 ORDER BY ID",
            SOURCE_ID,
            1
        );
        FederationSqlPlan coldPlan = engine.compile(request);
        FederationSqlPlan hotPlan = engine.compile(request);

        Assert.assertNotSame(coldPlan, hotPlan);
        FederationQueryMetricsSnapshot coldMetrics;
        try (FederationResultCursor cursor = engine.execute(
            coldPlan,
            SqlExecutionContext.of(List.of())
        )) {
            while (cursor.next()) {
                cursor.row();
            }
            coldMetrics = cursor.metrics();
        }
        FederationQueryMetricsSnapshot hotMetrics;
        try (FederationResultCursor cursor = engine.execute(
            hotPlan,
            SqlExecutionContext.of(List.of())
        )) {
            while (cursor.next()) {
                cursor.row();
            }
            hotMetrics = cursor.metrics();
        }

        Assert.assertFalse(coldMetrics.planCacheHit());
        Assert.assertTrue(hotMetrics.planCacheHit());
    }

    /**
     * 验证 DML 在获取执行连接前被只读基线拒绝。
     */
    @Test
    public void shouldRejectDataModificationSql() {
        try {
            engine.compile(SqlCompileRequest.of("DELETE FROM PERSON", SOURCE_ID, 1));
            Assert.fail("DML should be rejected");
        } catch (FederationSqlException exception) {
            Assert.assertEquals(FederationSqlErrorCode.SQL_NOT_READ_ONLY, exception.errorCode());
        }
        Assert.assertEquals(0, dataSource.activeConnections());
    }

    /**
     * 验证参数数量错误具有稳定错误码且不会借出执行连接。
     */
    @Test
    public void shouldRejectParameterCountMismatch() {
        try {
            engine.query(SqlQueryCommand.of(
                "SELECT NAME FROM PERSON WHERE ID = ?",
                SOURCE_ID,
                1,
                List.of()
            ));
            Assert.fail("parameter mismatch should be rejected");
        } catch (FederationSqlException exception) {
            Assert.assertEquals(FederationSqlErrorCode.PARAMETER_COUNT_MISMATCH, exception.errorCode());
        }

        FederationSqlPlan integerPlan = engine.compile(new SqlCompileRequest(
            "SELECT NAME FROM PERSON WHERE ID = ?",
            SOURCE_ID,
            1,
            List.of(Types.INTEGER),
            "policy-v1"
        ));
        try {
            engine.execute(
                integerPlan,
                SqlExecutionContext.of(List.of(new SqlParameter(Types.VARCHAR, "1")))
            );
            Assert.fail("parameter type mismatch should be rejected");
        } catch (FederationSqlException exception) {
            Assert.assertEquals(FederationSqlErrorCode.PARAMETER_COUNT_MISMATCH, exception.errorCode());
        }
        Assert.assertEquals(0, dataSource.activeConnections());
    }

    /**
     * 验证当前行复制允许数据库 NULL，且提前关闭游标立即归还连接。
     */
    @Test
    public void shouldReturnNullValuesAndReleaseConnectionOnEarlyClose() {
        FederationResultCursor cursor = engine.query(SqlQueryCommand.of(
            "SELECT CAST(NULL AS VARCHAR) AS OPTIONAL_VALUE FROM PERSON ORDER BY ID",
            SOURCE_ID,
            1,
            List.of()
        ));
        Assert.assertEquals(1, dataSource.activeConnections());
        Assert.assertTrue(cursor.next());
        Assert.assertEquals(Arrays.asList((Object) null), cursor.row());
        cursor.close();
        Assert.assertEquals(0, dataSource.activeConnections());
    }

    /**
     * 验证大文本和二进制列可通过流读取，并在关闭游标后归还连接。
     *
     * @throws Exception 流读取失败
     */
    @Test
    public void shouldExposeCharacterAndBinaryStreams() throws Exception {
        try (FederationResultCursor cursor = engine.query(SqlQueryCommand.of(
            "SELECT CAST('hello' AS VARCHAR), CAST(X'0102' AS VARBINARY) "
                + "FROM PERSON FETCH NEXT 1 ROWS ONLY",
            SOURCE_ID,
            1,
            List.of()
        ))) {
            Assert.assertTrue(cursor.next());
            try (Reader reader = cursor.getCharacterStream(1);
                 InputStream input = cursor.getBinaryStream(2)) {
                char[] text = new char[5];
                Assert.assertEquals(5, reader.read(text));
                Assert.assertEquals("hello", new String(text));
                Assert.assertArrayEquals(new byte[] {1, 2}, input.readAllBytes());
            }
        }
        Assert.assertEquals(0, dataSource.activeConnections());
    }

    /**
     * 验证计划缓存命中后的编译门面 p95 保持在 5ms 性能预算内。
     */
    @Test
    public void shouldKeepHotCompileFacadeWithinBudget() {
        SqlCompileRequest request = SqlCompileRequest.of(
            "SELECT ID, NAME FROM PERSON WHERE ID > 0 ORDER BY ID",
            SOURCE_ID,
            1
        );
        for (int index = 0; index < 10_000; index++) {
            engine.compile(request);
        }

        long[] durations = new long[10_000];
        for (int index = 0; index < durations.length; index++) {
            long started = System.nanoTime();
            engine.compile(request);
            durations[index] = System.nanoTime() - started;
        }
        Arrays.sort(durations);
        long p95Nanos = durations[(int) (durations.length * 0.95)];
        Assert.assertTrue(
            "hot compile p95 was " + p95Nanos + "ns",
            p95Nanos <= 5_000_000L
        );
    }

    /**
     * 验证 NULL 必须携带真实 JDBC 类型，并能通过 PreparedStatement 正确绑定。
     */
    @Test
    public void shouldRequireAndExecuteExplicitlyTypedNullParameter() {
        try {
            SqlParameter.of(null);
            Assert.fail("untyped NULL should be rejected");
        } catch (IllegalArgumentException expected) {
            Assert.assertTrue(expected.getMessage().contains("jdbcType"));
        }

        try (FederationResultCursor cursor = engine.query(SqlQueryCommand.of(
            "SELECT COALESCE(CAST(? AS VARCHAR), NAME) FROM PERSON WHERE ID = 1",
            SOURCE_ID,
            1,
            List.of(new SqlParameter(Types.VARCHAR, null))
        ))) {
            Assert.assertTrue(cursor.next());
            Assert.assertEquals("Alice", cursor.getObject(1));
            Assert.assertFalse(cursor.next());
        }
        Assert.assertEquals(0, dataSource.activeConnections());
    }

    /**
     * 验证声明的 JDBC 类型通过 Calcite CAST 参与动态参数校验和结果类型推导。
     */
    @Test
    public void shouldApplyDeclaredParameterTypeDuringCompilation() {
        String sql = "SELECT ? AS TYPED_VALUE FROM PERSON FETCH NEXT 1 ROWS ONLY";
        FederationSqlPlan plan = engine.compile(new SqlCompileRequest(
            sql,
            SOURCE_ID,
            1,
            List.of(Types.INTEGER),
            "default"
        ));

        Assert.assertEquals(Types.INTEGER, plan.columns().get(0).jdbcType());
        Assert.assertTrue(plan.normalizedSql().toUpperCase().contains("CAST(? AS INTEGER)"));
        Assert.assertFalse(
            plan.executableSql(),
            plan.executableSql().toUpperCase().contains("CAST(? AS INTEGER)")
        );
        try (FederationResultCursor cursor = engine.query(SqlQueryCommand.of(
            sql,
            SOURCE_ID,
            1,
            List.of(new SqlParameter(Types.INTEGER, 7))
        ))) {
            Assert.assertTrue(cursor.next());
            Assert.assertEquals(7, cursor.getObject(1));
            Assert.assertFalse(cursor.next());
        }
        Assert.assertEquals(0, dataSource.activeConnections());
    }

    /**
     * 验证 JDBC 4.2 时区参数和 UUID 参数能参与 Calcite 推导并正常执行。
     */
    @Test
    public void shouldCompileAndExecuteJdbc42ScalarParameters() {
        java.time.OffsetTime offsetTime = java.time.OffsetTime.parse("12:00:00+08:00");
        java.time.OffsetDateTime offsetTimestamp = java.time.OffsetDateTime.parse(
            "2026-08-21T12:00:00+08:00"
        );
        java.util.UUID uuid = java.util.UUID.randomUUID();
        String sql = "SELECT ? AS OFFSET_TIME_VALUE, ? AS OFFSET_TIMESTAMP_VALUE, "
            + "? AS UUID_VALUE FROM PERSON FETCH NEXT 1 ROWS ONLY";
        SqlQueryCommand command = SqlQueryCommand.of(
            sql,
            SOURCE_ID,
            1,
            List.of(
                SqlParameter.of(offsetTime),
                SqlParameter.of(offsetTimestamp),
                SqlParameter.of(uuid)
            )
        );

        FederationSqlPlan plan = engine.compile(new SqlCompileRequest(
            sql,
            SOURCE_ID,
            1,
            List.of(
                Types.TIME_WITH_TIMEZONE,
                Types.TIMESTAMP_WITH_TIMEZONE,
                Types.OTHER
            ),
            "default"
        ));
        Assert.assertEquals(Types.TIME_WITH_TIMEZONE, plan.columns().get(0).jdbcType());
        Assert.assertEquals(Types.TIMESTAMP_WITH_TIMEZONE, plan.columns().get(1).jdbcType());
        Assert.assertEquals(Types.OTHER, plan.columns().get(2).jdbcType());

        try (FederationResultCursor cursor = engine.query(command)) {
            Assert.assertTrue(cursor.next());
            Assert.assertEquals(offsetTime, cursor.getObject(1));
            Assert.assertEquals(offsetTimestamp, cursor.getObject(2));
            Assert.assertEquals(uuid, cursor.getObject(3));
            Assert.assertFalse(cursor.next());
        }
        Assert.assertEquals(0, dataSource.activeConnections());
    }

    /**
     * 验证仅用于 Calcite 推导的类型 CAST 不会下推并改变小数参数语义。
     */
    @Test
    public void shouldPreserveDecimalParameterValueDuringExecution() {
        BigDecimal value = new BigDecimal("123.45");
        String sql = "SELECT ? AS DECIMAL_VALUE FROM PERSON FETCH NEXT 1 ROWS ONLY";
        FederationSqlPlan plan = engine.compile(new SqlCompileRequest(
            sql,
            SOURCE_ID,
            1,
            List.of(Types.DECIMAL),
            "default"
        ));
        Assert.assertFalse(
            plan.executableSql(),
            plan.executableSql().toUpperCase().contains("CAST(? AS DECIMAL)")
        );
        SqlQueryCommand command = SqlQueryCommand.of(
            sql,
            SOURCE_ID,
            1,
            List.of(SqlParameter.of(value))
        );

        try (FederationResultCursor cursor = engine.query(command)) {
            Assert.assertTrue(cursor.next());
            Assert.assertEquals(0, value.compareTo((BigDecimal) cursor.getObject(1)));
            Assert.assertFalse(cursor.next());
        }
        Assert.assertEquals(0, dataSource.activeConnections());
    }

    /**
     * 验证公开执行选项不能关闭正式查询的 JDBC 只读防线。
     */
    @Test(expected = IllegalArgumentException.class)
    public void shouldRejectDisabledReadOnlyExecution() {
        new SqlExecutionOptions(100, 0, 30, false);
    }

    /**
     * 验证缓存命中和高级 execute 均会重新执行动态策略校验。
     */
    @Test
    public void shouldRevalidatePolicyOnCacheHitAndExecute() {
        AtomicBoolean allowed = new AtomicBoolean(true);
        AtomicInteger evaluations = new AtomicInteger();
        RuntimeFingerprint fingerprint = new RuntimeFingerprint("H2", "2", "H2 JDBC Driver", "2", "1");
        FederationSqlEngine policyEngine = FederationSqlEngines.builder()
            .dataSourceResolver(definition -> FederationDataSourceHandles.shared(dataSource, fingerprint))
            .policy(context -> {
                evaluations.incrementAndGet();
                if (!allowed.get()) {
                    throw new FederationSqlException(
                        FederationSqlErrorCode.INVALID_ARGUMENT,
                        "query is denied by the current policy"
                    );
                }
            })
            .build();
        try {
            policyEngine.sources().apply(definition(1), SourceApplyOptions.prewarmNow());
            SqlCompileRequest request = SqlCompileRequest.of("SELECT ID FROM PERSON", SOURCE_ID, 1);
            FederationSqlPlan plan = policyEngine.compile(request);
            FederationSqlPlan cached = policyEngine.compile(request);
            Assert.assertNotSame(plan, cached);
            Assert.assertSame(plan.relRoot(), cached.relRoot());
            Assert.assertEquals(2, evaluations.get());

            allowed.set(false);
            assertPolicyDenied(() -> policyEngine.compile(request));
            assertPolicyDenied(() -> policyEngine.execute(plan, SqlExecutionContext.of(List.of())));
            Assert.assertEquals(4, evaluations.get());
        } finally {
            policyEngine.close();
        }
        Assert.assertEquals(0, dataSource.activeConnections());
    }

    /**
     * 验证实际 Runtime 相同时，不同最小 revision 的计划仍保留各自策略上下文。
     */
    @Test
    public void shouldIsolatePolicyContextByMinimumRevision() {
        AtomicBoolean denyRevisionOne = new AtomicBoolean();
        RuntimeFingerprint fingerprint = new RuntimeFingerprint("H2", "2", "H2 JDBC Driver", "2", "1");
        FederationSqlEngine policyEngine = FederationSqlEngines.builder()
            .dataSourceResolver(definition -> FederationDataSourceHandles.shared(dataSource, fingerprint))
            .policy(context -> {
                if (denyRevisionOne.get() && context.request().minimumRevision() == 1) {
                    throw new FederationSqlException(
                        FederationSqlErrorCode.INVALID_ARGUMENT,
                        "minimum revision 1 is denied"
                    );
                }
            })
            .build();
        try {
            policyEngine.sources().apply(definition(2), SourceApplyOptions.prewarmNow());
            String sql = "SELECT ID FROM PERSON WHERE ID = 1";
            FederationSqlPlan revisionOnePlan = policyEngine.compile(
                SqlCompileRequest.of(sql, SOURCE_ID, 1)
            );
            FederationSqlPlan revisionTwoPlan = policyEngine.compile(
                SqlCompileRequest.of(sql, SOURCE_ID, 2)
            );
            Assert.assertNotSame(revisionOnePlan, revisionTwoPlan);

            denyRevisionOne.set(true);
            assertPolicyDenied(() -> policyEngine.execute(
                revisionOnePlan,
                SqlExecutionContext.of(List.of())
            ));
            try (FederationResultCursor cursor = policyEngine.execute(
                revisionTwoPlan,
                SqlExecutionContext.of(List.of())
            )) {
                Assert.assertTrue(cursor.next());
                Assert.assertEquals(1, cursor.getObject(1));
                Assert.assertFalse(cursor.next());
            }
        } finally {
            policyEngine.close();
        }
        Assert.assertEquals(0, dataSource.activeConnections());
    }

    /**
     * 验证调用方实现的 Plan 接口无法绕过当前 Engine 的签发与只读策略。
     */
    @Test
    public void shouldRejectPlanNotAuthorizedByCurrentEngine() {
        FederationSqlPlan compiled = engine.compile(SqlCompileRequest.of(
            "SELECT ID FROM PERSON",
            SOURCE_ID,
            1
        ));
        FederationSqlPlan forged = (FederationSqlPlan) Proxy.newProxyInstance(
            getClass().getClassLoader(),
            new Class<?>[] {FederationSqlPlan.class},
            (proxy, method, arguments) -> {
                if ("executableSql".equals(method.getName())) {
                    return "DELETE FROM PERSON";
                }
                try {
                    return method.invoke(compiled, arguments);
                } catch (InvocationTargetException exception) {
                    throw exception.getCause();
                }
            }
        );

        try {
            engine.execute(forged, SqlExecutionContext.of(List.of()));
            Assert.fail("forged plan should be rejected");
        } catch (FederationSqlException exception) {
            Assert.assertEquals(
                exception.getMessage(),
                FederationSqlErrorCode.INVALID_ARGUMENT,
                exception.errorCode()
            );
        }
        Assert.assertEquals(0, dataSource.activeConnections());
    }

    /**
     * 验证查询范围未变化时，无关 Source 新增或 revision 更新都不会冲掉热计划。
     */
    @Test
    public void shouldKeepPlanWhenUnrelatedSourceChanges() {
        RuntimeFingerprint fingerprint = new RuntimeFingerprint("H2", "2", "H2 JDBC Driver", "2", "1");
        FederationSqlEngine cacheEngine = FederationSqlEngines.builder()
            .dataSourceResolver(definition -> FederationDataSourceHandles.shared(dataSource, fingerprint))
            .maximumPlanCacheEntries(32)
            .build();
        try {
            cacheEngine.sources().apply(definition(SOURCE_ID, 1), SourceApplyOptions.prewarmNow());
            SqlCompileRequest request = SqlCompileRequest.of("SELECT ID FROM PERSON", SOURCE_ID, 1);
            FederationSqlPlan initial = cacheEngine.compile(request);
            FederationSqlPlan initialHit = cacheEngine.compile(request);
            Assert.assertNotSame(initial, initialHit);
            Assert.assertSame(initial.relRoot(), initialHit.relRoot());

            SourceId other = new SourceId("other");
            cacheEngine.sources().apply(definition(other, 1));
            FederationSqlPlan membershipChanged = cacheEngine.compile(request);
            Assert.assertSame(initial.relRoot(), membershipChanged.relRoot());

            cacheEngine.sources().apply(definition(other, 2));
            FederationSqlPlan revisionChanged = cacheEngine.compile(request);
            Assert.assertSame(membershipChanged.relRoot(), revisionChanged.relRoot());
        } finally {
            cacheEngine.close();
        }
        Assert.assertEquals(0, dataSource.activeConnections());
    }

    /**
     * 验证取消已返回的游标会同步归还 JDBC 连接和 Core 查询资源。
     */
    @Test
    public void shouldReleaseReturnedCursorResourcesOnCancel() {
        FederationResultCursor cursor = engine.query(SqlQueryCommand.of(
            "SELECT ID, NAME FROM PERSON ORDER BY ID",
            SOURCE_ID,
            1,
            List.of()
        ));
        Assert.assertEquals(1, dataSource.activeConnections());
        Assert.assertTrue(engine.cancel(cursor.queryId()));
        Assert.assertTrue(cursor.metrics().cancelled());
        Assert.assertFalse(cursor.metrics().complete());
        long frozenExecutionNanos = cursor.metrics().executionNanos();
        long frozenFragmentNanos = cursor.metrics().fragments().get(0).elapsedNanos();
        Assert.assertEquals(-1, cursor.metrics().returnedBytes());
        java.util.concurrent.locks.LockSupport.parkNanos(
            java.util.concurrent.TimeUnit.MILLISECONDS.toNanos(2)
        );
        Assert.assertEquals(frozenExecutionNanos, cursor.metrics().executionNanos());
        Assert.assertEquals(
            frozenFragmentNanos,
            cursor.metrics().fragments().get(0).elapsedNanos()
        );
        Assert.assertEquals(0, dataSource.activeConnections());
        cursor.close();
    }

    /**
     * 验证 Statement 登记后、驱动真正执行前到达的取消会关闭 Statement 并阻止执行。
     *
     * @throws Exception 并发执行失败
     */
    @Test
    public void shouldCancelBetweenStatementRegistrationAndDriverExecution()
        throws Exception {
        FederationSqlPlan plan = engine.compile(SqlCompileRequest.of(
            "SELECT ID FROM PERSON ORDER BY ID",
            SOURCE_ID,
            1
        ));
        QueryId queryId = new QueryId("cancel-before-driver-execute");
        ExecutionGate gate = dataSource.gateNextExecution();
        ExecutorService executor = Executors.newSingleThreadExecutor();
        try {
            Future<FederationSqlErrorCode> query = executor.submit(() -> {
                try (FederationResultCursor ignored = engine.execute(
                    plan,
                    new SqlExecutionContext(
                        queryId,
                        List.of(),
                        SqlExecutionOptions.defaults(),
                        Duration.ofSeconds(5)
                    )
                )) {
                    return null;
                } catch (FederationSqlException exception) {
                    return exception.errorCode();
                }
            });

            Assert.assertTrue(gate.awaitExecution());
            Assert.assertTrue(engine.cancel(queryId));
            Assert.assertEquals(
                FederationSqlErrorCode.QUERY_CANCELLED,
                query.get(2, TimeUnit.SECONDS)
            );
            Assert.assertTrue(gate.statementClosed());
            Assert.assertEquals(0, dataSource.activeConnections());
        } finally {
            gate.release();
            executor.shutdownNow();
        }
    }

    /**
     * 验证 Scope 总时限可以关闭阻塞在 JDBC 执行边界的 Statement，并返回 QUERY_TIMEOUT。
     *
     * @throws Exception 并发执行失败
     */
    @Test
    public void shouldEnforceHardExecutionDeadlineAcrossBlockingJdbc() throws Exception {
        FederationQueryScopeDefinition deadlineScope = FederationQueryScopeDefinition.virtual(
            "deadline-scope",
            1,
            Map.of(
                "MAIN",
                FederationSourceBindingDefinition.of(SOURCE_ID, 1)
            ),
            "MAIN",
            new FederationExecutionPolicy(1, 1, 1, 1_000, 64_000, 100)
        );
        FederationSqlPlan plan = engine.compile(SqlCompileRequest.of(
            "SELECT ID FROM PERSON ORDER BY ID",
            deadlineScope
        ));
        ExecutionGate gate = dataSource.gateNextExecution();
        ExecutorService executor = Executors.newSingleThreadExecutor();
        long started = System.nanoTime();
        try {
            Future<FederationSqlErrorCode> query = executor.submit(() -> {
                try (FederationResultCursor ignored = engine.execute(
                    plan,
                    new SqlExecutionContext(
                        new QueryId("hard-deadline"),
                        List.of(),
                        SqlExecutionOptions.defaults(),
                        Duration.ofSeconds(5)
                    )
                )) {
                    return null;
                } catch (FederationSqlException exception) {
                    return exception.errorCode();
                }
            });

            Assert.assertTrue(gate.awaitExecution());
            Assert.assertEquals(
                FederationSqlErrorCode.QUERY_TIMEOUT,
                query.get(2, TimeUnit.SECONDS)
            );
            Assert.assertTrue(gate.statementClosed());
            Assert.assertTrue(
                TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - started) < 2_000
            );
            Assert.assertEquals(0, dataSource.activeConnections());
        } finally {
            gate.release();
            executor.shutdownNow();
        }
    }

    /**
     * 验证 QueryId 在准入等待前已经登记，取消可以及时中断默认本地准入。
     *
     * @throws Exception 并发等待失败
     */
    @Test
    public void shouldCancelQueryWaitingForAdmission() throws Exception {
        LocalFederationQueryAdmissionController delegate =
            new LocalFederationQueryAdmissionController(1);
        CountDownLatch secondAdmission = new CountDownLatch(1);
        AtomicInteger admissionCalls = new AtomicInteger();
        FederationQueryAdmissionController admission = new FederationQueryAdmissionController() {
            @Override
            public FederationQueryPermit acquire(
                SourceId sourceId,
                QueryId queryId,
                Duration timeout
            ) {
                return delegate.acquire(sourceId, queryId, timeout);
            }

            @Override
            public FederationQueryPermit acquire(
                SourceId sourceId,
                QueryId queryId,
                Duration timeout,
                BooleanSupplier cancellationRequested
            ) {
                if (admissionCalls.incrementAndGet() == 2) {
                    secondAdmission.countDown();
                }
                return delegate.acquire(sourceId, queryId, timeout, cancellationRequested);
            }

            @Override
            public void close() {
                delegate.close();
            }
        };
        RuntimeFingerprint fingerprint = new RuntimeFingerprint("H2", "2", "H2 JDBC Driver", "2", "1");
        FederationSqlEngine admissionEngine = FederationSqlEngines.builder()
            .dataSourceResolver(definition -> FederationDataSourceHandles.shared(dataSource, fingerprint))
            .admissionController(admission)
            .build();
        ExecutorService executor = Executors.newSingleThreadExecutor();
        FederationResultCursor first = null;
        try {
            admissionEngine.sources().apply(definition(1), SourceApplyOptions.prewarmNow());
            first = admissionEngine.query(SqlQueryCommand.of(
                "SELECT ID FROM PERSON ORDER BY ID",
                SOURCE_ID,
                1,
                List.of()
            ));
            Assert.assertEquals(1, dataSource.activeConnections());

            QueryId waitingId = new QueryId("waiting-for-admission");
            SqlQueryCommand waitingCommand = new SqlQueryCommand(
                waitingId,
                "SELECT ID FROM PERSON ORDER BY ID",
                SOURCE_ID,
                1,
                List.of(),
                SqlExecutionOptions.defaults(),
                Duration.ofSeconds(5).toMillis(),
                "default"
            );
            Future<FederationSqlErrorCode> waiting = executor.submit(() -> {
                try (FederationResultCursor ignored = admissionEngine.query(waitingCommand)) {
                    return null;
                } catch (FederationSqlException exception) {
                    return exception.errorCode();
                }
            });
            Assert.assertTrue(secondAdmission.await(2, TimeUnit.SECONDS));
            Assert.assertTrue(admissionEngine.cancel(waitingId));
            Assert.assertEquals(
                FederationSqlErrorCode.QUERY_CANCELLED,
                waiting.get(2, TimeUnit.SECONDS)
            );
            Assert.assertEquals(1, dataSource.activeConnections());
        } finally {
            if (first != null) {
                first.close();
            }
            executor.shutdownNow();
            admissionEngine.close();
        }
        Assert.assertEquals(0, dataSource.activeConnections());
    }

    /**
     * 验证结果消费超时具有稳定错误码，并关闭全部 JDBC 资源。
     */
    @Test
    public void shouldClassifyResultTimeoutAndCloseResources() {
        AtomicBoolean resultSetClosed = new AtomicBoolean();
        AtomicBoolean statementClosed = new AtomicBoolean();
        AtomicBoolean connectionClosed = new AtomicBoolean();
        AtomicBoolean unregistered = new AtomicBoolean();
        ResultSet resultSet = (ResultSet) Proxy.newProxyInstance(
            getClass().getClassLoader(),
            new Class<?>[] {ResultSet.class},
            (proxy, method, arguments) -> {
                if ("next".equals(method.getName())) {
                    throw new SQLTimeoutException("simulated read timeout");
                }
                if ("close".equals(method.getName())) {
                    resultSetClosed.set(true);
                }
                return null;
            }
        );
        PreparedStatement statement = (PreparedStatement) Proxy.newProxyInstance(
            getClass().getClassLoader(),
            new Class<?>[] {PreparedStatement.class},
            (proxy, method, arguments) -> {
                if ("close".equals(method.getName())) {
                    statementClosed.set(true);
                }
                return null;
            }
        );
        Connection connection = (Connection) Proxy.newProxyInstance(
            getClass().getClassLoader(),
            new Class<?>[] {Connection.class},
            (proxy, method, arguments) -> {
                if ("close".equals(method.getName())) {
                    connectionClosed.set(true);
                }
                return null;
            }
        );
        StatementLifecycle lifecycle = new StatementLifecycle() {
            @Override
            public void register(Statement registeredStatement) {
            }

            @Override
            public void unregister(Statement registeredStatement) {
                unregistered.set(true);
            }
        };
        JdbcFederationResultCursor cursor = new JdbcFederationResultCursor(
            new QueryId("timeout-query"),
            List.of(new FederationColumn(1, "ID", Types.INTEGER, "INTEGER", false)),
            resultSet,
            statement,
            connection,
            lifecycle
        );

        try {
            cursor.next();
            Assert.fail("timeout should be reported");
        } catch (FederationSqlException exception) {
            Assert.assertEquals(FederationSqlErrorCode.QUERY_TIMEOUT, exception.errorCode());
        }
        Assert.assertTrue(resultSetClosed.get());
        Assert.assertTrue(statementClosed.get());
        Assert.assertTrue(connectionClosed.get());
        Assert.assertTrue(unregistered.get());
    }

    /**
     * 验证 SELECT 列中的 schema.table.column 不会因 SourceId 同名而被误判为跨源。
     */
    @Test
    public void shouldDetectCrossSourceOnlyFromRelationReferences() {
        RuntimeFingerprint fingerprint = new RuntimeFingerprint("H2", "2", "H2 JDBC Driver", "2", "1");
        FederationSqlEngine localEngine = FederationSqlEngines.builder()
            .dataSourceResolver(definition -> FederationDataSourceHandles.shared(dataSource, fingerprint))
            .crossSourceEnabled(false)
            .build();
        try {
            localEngine.sources().apply(new FederationSourceDefinition(
                SOURCE_ID,
                1,
                JdbcFederationSqlAdapterProvider.ADAPTER_ID,
                List.of(
                    new JdbcSchemaDefinition("APP", null, "PUBLIC"),
                    new JdbcSchemaDefinition("AUX", null, "PUBLIC")
                ),
                Map.of()
            ), SourceApplyOptions.prewarmNow());
            localEngine.sources().apply(definition(new SourceId("APP"), 1));

            FederationSqlPlan plan = localEngine.compile(SqlCompileRequest.of(
                "SELECT APP.PERSON.ID FROM APP.PERSON",
                SOURCE_ID,
                1
            ));
            Assert.assertTrue(plan.executable());
            Assert.assertEquals(Set.of(SOURCE_ID), plan.referencedSources());
            try {
                FederationQueryScopeDefinition scope = new FederationQueryScopeDefinition(
                    "cross-source-detection",
                    1,
                    Map.of(
                        "MAIN",
                        FederationSourceBindingDefinition.of(SOURCE_ID, 1),
                        "REMOTE",
                        FederationSourceBindingDefinition.of(new SourceId("APP"), 1)
                    ),
                    "MAIN",
                    FederationExecutionPolicy.basic()
                );
                localEngine.compile(SqlCompileRequest.of(
                    "SELECT p.ID FROM MAIN.APP.PERSON p "
                        + "JOIN REMOTE.APP.PERSON r ON p.ID = r.ID",
                    scope
                ));
                Assert.fail("qualified relation should be detected as cross-source");
            } catch (FederationSqlException exception) {
                Assert.assertEquals(FederationSqlErrorCode.CROSS_SOURCE_DISABLED, exception.errorCode());
            }
        } finally {
            localEngine.close();
        }
    }

    /**
     * 断言策略拒绝使用稳定的参数错误码返回。
     *
     * @param action 待执行动作
     */
    private static void assertPolicyDenied(Runnable action) {
        try {
            action.run();
            Assert.fail("policy should reject the query");
        } catch (FederationSqlException exception) {
            Assert.assertEquals(FederationSqlErrorCode.INVALID_ARGUMENT, exception.errorCode());
        }
    }

    private static FederationSourceDefinition definition(long revision) {
        return definition(SOURCE_ID, revision);
    }

    private static FederationSourceDefinition definition(SourceId sourceId, long revision) {
        return new FederationSourceDefinition(
            sourceId,
            revision,
            JdbcFederationSqlAdapterProvider.ADAPTER_ID,
            List.of(new JdbcSchemaDefinition("APP", null, "PUBLIC")),
            Map.of()
        );
    }

    /**
     * 统计借出连接且保持 DataSource 其余行为透明的测试包装器。
     */
    private static final class TrackingDataSource implements DataSource {

        private final DataSource delegate;
        private final AtomicInteger activeConnections = new AtomicInteger();
        private final AtomicReference<ExecutionGate> nextExecutionGate =
            new AtomicReference<>();

        private TrackingDataSource(DataSource delegate) {
            this.delegate = delegate;
        }

        private int activeConnections() {
            return activeConnections.get();
        }

        /**
         * 为下一条 PreparedStatement 安装执行门闩。
         *
         * @return 执行门闩
         */
        private ExecutionGate gateNextExecution() {
            ExecutionGate gate = new ExecutionGate();
            if (!nextExecutionGate.compareAndSet(null, gate)) {
                throw new IllegalStateException("an execution gate is already pending");
            }
            return gate;
        }

        @Override
        public Connection getConnection() throws SQLException {
            return track(delegate.getConnection());
        }

        @Override
        public Connection getConnection(String username, String password) throws SQLException {
            return track(delegate.getConnection(username, password));
        }

        private Connection track(Connection connection) {
            activeConnections.incrementAndGet();
            AtomicBoolean closed = new AtomicBoolean();
            return (Connection) Proxy.newProxyInstance(
                getClass().getClassLoader(),
                new Class<?>[] {Connection.class},
                (proxy, method, arguments) -> {
                    if ("close".equals(method.getName())) {
                        if (closed.compareAndSet(false, true)) {
                            try {
                                connection.close();
                            } finally {
                                activeConnections.decrementAndGet();
                            }
                        }
                        return null;
                    }
                    try {
                        Object result = method.invoke(connection, arguments);
                        if (result instanceof PreparedStatement preparedStatement) {
                            ExecutionGate gate = nextExecutionGate.getAndSet(null);
                            return gate == null
                                ? preparedStatement
                                : gate.wrap(preparedStatement);
                        }
                        return result;
                    } catch (InvocationTargetException exception) {
                        throw exception.getCause();
                    }
                }
            );
        }

        @Override
        public PrintWriter getLogWriter() throws SQLException {
            return delegate.getLogWriter();
        }

        @Override
        public void setLogWriter(PrintWriter out) throws SQLException {
            delegate.setLogWriter(out);
        }

        @Override
        public void setLoginTimeout(int seconds) throws SQLException {
            delegate.setLoginTimeout(seconds);
        }

        @Override
        public int getLoginTimeout() throws SQLException {
            return delegate.getLoginTimeout();
        }

        @Override
        public Logger getParentLogger() {
            return Logger.getGlobal();
        }

        @Override
        public <T> T unwrap(Class<T> iface) throws SQLException {
            return delegate.unwrap(iface);
        }

        @Override
        public boolean isWrapperFor(Class<?> iface) throws SQLException {
            return delegate.isWrapperFor(iface);
        }
    }

    /**
     * 将 PreparedStatement 暂停在驱动执行边界，用于复现取消空窗。
     */
    private static final class ExecutionGate {

        private final CountDownLatch executionEntered = new CountDownLatch(1);
        private final CountDownLatch executionReleased = new CountDownLatch(1);
        private final AtomicBoolean statementClosed = new AtomicBoolean();

        /**
         * 包装待执行 Statement。
         *
         * @param delegate 原始 Statement
         * @return 受门闩控制的 Statement
         */
        private PreparedStatement wrap(PreparedStatement delegate) {
            return (PreparedStatement) Proxy.newProxyInstance(
                getClass().getClassLoader(),
                new Class<?>[] {PreparedStatement.class},
                (proxy, method, arguments) -> {
                    if ("executeQuery".equals(method.getName())) {
                        executionEntered.countDown();
                        try {
                            if (!executionReleased.await(2, TimeUnit.SECONDS)) {
                                throw new SQLException("execution gate timed out");
                            }
                        } catch (InterruptedException exception) {
                            Thread.currentThread().interrupt();
                            throw new SQLException("execution gate interrupted", exception);
                        }
                        if (statementClosed.get()) {
                            throw new SQLException("statement closed before execution");
                        }
                    }
                    if ("close".equals(method.getName())) {
                        statementClosed.set(true);
                        executionReleased.countDown();
                    }
                    try {
                        return method.invoke(delegate, arguments);
                    } catch (InvocationTargetException exception) {
                        throw exception.getCause();
                    }
                }
            );
        }

        /**
         * 等待查询抵达驱动执行边界。
         *
         * @return 是否在超时前到达
         * @throws InterruptedException 等待被中断
         */
        private boolean awaitExecution() throws InterruptedException {
            return executionEntered.await(2, TimeUnit.SECONDS);
        }

        /**
         * 返回 Statement 是否已由取消路径关闭。
         *
         * @return 是否关闭
         */
        private boolean statementClosed() {
            return statementClosed.get();
        }

        /**
         * 释放驱动执行门闩。
         */
        private void release() {
            executionReleased.countDown();
        }
    }
}
