package duckdb import ( "context" "database/sql" "os" "testing" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) type SlowlogInDL struct { ID int64 `db:"id"` Time float64 `db:"time"` Digest string `db:"digest"` Date string `db:"date"` // partition column Instance sql.NullString `db:"instance"` TxnStartTS sql.NullString `db:"txn_start_ts"` // uint64 as string User sql.NullString `db:"user"` Host sql.NullString `db:"host"` ConnID sql.NullString `db:"conn_id"` // uint64 as string SessionAlias sql.NullString `db:"session_alias"` ExecRetryCount sql.NullString `db:"exec_retry_count"` // uint64 as string ExecRetryTime sql.NullFloat64 `db:"exec_retry_time"` QueryTime sql.NullFloat64 `db:"query_time"` ParseTime sql.NullFloat64 `db:"parse_time"` CompileTime sql.NullFloat64 `db:"compile_time"` RewriteTime sql.NullFloat64 `db:"rewrite_time"` PreprocSubqueries sql.NullString `db:"preproc_subqueries"` PreprocSubqueriesTime sql.NullFloat64 `db:"preproc_subqueries_time"` OptimizeTime sql.NullFloat64 `db:"optimize_time"` WaitTS sql.NullFloat64 `db:"wait_ts"` PrewriteTime sql.NullFloat64 `db:"prewrite_time"` WaitPrewriteBinlogTime sql.NullFloat64 `db:"wait_prewrite_binlog_time"` CommitTime sql.NullFloat64 `db:"commit_time"` GetCommitTSTime sql.NullFloat64 `db:"get_commit_ts_time"` CommitBackoffTime sql.NullFloat64 `db:"commit_backoff_time"` BackoffTypes sql.NullString `db:"backoff_types"` ResolveLockTime sql.NullFloat64 `db:"resolve_lock_time"` LocalLatchWaitTime sql.NullFloat64 `db:"local_latch_wait_time"` WriteKeys sql.NullInt64 `db:"write_keys"` WriteSize sql.NullInt64 `db:"write_size"` PrewriteRegion sql.NullInt64 `db:"prewrite_region"` TxnRetry sql.NullInt64 `db:"txn_retry"` CopTime sql.NullFloat64 `db:"cop_time"` ProcessTime sql.NullFloat64 `db:"process_time"` WaitTime sql.NullFloat64 `db:"wait_time"` BackoffTime sql.NullFloat64 `db:"backoff_time"` LockkeysTime sql.NullFloat64 `db:"lockkeys_time"` RequestCount sql.NullString `db:"request_count"` // uint64 as string TotalKeys sql.NullString `db:"total_keys"` // uint64 as string ProcessKeys sql.NullString `db:"process_keys"` // uint64 as string RocksdbDeleteSkippedCount sql.NullString `db:"rocksdb_delete_skipped_count"` // uint64 as string RocksdbKeySkippedCount sql.NullString `db:"rocksdb_key_skipped_count"` // uint64 as string RocksdbBlockCacheHitCount sql.NullString `db:"rocksdb_block_cache_hit_count"` // uint64 as string RocksdbBlockReadCount sql.NullString `db:"rocksdb_block_read_count"` // uint64 as string RocksdbBlockReadByte sql.NullString `db:"rocksdb_block_read_byte"` // uint64 as string DB sql.NullString `db:"db"` IndexNames sql.NullString `db:"index_names"` IsInternal sql.NullBool `db:"is_internal"` Stats sql.NullString `db:"stats"` CopProcAvg sql.NullFloat64 `db:"cop_proc_avg"` CopProcP90 sql.NullFloat64 `db:"cop_proc_p90"` CopProcMax sql.NullFloat64 `db:"cop_proc_max"` CopProcAddr sql.NullString `db:"cop_proc_addr"` CopWaitAvg sql.NullFloat64 `db:"cop_wait_avg"` CopWaitP90 sql.NullFloat64 `db:"cop_wait_p90"` CopWaitMax sql.NullFloat64 `db:"cop_wait_max"` CopWaitAddr sql.NullString `db:"cop_wait_addr"` MemMax sql.NullInt64 `db:"mem_max"` DiskMax sql.NullInt64 `db:"disk_max"` KvTotal sql.NullFloat64 `db:"kv_total"` PdTotal sql.NullFloat64 `db:"pd_total"` BackoffTotal sql.NullFloat64 `db:"backoff_total"` WriteSQLResponseTotal sql.NullFloat64 `db:"write_sql_response_total"` ResultRows sql.NullInt64 `db:"result_rows"` Warnings sql.NullString `db:"warnings"` BackoffDetail sql.NullString `db:"backoff_detail"` Prepared sql.NullBool `db:"prepared"` Succ sql.NullBool `db:"succ"` IsExplicitTxn sql.NullBool `db:"is_explicit_txn"` IsWriteCacheTable sql.NullBool `db:"is_write_cache_table"` PlanFromCache sql.NullBool `db:"plan_from_cache"` PlanFromBinding sql.NullBool `db:"plan_from_binding"` HasMoreResults sql.NullBool `db:"has_more_results"` ResourceGroup sql.NullString `db:"resource_group"` RequestUnitRead sql.NullFloat64 `db:"request_unit_read"` RequestUnitWrite sql.NullFloat64 `db:"request_unit_write"` TimeQueuedByRC sql.NullFloat64 `db:"time_queued_by_rc"` PlanDigest sql.NullString `db:"plan_digest"` PrevStmt sql.NullString `db:"prev_stmt"` ServerlessTenantID sql.NullString `db:"serverless_tenant_id"` ServerlessProjectID sql.NullString `db:"serverless_project_id"` ServerlessClusterID sql.NullString `db:"serverless_cluster_id"` WorkerRole sql.NullString `db:"worker_role"` IsTiflash sql.NullBool `db:"is_tiflash"` RU sql.NullFloat64 `db:"ru"` Query sql.NullString `db:"query"` BinaryPlan sql.NullString `db:"binary_plan"` Plan sql.NullString `db:"plan"` DecodedPlan sql.NullString `db:"decoded_plan"` } var FieldCastTypeMap = map[string]string{ "txn_start_ts": "UBIGINT", "conn_id": "UBIGINT", "exec_retry_count": "UBIGINT", "preproc_subqueries": "UBIGINT", "write_keys": "UBIGINT", "write_size": "UBIGINT", "prewrite_region": "UBIGINT", "txn_retry": "UBIGINT", "request_count": "UBIGINT", "total_keys": "UBIGINT", "process_keys": "UBIGINT", "rocksdb_delete_skipped_count": "UBIGINT", "rocksdb_key_skipped_count": "UBIGINT", "rocksdb_block_cache_hit_count": "UBIGINT", "rocksdb_block_read_count": "UBIGINT", "rocksdb_block_read_byte": "UBIGINT", "mem_max": "UBIGINT", "disk_max": "UBIGINT", "result_rows": "UBIGINT", } func TestNewQueryDuckDBFetcher_InitSuccess(t *testing.T) { os.Setenv("AWS_ACCESS_KEY_ID", "test_access_key") os.Setenv("AWS_SECRET_ACCESS_KEY", "test_secret_key") defer func() { os.Unsetenv("AWS_ACCESS_KEY_ID") os.Unsetenv("AWS_SECRET_ACCESS_KEY") }() ctx := context.Background() fetcher, err := NewQueryDuckDBFetcher(ctx, "aws", "us-east-1") require.NoError(t, err) assert.NotNil(t, fetcher) assert.NotNil(t, fetcher.db) assert.NotNil(t, fetcher.dbx) assert.Equal(t, "aws", fetcher.vendor) assert.Equal(t, "us-east-1", fetcher.region) err = fetcher.Close(ctx) assert.NoError(t, err) } func TestNewQueryDuckDBFetcher_InitSuccess_GCP(t *testing.T) { ctx := context.Background() fetcher, err := NewQueryDuckDBFetcher(ctx, "gcp", "us-central1") require.NoError(t, err) assert.NotNil(t, fetcher) assert.Equal(t, "gcp", fetcher.vendor) err = fetcher.Close(ctx) assert.NoError(t, err) } func TestNewQueryDuckDBFetcher_InitSuccess_Azure(t *testing.T) { ctx := context.Background() fetcher, err := NewQueryDuckDBFetcher(ctx, "azure", "eastus") require.NoError(t, err) assert.NotNil(t, fetcher) assert.Equal(t, "azure", fetcher.vendor) err = fetcher.Close(ctx) assert.NoError(t, err) } func TestNewQueryDuckDBFetcher_InitError_UnsupportedVendor(t *testing.T) { ctx := context.Background() fetcher, err := NewQueryDuckDBFetcher(ctx, "unsupported", "us-east-1") require.Error(t, err) assert.Contains(t, err.Error(), "unsupported vendor: unsupported") assert.Nil(t, fetcher) } func TestNewQueryDuckDBFetcher_InitMemoryLimits(t *testing.T) { tests := []struct { name string podMemLimit string duckdbMemPercent string expectedMinMemoryMB int expectError bool errorContains string }{ {"Valid config", "4096", "75", 3072, false, ""}, {"Low pod limit, default duckdb percent", "1000", "", 2048, false, ""}, {"Invalid pod limit", "abc", "80", 2048, false, ""}, // Uses default, then adjusts to min {"Invalid duckdb percent", "4096", "abc", 3276, false, ""}, // Uses default 80% {"Duckdb percent too high", "4096", "110", 3276, false, ""}, // Uses default 80% {"Duckdb percent too low", "4096", "0", 3276, false, ""}, // Uses default 80% {"Calculated memory below minimum", "2000", "50", 2048, false, ""}, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { os.Setenv("POD_MEMORY_LIMITS_MB", tt.podMemLimit) os.Setenv("DUCKDB_MEMORY_PERCENTAGE", tt.duckdbMemPercent) defer func() { os.Unsetenv("POD_MEMORY_LIMITS_MB") os.Unsetenv("DUCKDB_MEMORY_PERCENTAGE") }() ctx := context.Background() fetcher, err := NewQueryDuckDBFetcher(ctx, "aws", "us-east-1") if tt.expectError { require.Error(t, err) assert.Contains(t, err.Error(), tt.errorContains) assert.Nil(t, fetcher) } else { require.NoError(t, err) assert.NotNil(t, fetcher) // We can't directly read the memory_limit from DuckDB via its API // but we can assert that initialization succeeded. // For more precise testing of SET commands, mocking would be needed. err = fetcher.Close(ctx) assert.NoError(t, err) } }) } } func TestSlowQueryDuckDBFetcher_Close(t *testing.T) { os.Setenv("AWS_ACCESS_KEY_ID", "test_access_key") os.Setenv("AWS_SECRET_ACCESS_KEY", "test_secret_key") defer func() { os.Unsetenv("AWS_ACCESS_KEY_ID") os.Unsetenv("AWS_SECRET_ACCESS_KEY") }() ctx := context.Background() fetcher, err := NewQueryDuckDBFetcher(ctx, "aws", "us-east-1") require.NoError(t, err) require.NotNil(t, fetcher) err = fetcher.Close(ctx) assert.NoError(t, err) err = fetcher.Close(ctx) assert.NoError(t, err) } func TestSlowQueryDuckDBFetcher_SelectSlowlogInDL(t *testing.T) { os.Setenv("AWS_ACCESS_KEY_ID", "test_access_key") os.Setenv("AWS_SECRET_ACCESS_KEY", "test_secret_key") defer func() { os.Unsetenv("AWS_ACCESS_KEY_ID") os.Unsetenv("AWS_SECRET_ACCESS_KEY") }() ctx := context.Background() fetcher, err := NewQueryDuckDBFetcher(ctx, "aws", "us-east-1") require.NoError(t, err) require.NotNil(t, fetcher) defer func() { err := fetcher.Close(ctx) assert.NoError(t, err) }() createTableSQL := ` CREATE TABLE slowlogs ( id BIGINT, time DOUBLE, digest VARCHAR, date VARCHAR, query VARCHAR );` _, err = fetcher.dbx.ExecContext(ctx, createTableSQL) require.NoError(t, err) insertSQL := ` INSERT INTO slowlogs (id, time, digest, date, query) VALUES (1, 10.5, 'digest_abc', '2023-01-01', 'SELECT * FROM users;'), (2, 20.0, 'digest_def', '2023-01-01', 'INSERT INTO products VALUES (1, ''test'');'), (3, 5.2, 'digest_ghi', '2023-01-02', 'UPDATE orders SET status=''completed'';');` _, err = fetcher.dbx.ExecContext(ctx, insertSQL) require.NoError(t, err) t.Run("successful query", func(t *testing.T) { query := "SELECT id, time, digest, date, query FROM slowlogs WHERE date = '2023-01-01' ORDER BY id;" var results []*SlowlogInDL err := fetcher.SelectSlowlogInDL(ctx, &results, query) require.NoError(t, err) assert.Len(t, results, 2) assert.Equal(t, int64(1), results[0].ID) assert.InDelta(t, 10.5, results[0].Time, 0.001) assert.Equal(t, "digest_abc", results[0].Digest) assert.Equal(t, "2023-01-01", results[0].Date) assert.Equal(t, sql.NullString{String: "SELECT * FROM users;", Valid: true}, results[0].Query) }) t.Run("no results found", func(t *testing.T) { query := "SELECT id, time, digest, date, query FROM slowlogs WHERE date = '2024-01-01';" var results []*SlowlogInDL err := fetcher.SelectSlowlogInDL(ctx, &results, query) require.NoError(t, err) assert.Empty(t, results) }) t.Run("invalid query", func(t *testing.T) { query := "SELECT non_existent_column FROM slowlogs;" var results []*SlowlogInDL err := fetcher.SelectSlowlogInDL(ctx, &results, query) require.Error(t, err) assert.Contains(t, err.Error(), "failed to query and scan with sqlx") assert.Nil(t, results) }) t.Run("context cancelled", func(t *testing.T) { ctxCancelled, cancel := context.WithCancel(ctx) cancel() query := "SELECT id, time, digest, date, query FROM slowlogs;" var results []*SlowlogInDL err := fetcher.SelectSlowlogInDL(ctxCancelled, &results, query) require.Error(t, err) assert.Contains(t, err.Error(), "context canceled") assert.Nil(t, results) }) } func TestSlowQueryDuckDBFetcher_SelectSlowlogInDLWithoutFilterPushdown(t *testing.T) { os.Setenv("AWS_ACCESS_KEY_ID", "test_access_key") os.Setenv("AWS_SECRET_ACCESS_KEY", "test_secret_key") defer func() { os.Unsetenv("AWS_ACCESS_KEY_ID") os.Unsetenv("AWS_SECRET_ACCESS_KEY") }() ctx := context.Background() fetcher, err := NewQueryDuckDBFetcher(ctx, "aws", "us-east-1") require.NoError(t, err) require.NotNil(t, fetcher) defer func() { err := fetcher.Close(ctx) assert.NoError(t, err) }() mockTableAlias := "mock_slowlogs" mockTablePath := "s3://my-bucket/path/to/slowlogs" // Manually add to attachedTables to simulate successful attachment for testing query execution // In a real scenario, attachTable would handle this. fetcher.mu.Lock() fetcher.attachedTables[mockTableAlias] = mockTablePath fetcher.mu.Unlock() // Now create a temporary table in memory to simulate data for querying createTableSQL := ` CREATE TABLE mock_slowlogs ( id BIGINT, time DOUBLE, digest VARCHAR, date VARCHAR, query VARCHAR );` _, err = fetcher.dbx.ExecContext(ctx, createTableSQL) require.NoError(t, err) insertSQL := ` INSERT INTO mock_slowlogs (id, time, digest, date, query) VALUES (1, 10.5, 'digest_abc', '2023-01-01', 'SELECT * FROM users_no_pushdown;'), (2, 20.0, 'digest_def', '2023-01-01', 'INSERT INTO products_no_pushdown VALUES (1, ''test'');');` _, err = fetcher.dbx.ExecContext(ctx, insertSQL) require.NoError(t, err) t.Run("successful query without filter pushdown", func(t *testing.T) { query := "SELECT id, time, digest, date, query FROM mock_slowlogs WHERE date = '2023-01-01' ORDER BY id;" var results []*SlowlogInDL err := fetcher.SelectSlowlogInDLWithoutFilterPushdown(ctx, &results, mockTableAlias, mockTablePath, query) require.NoError(t, err) assert.Len(t, results, 2) assert.Equal(t, int64(1), results[0].ID) assert.InDelta(t, 10.5, results[0].Time, 0.001) assert.Equal(t, "digest_abc", results[0].Digest) assert.Equal(t, "2023-01-01", results[0].Date) assert.Equal(t, sql.NullString{String: "SELECT * FROM users_no_pushdown;", Valid: true}, results[0].Query) }) t.Run("invalid query without filter pushdown", func(t *testing.T) { query := "SELECT non_existent_column FROM mock_slowlogs;" var results []*SlowlogInDL err := fetcher.SelectSlowlogInDLWithoutFilterPushdown(ctx, &results, mockTableAlias, mockTablePath, query) require.Error(t, err) assert.Contains(t, err.Error(), "failed to query and scan with sqlx") assert.Nil(t, results) }) }