# Premium Import Into Async File Discovery Implementation Plan

## Analysis

### Impact Table

| Area | File | Change | Verification | Success Criteria |
|------|------|--------|--------------|------------------|
| NG agent async contract for file discovery | `dataflow-agent-ng/idl/idl/agent/ng_agent/agent.proto` | Add async create/get RPCs and response messages for import file discovery, mirroring the existing estimate-data-size task pattern instead of relying on a long-running synchronous `GetImportFiles` call | `make generate-idl` | Generated Go/OpenAPI artifacts contain the new async file-discovery RPCs and messages without manual edits |
| NG agent async execution | `dataflow-agent-ng/agent/grpc/services.go`, `dataflow-agent-ng/agent/grpc/import_task.go` | Refactor current `GetImportFiles` scan logic into a reusable helper and run it through `importTaskManager` with deterministic task IDs and task retention suitable for Premium polling | `go test ./dataflow-agent-ng/agent/grpc -run '^(TestGetImportFiles|TestCreateImportFilesTask|TestGetImportFilesTask)'` | File discovery can be started once, polled repeatedly, and returns the same result/error until completion |
| NG regional client wrapper | `internal/handler/agent/regional-ng.go`, `internal/handler/agent/regional_ng_test.go` | Add `CreateImportFilesTask` / `GetImportFilesTask` client helpers so dataflow-service can use the async agent flow | `go test ./internal/handler/agent -run '^(TestNGRegionalClientCreateImportDataSizeEstimateTaskIncludesTableMeta|TestNGRegionalClientCreateImportFilesTask|TestNGRegionalClientGetImportFilesTask)$'` | Client methods hit the expected endpoints and preserve table meta payloads |
| Premium import-into worker prepare stage | `internal/model/import/import_into.go`, `internal/model/import/import_into_prepare_test.go` | Split "file discovery prepared" from "size estimated", poll async file discovery before estimate-size, and make `Preparing` re-entrant without using partial subtask metadata as a false-ready signal | `go test ./internal/model/import -run '^(TestApplyImportIntoEstimatedSourceSize|TestValidateImportIntoEstimateResult|TestCollectImportIntoEstimateTableMetas|TestImportIntoFileDiscoveryStateHelpers|TestBuildImportIntoGetFilesTaskID)$'` | A Premium import can stay in `preparing` while file discovery is still running, then move to estimate-size only after complete file metadata is available |
| Premium import source check API | `idl/idl/dataflow/adminapi/v1beta2/import_into.proto`, `internal/openapi/service/importadmin/import_into.go`, `internal/openapi/service/importadmin/import_into_test.go` | Change `CheckImportSource` from a timeout-avoidance sync scan with `2000` file cap into a pollable async contract backed by the new agent task flow | `make generate-idl && go test ./internal/openapi/service/importadmin -run '^(TestCheckImportSource|TestFromSourceURI_Azure_StripsSecrets|TestNormalizeAzureBlobInput|TestJoinSourceURI_PreservesQuery)$'` | Premium callers can poll the same API until file discovery completes; the scan no longer depends on the `2000`-file cap for timeout protection |

### Dead Code Cascade

- The synchronous-only `CheckImportSource -> agentCli.GetImportFiles(..., 2000, ...)` timeout workaround in `internal/openapi/service/importadmin/import_into.go` should be removed once async polling is in place.
- Any one-off status mapping that only exists to describe "truncated because the sync call would time out" should be removed or reduced to true product truncation semantics. If product still needs truncation for UI reasons, keep it explicit and separate from timeout handling.
- The current `isImportIntoPrepared()` helper in `internal/model/import/import_into.go` is too coarse for a two-phase async prepare flow because any partial subtask metadata makes the whole import look prepared. It should be replaced by narrower predicates rather than extended.

### Symbols To Keep / Introduce / Remove

Keep:

- `buildImportIntoEstimateTaskID`
- `collectImportIntoEstimateTableMetas`
- `applyImportIntoEstimatedSourceSize`
- `importAsyncTaskManager`

Introduce:

- `CreateImportFilesTask`
- `GetImportFilesTask`
- `buildImportIntoGetFilesTaskID`
- `buildCheckImportSourceTaskID`
- `populateImportIntoSubTasksFromFileDiscovery`
- `isImportIntoFileDiscoveryPrepared`
- `ImportFilesTaskStatus` or equivalent proto enum/message names matching the estimate task style

Remove or replace:

- `isImportIntoPrepared`
- The hard-coded `2000` max-scan argument in Premium `CheckImportSource`
- Any "truncated because sync may time out" branch that becomes redundant after async polling

### Scope Notes

