fix: avoid mutating the caller's common_attributes in timestream.write

[hsusul]

Feature or Bugfix

• Bugfix

Detail

wr.timestream.write mutates the common_attributes dictionary supplied by the caller. Reusing that dictionary for a later write silently discards the later call's version, time_unit and measure_name/measure_col arguments.

Affected API: awswrangler.timestream.write (via awswrangler/timestream/_write.py::_sanitize_common_attributes)

Root cause

_sanitize_common_attributes resolves its defaults with setdefault on the object it was handed, not on a copy:

common_attributes = {} if not common_attributes else common_attributes
common_attributes.setdefault("Version", version)
common_attributes.setdefault("TimeUnit", _check_time_unit(...))
...
common_attributes.setdefault("MeasureName", measure_name)

After one write, the caller's dictionary permanently carries Version, TimeUnit and MeasureName. Because common_attributes takes precedence over the other arguments by design, every subsequent write reusing that dictionary is pinned to the first call's values.

Reproduction (no AWS credentials — moto only)

import datetime, boto3, moto, pandas as pd, awswrangler as wr

with moto.mock_aws():
    session = boto3.Session(region_name="us-east-1")
    client = session.client("timestream-write")
    client.create_database(DatabaseName="sampleDB")
    client.create_table(DatabaseName="sampleDB", TableName="sampleTable")

    captured = []
    session.events.register(
        "provide-client-params.timestream-write.WriteRecords",
        lambda params, **kwargs: captured.append(params),
    )

    df = pd.DataFrame({
        "time": [datetime.datetime(2024, 1, 1, tzinfo=datetime.timezone.utc)],
        "cpu": [1.0],
        "mem": [2.0],
    })
    common_attributes = {"Dimensions": [{"Name": "host", "Value": "h1", "DimensionValueType": "VARCHAR"}]}

    wr.timestream.write(df=df, database="sampleDB", table="sampleTable", time_col="time",
                        measure_col="cpu", common_attributes=common_attributes,
                        boto3_session=session, use_threads=False)
    wr.timestream.write(df=df, database="sampleDB", table="sampleTable", time_col="time",
                        measure_col="mem", version=5, time_unit="MICROSECONDS",
                        common_attributes=common_attributes,
                        boto3_session=session, use_threads=False)

    print(common_attributes)
    print(captured[1]["CommonAttributes"], captured[1]["Records"])

Current behavior

{'Dimensions': [...], 'Version': 1, 'TimeUnit': 'MILLISECONDS', 'MeasureName': 'cpu'}
{'Dimensions': [...], 'Version': 1, 'TimeUnit': 'MILLISECONDS', 'MeasureName': 'cpu'} [{'Time': '1704067200000', ...}]

The caller's dictionary has been modified, and the second write ignores measure_col="mem", version=5 and time_unit="MICROSECONDS". The mem values are stored under the measure name cpu, at version 1 (so a genuine upsert is silently rejected), and the timestamp is emitted in milliseconds while the request declares nothing to the contrary — a 1000x error against the requested unit.

Corrected behavior

{'Dimensions': [...]}
{'Dimensions': [...], 'Version': 5, 'TimeUnit': 'MICROSECONDS', 'MeasureName': 'mem'} [{'Time': '1704067200000000', ...}]

Implementation

One-line change in _sanitize_common_attributes: copy the dictionary before resolving defaults.

-    common_attributes = {} if not common_attributes else common_attributes
+    # Copy to avoid mutating the caller's dictionary with the defaults resolved below
+    common_attributes = {} if not common_attributes else dict(common_attributes)

A shallow copy is sufficient — only top-level keys are written; nested values such as Dimensions are read, never modified. Precedence semantics, validation and error behavior are unchanged, and no public API changed.

Regression tests

Added to tests/unit/test_moto.py (the module that runs in the PR test job), using a new moto_timestream_session fixture. No AWS credentials required.

• test_timestream_write_does_not_mutate_common_attributes — the caller's dictionary is byte-for-byte unchanged after a successful write.
• test_timestream_write_common_attributes_reused_across_calls — reusing one dictionary across two writes: the second call's measure_col, version and time_unit are honored, asserted on the captured WriteRecords parameters including the formatted epoch value (timezone-aware input, so the assertion is independent of the local timezone).
• test_timestream_write_common_attributes_take_precedence — guard that the copy did not change precedence: MeasureName, TimeUnit and Version supplied inside common_attributes still win over measure_name, time_unit and version.

The first two fail on current main and pass with the fix; the third passes both ways.

Validation (macOS, Python 3.13.5, moto 5.2.2, boto3 1.42.68)

Run exactly as in .github/workflows/minimal-tests.yml:

• pytest tests/unit/test_metadata.py tests/unit/test_session.py tests/unit/test_utils.py — 20 passed
• pytest -n 4 tests/unit/test_moto.py — 49 passed
• ruff format --check . — 280 files already formatted
• ruff check . — all checks passed
• mypy awswrangler — 19 errors, identical to the pre-existing count on main; none in awswrangler/timestream
• doc8 --max-line-length 120 docs/source — exit 0
• uv lock --check — up to date
• git diff --check — clean

Not run: the AWS integration suites (tests/unit/test_timestream.py and the other tests/unit/* modules), which require a live AWS account and the test_infra CDK stacks. The behavior they would cover is exercised locally through moto and through assertions on the exact WriteRecords request parameters.

Compatibility

No public API, signature, default or dtype change. Callers that relied on reading back the resolved defaults from their own common_attributes dictionary after a write — undocumented behavior — would no longer see them; the resolved values are still sent in the request as before.

Relates

• No existing issue; found while reading awswrangler/timestream/_write.py.

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.