Feat: BQ io_manager to allow allow column case presevetion#33973
Feat: BQ io_manager to allow allow column case presevetion#33973adschem wants to merge 3 commits into
Conversation
Greptile SummaryThis PR adds an opt-in
Confidence Score: 5/5Safe to merge — purely additive with default=True preserving all existing behavior unchanged. The change is purely opt-in: the default value keeps every existing pipeline's column-casing behavior identical. The implementation follows the existing pattern for other config values and both write and read paths are covered by new unit tests. No files require special attention; the base-class placement of
|
| Filename | Overview |
|---|---|
| python_modules/libraries/dagster-gcp/dagster_gcp/bigquery/io_manager.py | Adds uppercase_column_names: bool field (default True) to BigQueryIOManager. The field is read by the Pandas type handler but is placed on the shared base class, making it visible to all BigQuery IO manager users regardless of their type handler. |
| python_modules/libraries/dagster-gcp-pandas/dagster_gcp_pandas/bigquery/bigquery_pandas_type_handler.py | Adds _should_uppercase_column_names helper and wires it into both handle_output (conditional rename(columns=str.upper)) and load_input (conditional str.lower on result columns). Logic follows existing patterns for reading location/timeout from resource_config, and default=True preserves backward compatibility. |
| python_modules/libraries/dagster-gcp-pandas/dagster_gcp_pandas_tests/bigquery/test_type_handler.py | Adds test_handle_output_preserve_column_case and test_load_input_preserve_column_case, covering both sides of the new flag (write and read) with mocked connections and mixed-case column names. |
Flowchart
%%{init: {'theme': 'neutral'}}%%
flowchart TD
A[BigQueryIOManager config\nuppercase_column_names: bool\ndefault = True] --> B{_should_uppercase_column_names\ncontext.resource_config}
B -->|True / not set| C[handle_output:\nobj.rename columns=str.upper]
B -->|False| D[handle_output:\npass obj as-is]
C --> E[load_table_from_dataframe\nUPPERCASE columns in BQ]
D --> F[load_table_from_dataframe\noriginal-case columns in BQ]
E --> G[load_input:\nresult.columns = map str.lower]
F --> H[load_input:\nno column rename]
G --> I[Return DataFrame\nlowercase columns]
H --> J[Return DataFrame\noriginal-case columns]
%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
flowchart TD
A[BigQueryIOManager config\nuppercase_column_names: bool\ndefault = True] --> B{_should_uppercase_column_names\ncontext.resource_config}
B -->|True / not set| C[handle_output:\nobj.rename columns=str.upper]
B -->|False| D[handle_output:\npass obj as-is]
C --> E[load_table_from_dataframe\nUPPERCASE columns in BQ]
D --> F[load_table_from_dataframe\noriginal-case columns in BQ]
E --> G[load_input:\nresult.columns = map str.lower]
F --> H[load_input:\nno column rename]
G --> I[Return DataFrame\nlowercase columns]
H --> J[Return DataFrame\noriginal-case columns]
Reviews (2): Last reviewed commit: "adding in the load_input test feedback f..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| def test_handle_output_preserve_column_case(): | ||
| handler = BigQueryPandasTypeHandler() | ||
| df = pd.DataFrame({"Foo": ["a", "b"], "bar": [1, 2]}) | ||
| connection = MagicMock() | ||
| output_context = build_output_context(resource_config={"uppercase_column_names": False}) | ||
|
|
||
| handler.handle_output( | ||
| output_context, | ||
| TableSlice(table="my_table", schema="my_schema", database="my_db"), | ||
| df, | ||
| connection, | ||
| ) | ||
|
|
||
| loaded_df = connection.load_table_from_dataframe.call_args.kwargs["dataframe"] | ||
| assert list(loaded_df.columns) == ["Foo", "bar"] |
There was a problem hiding this comment.
load_input with uppercase_column_names=False
The new test only exercises handle_output; the symmetric change in load_input (skipping str.lower when the flag is False) has no unit test. A test with connection.query(...) mocked to return a DataFrame whose columns are already mixed-case (e.g., "Foo" / "bar") would confirm that those column names are returned unchanged when uppercase_column_names=False, and lowercased only when the flag is True (existing default). Without this coverage, the load_input side of the feature can be silently broken in a future refactor without any test catching it.
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
1 participant
Summary & Motivation
Test Plan
Changelog