- `CreateImportIntoTask` is already async at the service boundary because it only persists the import and lets `ImportIntoWorker` drive `Preparing`; this plan does not add another worker.
- The real synchronous timeout boundary today is `CheckImportSource`, while the worker-side `onImportIntoPreparing` still performs a synchronous file scan inside the background loop. Both call sites should share the same async file-discovery primitive so behavior stays consistent.
- The estimate-size async flow in `internal/model/import/import_into.go` is the canonical pattern to follow. The new file-discovery task should mirror its task lifecycle closely enough that future maintenance stays obvious.
- Future Premium import phases such as `CreateSchemas` and import execution should follow the same orchestration model, but they are explicitly out of scope for this change. This plan only makes `getImportFiles` async and reshapes the prepare flow so later async steps can be inserted cleanly without another state-machine rewrite.
- No schema migration is planned by default. Task IDs can be deterministic and derivable from immutable request/import inputs, so the worker and API can recreate/poll tasks without persisting a new database column. If implementation proves that task retention or collision handling requires durable storage, reevaluate before coding.

### Unified Async Direction

To keep the overall Premium import code readable as more steps move from sync to async, treat `Preparing` as an orchestration stage composed of ordered async sub-steps instead of open-coded one-off waits:

1. `getImportFiles`
2. `estimate size`
3. `create schema`
4. `import into` submission

The design direction should be:

- One orchestration owner: `internal/model/import/import_into.go` remains the only place that decides which step is next for a Premium import task.
- One async pattern: every long-running remote action should expose `CreateXTask` + `GetXTask` semantics, regardless of whether the caller is OpenAPI (`CheckImportSource`) or the worker (`Preparing` / future `ScalingUp`).
- One state-machine vocabulary: each async step should use the same three responsibilities and names in worker code:
  - `resolveXxxTaskState`: only decide `succeeded / wait / terminal failure`
  - `buildXxxFailureMessage`: only build the persisted terminal message
  - `applyXxxTaskResult`: only apply the succeeded payload into business state
- One persistence boundary: DB state stores business progress and final artifacts, not opaque transport details. The worker should be able to recompute or recreate remote async task IDs deterministically.
- One readiness rule per phase: do not use broad helpers like `isImportIntoPrepared()` that mix several milestones together. Each async sub-step should have its own "completed" predicate.
- One insertion path for future async work: after this change, adding async `create schema` should mean inserting another sub-step after estimate-size, not duplicating polling logic in a different status handler.

For this PR, only step `1. getImportFiles` is being converted. The code structure should still be updated so steps `3` and `4` can adopt the same pattern later. Follow this rollout order when touching the worker:

1. Unify naming and responsibility boundaries for `resolve / build / apply`.
2. Move error-message construction out of the orchestration `switch` blocks into dedicated helpers, including estimate-size task handling.
3. Make each async polling segment explicitly follow `succeeded / wait / terminal failure`.
4. Only after those boundaries are stable, consider regrouping `Preparing` into clearer sections or files.

## Execution Steps

1. Define the async file-discovery agent API.
   - Add `CreateImportFilesTask` and `GetImportFilesTask` to `dataflow-agent-ng/idl/idl/agent/ng_agent/agent.proto`.
   - Reuse the existing `GetImportFilesResp.TableMeta` as the completed result payload instead of inventing a second table-meta shape.
   - Add a task-status enum parallel to `ImportDataSizeEstimateTaskStatus`, and decide whether polling errors stay in `error_msg` only or also need a top-level failure status; prefer matching the estimate-size contract exactly.
   - Verification:
     - `make generate-idl`
     - Expected: generated pb/openapi files update cleanly and no manual edits are needed under generated paths.

2. Refactor NG agent file-scan logic behind a shared helper and async task wrapper.
   - Extract the body of `(*AgentServer).GetImportFiles` into a pure helper such as `getImportFiles(ctx, req)`.
   - Keep the existing synchronous RPC for compatibility, but make it call the shared helper so sync and async paths cannot drift.
   - Implement `CreateImportFilesTask` with `importTaskManager.CreateIfAbsent(...)`.
   - Implement `GetImportFilesTask` by reading the same task manager and returning a typed result when complete.
   - Review `import_task.go` retention. If the current `10m` TTL is too short for Premium polling or large-bucket scans, increase it or make it configurable before wiring callers onto it.
   - Verification:
     - `go test ./dataflow-agent-ng/agent/grpc -run '^(TestGetImportFiles_WithTableMeta_SchemaExists|TestGetImportFiles_WithTableMeta_SchemaNotFound|TestGetImportFiles_WithTableMeta_MultipleTables|TestCreateImportFilesTask|TestGetImportFilesTask|TestImportAsyncTaskManager)'`
     - Expected: existing synchronous tests still pass; new async tests cover running, succeeded, failed, and duplicate-create behavior.

