⚡️ Speed up function _get_safe_command by 98%

[codeflash-ai]

📄 98% (0.98x) speedup for _get_safe_command in sentry_sdk/integrations/redis/utils.py

⏱️ Runtime : 654 microseconds → 330 microseconds (best of 97 runs)

📝 Explanation and details

The optimized code achieves a 98% speedup by eliminating redundant function calls within the loop through two key optimizations:

1. Hoisted string operation outside loop:

• Moved name_low = name.lower() before the loop instead of recalculating it on every iteration
• This converts O(n) string operations to O(1), particularly beneficial for commands with many arguments

2. Cached PII check result:

• Added send_pii = None variable to cache the result of should_send_default_pii()
• Only calls the function once when first needed (lazy evaluation), then reuses the cached boolean value
• Eliminates repeated expensive function calls that involve scope lookups and client traversal

The optimizations are most effective for:

• Commands with many arguments (200-280% faster): The cached PII check and hoisted .lower() call provide exponential benefits as argument count increases
• Non-sensitive commands with multiple args (25-40% faster): Avoids repeated PII checks for each non-key argument
• Commands exceeding max args limit (200-240% faster): Still processes efficiently despite large input due to reduced per-iteration overhead

These changes maintain identical functionality while dramatically reducing computational overhead in the critical path of Redis command processing, where this function is likely called frequently in high-throughput scenarios.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 49 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	93.8%

🌀 Generated Regression Tests and Runtime

import pytest
from sentry_sdk.integrations.redis.utils import _get_safe_command

# --- Function and dependencies for testing ---

# Constants used in _get_safe_command
SENSITIVE_DATA_SUBSTITUTE = "[Filtered]"
_MAX_NUM_ARGS = 10
_COMMANDS_INCLUDING_SENSITIVE_DATA = set([
    "auth", "set", "mset", "hmset", "config", "client", "acl", "psetex", "setex", "setnx"
])

# Simulate should_send_default_pii toggle
class DummyClient:
    def __init__(self, send_pii):
        self._send_pii = send_pii
    def should_send_default_pii(self):
        return self._send_pii
    def is_active(self):
        return True

class DummyScope:
    def __init__(self, client):
        self.client = client

_current_scope = DummyScope(DummyClient(False))  # Default: no PII
from sentry_sdk.integrations.redis.utils import _get_safe_command


# --- Pytest fixtures to toggle PII ---
@pytest.fixture
def set_pii_true():
    prev = _current_scope.client._send_pii
    _current_scope.client._send_pii = True
    yield
    _current_scope.client._send_pii = prev

# --- Unit tests ---

# 1. Basic Test Cases

def test_basic_single_arg_key_only():
    # Only key, not sensitive command, no PII
    codeflash_output = _get_safe_command("GET", ["mykey"]); cmd = codeflash_output # 2.44μs -> 2.42μs (1.12% faster)


def test_basic_multiple_args_no_pii():
    # Multiple args, not sensitive command, no PII
    codeflash_output = _get_safe_command("GET", ["mykey", "field1", "field2"]); cmd = codeflash_output # 12.5μs -> 10.0μs (24.8% faster)


def test_basic_sensitive_command_no_pii():
    # Sensitive command, all args filtered
    codeflash_output = _get_safe_command("AUTH", ["user", "pass"]); cmd = codeflash_output # 3.07μs -> 2.89μs (6.26% faster)


def test_basic_case_insensitive_sensitive_command():
    # Command name is case-insensitive
    codeflash_output = _get_safe_command("Set", ["mykey", "value"]); cmd = codeflash_output # 10.2μs -> 10.0μs (1.85% faster)

# 2. Edge Test Cases

def test_edge_empty_args():
    # No arguments
    codeflash_output = _get_safe_command("PING", []); cmd = codeflash_output # 1.13μs -> 1.38μs (18.2% slower)

def test_edge_empty_string_key():
    # Key is empty string
    codeflash_output = _get_safe_command("GET", [""]); cmd = codeflash_output # 2.40μs -> 2.39μs (0.418% faster)

def test_edge_none_arg():
    # Key is None
    codeflash_output = _get_safe_command("GET", [None]); cmd = codeflash_output # 2.24μs -> 2.25μs (0.134% slower)

def test_edge_non_string_args():
    # Args are int, float, bool
    codeflash_output = _get_safe_command("GET", [42, 3.14, True]); cmd = codeflash_output # 10.7μs -> 8.54μs (25.7% faster)


def test_edge_sensitive_command_with_non_string_args():
    # Sensitive command with non-string args
    codeflash_output = _get_safe_command("SET", [42, 3.14, True]); cmd = codeflash_output # 13.0μs -> 9.77μs (33.6% faster)

