| 252 | } |
| 253 | |
| 254 | func getRowCount(ctx context.Context, t *testing.T, driver db.Driver, _ storepb.Engine) int { |
| 255 | query := "SELECT COUNT(*) FROM test_table" |
| 256 | |
| 257 | // Use QueryConn to get results |
| 258 | conn, err := driver.GetDB().Conn(ctx) |
| 259 | require.NoError(t, err) |
| 260 | defer conn.Close() |
| 261 | |
| 262 | results, err := driver.QueryConn(ctx, conn, query, db.QueryContext{}) |
| 263 | require.NoError(t, err) |
| 264 | require.Len(t, results, 1) |
| 265 | require.Len(t, results[0].Rows, 1) |
| 266 | require.Len(t, results[0].Rows[0].Values, 1) |
| 267 | |
| 268 | // Extract count value |
| 269 | switch v := results[0].Rows[0].Values[0].Kind.(type) { |
| 270 | case *v1pb.RowValue_Int32Value: |
| 271 | return int(v.Int32Value) |
| 272 | case *v1pb.RowValue_Int64Value: |
| 273 | return int(v.Int64Value) |
| 274 | case *v1pb.RowValue_StringValue: |
| 275 | // Some databases return count as string |
| 276 | count := 0 |
| 277 | if _, err := fmt.Sscanf(v.StringValue, "%d", &count); err != nil { |
| 278 | t.Fatalf("Failed to parse count from string: %v", err) |
| 279 | } |
| 280 | return count |
| 281 | default: |
| 282 | t.Fatalf("Unexpected count value type: %T", v) |
| 283 | return 0 |
| 284 | } |
| 285 | } |
| 286 | |
| 287 | func cleanupTable(ctx context.Context, t *testing.T, driver db.Driver, _ storepb.Engine) { |
| 288 | _, err := driver.Execute(ctx, "DELETE FROM test_table", db.ExecuteOptions{}) |