fix(jetstream): use .get() for optional analysis_basis field

[abhi-0203]

Summary

Fixes #16535 — fetch_experiment_data crashes with KeyError: 'analysis_basis' when a statistics row omits the field.

Problem

In get_experiment_data, statistics rows are read from raw JSON and accessed with hard key lookups:

if point["analysis_basis"] == AnalysisBasis.ENROLLMENTS:

The analysis_basis field is Optional in the Statistic schema, so rows that omit it pass Pydantic validation cleanly but crash at the dictionary lookup. This causes fetch_experiment_data to fail for older experiments whose analysis results don't include this field.

Fix

Replace point["analysis_basis"] with point.get("analysis_basis") on both comparison lines (ENROLLMENTS and EXPOSURES). When the field is missing, .get() returns None, which won't match either basis — the point is silently skipped, which is the correct behavior for incomplete data.

Validation

• Verified the bug: point["analysis_basis"] raises KeyError when the field is absent
• Verified the fix: point.get("analysis_basis") returns None, allowing the comparison to safely skip the point
• No other code paths depend on the exception being raised
• Two-line change, no behavioral difference for experiments that already include analysis_basis

[jaredlockhart]

Hi @abhi-0203 ! Thanks for contributing! Please add tests for this case and verify with make code_format check and then we can get this merged 🎉

[freshstrangemusic]

Please also look at CONTRIBUTING.md for how we write our PRs

[abhi-0203]

Added tests and ran formatting. Tests verify:

• Missing analysis_basis doesn't raise KeyError
• Present analysis_basis works correctly

Used black + isort for formatting. Let me know if anything else is needed!

[abhi-0203]

Updated commit to follow CONTRIBUTING.md format:

• Commit: test(jetstream): add tests for analysis_basis handling
• Added Because/This commit/Fixes structure
• Tests added for both missing and present analysis_basis cases

Note: couldn't run full make check locally (no Docker), but formatting verified with black + isort.

[yashikakhurana]

Updated commit to follow CONTRIBUTING.md format:

* Commit: test(jetstream): add tests for analysis_basis handling

* Added Because/This commit/Fixes structure

* Tests added for both missing and present analysis_basis cases

Note: couldn't run full make check locally (no Docker), but formatting verified with black + isort.

Thank you @abhi-0203 It seems like some formatting still need to be done, give it another shot to format files please, if you encountered any issues please feel free to reach out