3. Add dataflow-service NG client wrappers for the new agent endpoints.
   - Extend `internal/handler/agent/regional-ng.go` with `CreateImportFilesTask(...)` and `GetImportFilesTask(...)`.
   - Keep the method signatures parallel to the estimate-size wrappers so the caller code reads the same way.
   - Add HTTP-path tests in `internal/handler/agent/regional_ng_test.go` for create/get plus table-meta serialization.
   - Verification:
     - `go test ./internal/handler/agent -run '^(TestNGRegionalClientCreateImportDataSizeEstimateTaskIncludesTableMeta|TestNGRegionalClientCreateImportFilesTask|TestNGRegionalClientGetImportFilesTask)$'`
     - Expected: tests assert the correct request path, payload, and decoded response fields.

4. Make Premium import-into `Preparing` explicitly two-phase and re-entrant.
  - Replace `isImportIntoPrepared()` with separate helpers:
    - one for "file discovery finished for all relevant subtasks"
    - one for "estimate-size already applied"
    - keep the helper naming neutral enough that a later "schema created" predicate can be added without another semantic reset
  - Align each async sub-step with the shared helper vocabulary:
    - `resolveXxxTaskState`
    - `buildXxxFailureMessage`
    - `applyXxxTaskResult`
  - Add deterministic task ID creation for worker file discovery, for example `import-into-get-files-<importID>`.
  - In `onImportIntoPreparing`:
    - create/poll async file discovery first
     - if the task is still running, persist no false-ready metadata and return `nil` so the worker retries on the next tick
     - on success, populate subtask file info and only then start or poll the existing estimate-size task
     - on failure, mark the import failed with a sanitized message consistent with the current import boundary
     - keep the step dispatch explicit, so a future async `CreateSchemas` step can be inserted after estimate-size without collapsing `Preparing` back into nested conditionals
   - Make sure custom-table and auto-detect flows both use the same completion helper so they cannot diverge semantically.
   - Verification:
     - `go test ./internal/model/import -run '^(TestApplyImportIntoEstimatedSourceSize|TestValidateImportIntoEstimateResult|TestCollectImportIntoEstimateTableMetas|TestImportIntoFileDiscoveryStateHelpers|TestBuildImportIntoGetFilesTaskID)$'`
     - Expected: helper tests prove partial subtask metadata no longer marks the whole import as prepared; deterministic task IDs stay stable across retries.

5. Convert Premium `CheckImportSource` to async polling semantics.
   - Decide the minimal user-facing contract:
     - preferred: keep the same endpoint, add `ENUM_RUNNING` to `CheckImportSourceStatus`, derive a deterministic task ID from `(clusterID, normalized source URI, normalized table metas)`, and let repeated identical requests poll the same underlying agent task
     - only add a new response/request task-id field if deterministic polling by identical request is insufficient
   - Remove the `2000` file cap from the async path; timeout protection should come from polling, not truncation.
   - Preserve existing normalization rules for S3/OSS/Azure URIs and wildcard joins before hashing task identity, otherwise repeated polls may miss the same agent task.
   - Keep `table_metas`, `total_scanned_objects`, and final `TRUNCATED` behavior only if product still wants truncation after a full async scan; otherwise return `SUCCESS` on full completion.
   - Verification:
     - `make generate-idl`
     - `go test ./internal/openapi/service/importadmin -run '^(TestCheckImportSourceRunning|TestCheckImportSourceCompleted|TestCheckImportSourceCustomTables|TestFromSourceURI_Azure_StripsSecrets|TestNormalizeAzureBlobInput|TestJoinSourceURI_PreservesQuery)$'`
     - Expected: first call returns running status without blocking on full scan; subsequent poll returns completed metas from the same task.

6. Run focused end-to-end verification for the full Premium flow.
   - Verify NG agent async file discovery in isolation.
   - Verify Premium service polling behavior.
   - Verify import-into prepare loop integration.
   - Recommended command set:
     - `go test ./dataflow-agent-ng/agent/grpc`
     - `go test ./internal/handler/agent`
     - `go test ./internal/model/import`
     - `go test ./internal/openapi/service/importadmin`
     - `make app-unit-test-local`
   - Expected:
     - all focused packages pass
     - `make app-unit-test-local` completes without new regressions in Premium import paths

## Risks

- Async polling without a persisted task ID depends on deterministic request normalization. If the service produces different normalized URIs across repeated requests, callers will create duplicate scans instead of polling.
- `importAsyncTaskManager` is in-memory. Agent restarts will lose running/completed file-discovery tasks, so callers must tolerate "task not found" by recreating the task idempotently.
- Extending `CheckImportSourceStatus` changes the Premium OpenAPI contract. FE or external callers that assume only `SUCCESS`/`TRUNCATED` may need coordinated rollout.
- Large completed results may now be returned after async completion instead of being hidden behind a sync truncation cap. Confirm whether FE can handle the full table/file metadata volume before removing truncation semantics entirely.
- If this PR only adds async `getImportFiles` but leaves the worker code structured around ad hoc milestone checks, the next async conversion (`CreateSchemas` or import submission) will reintroduce branching complexity quickly. The implementation should pay down that structure now even though only one sub-step changes behavior.