def test_edge_command_name_with_spaces():
    # Command name with spaces (should not happen, but test anyway)
    codeflash_output = _get_safe_command("GET VALUE", ["key"]); cmd = codeflash_output # 2.43μs -> 2.40μs (1.50% faster)

def test_edge_args_exceed_max_num_args():
    # More than _MAX_NUM_ARGS arguments
    args = ["key"] + ["v%d" % i for i in range(_MAX_NUM_ARGS + 2)]
    codeflash_output = _get_safe_command("GET", args); cmd = codeflash_output # 33.9μs -> 9.95μs (241% faster)
    # Only up to _MAX_NUM_ARGS+1 args (since i > _MAX_NUM_ARGS breaks)
    expected = "GET " + "'key' " + " ".join(["[Filtered]"] * _MAX_NUM_ARGS)


def test_edge_sensitive_command_exceed_max_num_args():
    args = ["key"] + ["v%d" % i for i in range(_MAX_NUM_ARGS + 2)]
    codeflash_output = _get_safe_command("SET", args); cmd = codeflash_output # 36.0μs -> 11.9μs (203% faster)
    expected = "SET " + " ".join(["[Filtered]"] * (_MAX_NUM_ARGS + 1))

def test_edge_command_name_empty_string():
    # Command name is empty string
    codeflash_output = _get_safe_command("", ["key"]); cmd = codeflash_output # 2.39μs -> 2.39μs (0.209% faster)

def test_edge_args_are_objects():
    # Arguments are objects
    class Dummy:
        def __repr__(self):
            return "<Dummy>"
    codeflash_output = _get_safe_command("GET", [Dummy(), Dummy()]); cmd = codeflash_output # 8.82μs -> 8.68μs (1.60% faster)


def test_large_many_args_no_pii():
    # Large number of args, not sensitive command, no PII
    args = ["key"] + ["val%d" % i for i in range(999)]
    codeflash_output = _get_safe_command("GET", args); cmd = codeflash_output # 35.8μs -> 11.8μs (204% faster)
    # Only up to _MAX_NUM_ARGS+1 args are processed
    expected = "GET " + "'key' " + " ".join(["[Filtered]"] * _MAX_NUM_ARGS)


def test_large_sensitive_command_many_args():
    args = ["key"] + ["val%d" % i for i in range(999)]
    codeflash_output = _get_safe_command("SET", args); cmd = codeflash_output # 34.2μs -> 11.0μs (211% faster)
    expected = "SET " + " ".join(["[Filtered]"] * (_MAX_NUM_ARGS + 1))

def test_large_args_all_none():
    # All arguments are None
    args = [None for _ in range(_MAX_NUM_ARGS + 2)]
    codeflash_output = _get_safe_command("GET", args); cmd = codeflash_output # 29.1μs -> 9.88μs (195% faster)
    expected = "GET None " + " ".join(["[Filtered]"] * _MAX_NUM_ARGS)


def test_large_args_all_empty_strings():
    args = [""] * (_MAX_NUM_ARGS + 2)
    codeflash_output = _get_safe_command("GET", args); cmd = codeflash_output # 30.7μs -> 12.1μs (154% faster)
    expected = "GET '' " + " ".join(["[Filtered]"] * _MAX_NUM_ARGS)


def test_large_sensitive_command_all_empty_strings():
    args = [""] * (_MAX_NUM_ARGS + 2)
    codeflash_output = _get_safe_command("SET", args); cmd = codeflash_output # 36.2μs -> 11.4μs (217% faster)
    expected = "SET " + " ".join(["[Filtered]"] * (_MAX_NUM_ARGS + 1))
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
from typing import Any, Sequence

# imports
import pytest
from sentry_sdk.integrations.redis.utils import _get_safe_command

# Setup for tests: define constants and context
SENSITIVE_DATA_SUBSTITUTE = "[Filtered]"
_MAX_NUM_ARGS = 10
_COMMANDS_INCLUDING_SENSITIVE_DATA = {"auth", "set", "hmset", "mset", "config", "client", "acl", "cluster"}

# Minimal mock for should_send_default_pii, which can be toggled in tests
class DummyClient:
    def __init__(self, send_pii):
        self._send_pii = send_pii
    def should_send_default_pii(self):
        return self._send_pii
    def is_active(self):
        return True

class DummyScope:
    def __init__(self, client):
        self.client = client

# ContextVar mock (since we don't have sentry_sdk's ContextVar)
class ContextVar:
    def __init__(self, name, default=None):
        self.value = default
    def get(self):
        return self.value
    def set(self, v):
        self.value = v

_current_scope = ContextVar("current_scope", default=None)
from sentry_sdk.integrations.redis.utils import _get_safe_command

# -------------------- UNIT TESTS --------------------

# Helper to set up the client context for tests
def set_pii_context(send_pii):
    client = DummyClient(send_pii)
    scope = DummyScope(client)
    _current_scope.set(scope)

# ----------- BASIC TEST CASES -----------

def test_basic_command_no_args():
    # Test: No arguments, not sensitive command, should just return the command name
    set_pii_context(False)
    codeflash_output = _get_safe_command("PING", []) # 1.13μs -> 1.40μs (19.1% slower)

def test_basic_command_with_args_pii_false():
    # Test: Not sensitive, args, pii disabled, only key revealed, rest filtered
    set_pii_context(False)
    codeflash_output = _get_safe_command("GET", ["mykey", "extra"]) # 8.43μs -> 8.10μs (4.07% faster)

def test_basic_command_with_args_pii_true():
    # Test: Not sensitive, args, pii enabled, all args revealed
    set_pii_context(True)
    codeflash_output = _get_safe_command("GET", ["mykey", "extra"]) # 7.52μs -> 7.30μs (2.96% faster)

def test_sensitive_command_pii_false():
    # Test: Sensitive command, args, pii disabled, all args filtered
    set_pii_context(False)
    codeflash_output = _get_safe_command("AUTH", ["user", "password"]) # 2.57μs -> 2.39μs (7.75% faster)

def test_sensitive_command_pii_true():
    # Test: Sensitive command, args, pii enabled, all args filtered (sensitivity trumps pii)
    set_pii_context(True)
    codeflash_output = _get_safe_command("AUTH", ["user", "password"]) # 2.31μs -> 2.27μs (1.76% faster)

def test_case_insensitive_sensitive_command():
    # Test: Sensitive command with upper/lowercase, should be filtered
    set_pii_context(False)
    codeflash_output = _get_safe_command("SeT", ["foo", "bar"]) # 8.25μs -> 7.90μs (4.52% faster)

def test_non_sensitive_command_with_numeric_args():
    # Test: Not sensitive, numeric args, pii disabled
    set_pii_context(False)
    codeflash_output = _get_safe_command("INCRBY", ["counter", 5]) # 7.83μs -> 7.33μs (6.79% faster)

def test_non_sensitive_command_with_numeric_args_pii_true():
    # Test: Not sensitive, numeric args, pii enabled
    set_pii_context(True)
    codeflash_output = _get_safe_command("INCRBY", ["counter", 5]) # 7.60μs -> 7.25μs (4.91% faster)

# ----------- EDGE TEST CASES -----------

def test_max_num_args_limit():
    # Test: Args exceeding _MAX_NUM_ARGS, should only process up to _MAX_NUM_ARGS+1
    set_pii_context(False)
    args = ["key"] + ["v"] * (_MAX_NUM_ARGS + 5)
    codeflash_output = _get_safe_command("GET", args); result = codeflash_output # 27.7μs -> 9.28μs (199% faster)
    # Only first _MAX_NUM_ARGS+1 args processed (key + _MAX_NUM_ARGS others)
    expected = "GET " + "'key' " + " ".join(["[Filtered]"] * _MAX_NUM_ARGS)

def test_empty_string_key():
    # Test: Key is empty string, should show repr of empty string
    set_pii_context(False)
    codeflash_output = _get_safe_command("GET", [""]) # 2.28μs -> 2.27μs (0.751% faster)

def test_none_key():
    # Test: Key is None, should show repr of None
    set_pii_context(False)
    codeflash_output = _get_safe_command("GET", [None]) # 2.20μs -> 2.15μs (2.18% faster)

def test_command_with_special_characters():
    # Test: Key/args with special characters
    set_pii_context(False)
    codeflash_output = _get_safe_command("GET", ["key\n", "val\t"]) # 7.96μs -> 7.84μs (1.60% faster)

def test_command_with_tuple_and_dict_arg():
    # Test: Key is tuple, arg is dict
    set_pii_context(True)
    codeflash_output = _get_safe_command("GET", [(1,2), {"a": "b"}]) # 7.91μs -> 8.13μs (2.67% slower)

def test_sensitive_command_with_no_args():
    # Test: Sensitive command, no args
    set_pii_context(False)
    codeflash_output = _get_safe_command("AUTH", []) # 1.08μs -> 1.32μs (18.2% slower)

def test_command_with_boolean_args():
    # Test: Boolean args, pii enabled
    set_pii_context(True)
    codeflash_output = _get_safe_command("GET", ["key", True, False]) # 11.2μs -> 8.11μs (38.7% faster)

def test_command_with_large_arg_value():
    # Test: Large string value as key
    set_pii_context(False)
    large_key = "x" * 100
    codeflash_output = _get_safe_command("GET", [large_key]) # 2.55μs -> 2.32μs (9.90% faster)

def test_command_with_unicode_characters():
    # Test: Unicode in key/args
    set_pii_context(True)
    codeflash_output = _get_safe_command("GET", ["ключ", "значение"]) # 8.40μs -> 8.39μs (0.071% faster)

def test_command_with_empty_args_list():
    # Test: Empty args list, should just return command name
    set_pii_context(True)
    codeflash_output = _get_safe_command("GET", []) # 1.09μs -> 1.29μs (15.5% slower)

def test_sensitive_command_with_many_args():
    # Test: Sensitive command with many args, all filtered
    set_pii_context(True)
    args = ["a", "b", "c", "d"]
    codeflash_output = _get_safe_command("SET", args) # 13.0μs -> 8.26μs (57.7% faster)

# ----------- LARGE SCALE TEST CASES -----------

def test_large_number_of_args_pii_false():
    # Test: Large number of args, pii disabled
    set_pii_context(False)
    args = ["key"] + [f"val{i}" for i in range(999)]
    codeflash_output = _get_safe_command("GET", args); result = codeflash_output # 33.6μs -> 8.83μs (280% faster)
    # Only up to _MAX_NUM_ARGS+1 args processed
    expected = "GET " + "'key' " + " ".join(["[Filtered]"] * _MAX_NUM_ARGS)

def test_large_number_of_args_pii_true():
    # Test: Large number of args, pii enabled
    set_pii_context(True)
    args = ["key"] + [f"val{i}" for i in range(999)]
    codeflash_output = _get_safe_command("GET", args); result = codeflash_output # 27.1μs -> 9.59μs (183% faster)
    # Only up to _MAX_NUM_ARGS+1 args processed
    expected = "GET " + "'key' " + " ".join([repr(f"val{i}") for i in range(_MAX_NUM_ARGS)])

def test_large_sensitive_command():
    # Test: Sensitive command, large number of args, all filtered
    set_pii_context(True)
    args = ["a"] * 999
    codeflash_output = _get_safe_command("SET", args); result = codeflash_output # 28.2μs -> 8.94μs (215% faster)
    expected = "SET " + " ".join(["[Filtered]"] * _MAX_NUM_ARGS)

def test_large_command_name():
    # Test: Large command name, should be preserved
    set_pii_context(False)
    large_cmd = "CMD" * 100
    args = ["key", "val"]
    expected = f"{large_cmd} 'key' [Filtered]"
    codeflash_output = _get_safe_command(large_cmd, args) # 7.72μs -> 7.22μs (6.92% faster)

def test_large_args_with_various_types():
    # Test: Large args list with mixed types, pii enabled
    set_pii_context(True)
    args = ["key"] + [i for i in range(_MAX_NUM_ARGS)]
    expected = "GET 'key' " + " ".join([repr(i) for i in range(_MAX_NUM_ARGS)])
    codeflash_output = _get_safe_command("GET", args) # 32.7μs -> 8.66μs (277% faster)

def test_large_args_sensitive_command():
    # Test: Sensitive command, large args, all filtered
    set_pii_context(False)
    args = ["a"] * 1000
    expected = "SET " + " ".join(["[Filtered]"] * _MAX_NUM_ARGS)
    codeflash_output = _get_safe_command("SET", args) # 27.7μs -> 8.81μs (215% faster)

# ----------- MUTATION TESTING CASES (robustness) -----------

def test_mutation_sensitive_command_not_filtered():
    # If sensitive command does not filter, this test fails
    set_pii_context(True)
    codeflash_output = _get_safe_command("AUTH", ["user", "pass"]); result = codeflash_output # 2.45μs -> 2.30μs (6.43% faster)

def test_mutation_non_sensitive_command_key_is_not_first():
    # If key is not shown only for first arg, this test fails
    set_pii_context(False)
    codeflash_output = _get_safe_command("GET", ["key1", "key2", "key3"]); result = codeflash_output # 10.3μs -> 7.79μs (32.6% faster)
    parts = result.split()
    for p in parts[2:]:
        pass

def test_mutation_non_sensitive_command_all_args_filtered_pii_true():
    # If pii is True but args are still filtered, this test fails
    set_pii_context(True)
    codeflash_output = _get_safe_command("GET", ["a", "b", "c"]); result = codeflash_output # 10.6μs -> 7.60μs (39.3% faster)

def test_mutation_sensitive_command_args_revealed():
    # If sensitive command reveals args, this test fails
    set_pii_context(True)
    codeflash_output = _get_safe_command("SET", ["foo", "bar"]); result = codeflash_output # 7.42μs -> 6.58μs (12.7% faster)

def test_mutation_max_num_args_exceeded():
    # If function does not stop at _MAX_NUM_ARGS+1, this test fails
    set_pii_context(False)
    args = ["key"] + ["v"] * (_MAX_NUM_ARGS + 100)
    codeflash_output = _get_safe_command("GET", args); result = codeflash_output # 27.7μs -> 8.97μs (209% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

To edit these changes git checkout codeflash/optimize-_get_safe_command-mg9l00v8 and push.