Skip to content
This repository was archived by the owner on Jun 5, 2025. It is now read-only.

Commit 02e3efd

Browse files
committed
simplify signature initialization for optimization
Remove unnecessary lazy loading mechanism from signatures since they are always loaded during server startup when CodegateSecrets is created. This simplifies the code and makes the initialization flow singular, which will reduce uneeded lazy loading on every call
1 parent 4c1b4be commit 02e3efd

File tree

6 files changed

+9
-46
lines changed

6 files changed

+9
-46
lines changed

src/codegate/cli.py

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -331,7 +331,6 @@ def serve(
331331
logger.exception("Unexpected error occurred")
332332
click.echo(f"Error: {e}", err=True)
333333
sys.exit(1)
334-
335334

336335

337336
async def run_servers(cfg: Config, app) -> None:

src/codegate/pipeline/secrets/secrets.py

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -26,9 +26,8 @@ class CodegateSecrets(PipelineStep):
2626
def __init__(self):
2727
"""Initialize the CodegateSecrets pipeline step."""
2828
super().__init__()
29-
# Initialize signatures eagerly
29+
# Initialize and load signatures immediately
3030
CodegateSignatures.initialize("signatures.yaml")
31-
CodegateSignatures._ensure_signatures_loaded()
3231

3332
@property
3433
def name(self) -> str:

src/codegate/pipeline/secrets/signatures.py

Lines changed: 8 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,6 @@ class CodegateSignatures:
3939
"""Main class for detecting secrets in text using regex patterns."""
4040

4141
_instance_lock: ClassVar[Lock] = Lock()
42-
_signatures_loaded: ClassVar[bool] = False
4342
_signature_groups: ClassVar[List[SignatureGroup]] = []
4443
_compiled_regexes: ClassVar[Dict[str, re.Pattern]] = {}
4544
_yaml_path: ClassVar[Optional[str]] = None
@@ -48,23 +47,22 @@ class CodegateSignatures:
4847
def reset(cls) -> None:
4948
"""Reset the cached patterns."""
5049
with cls._instance_lock:
51-
cls._signatures_loaded = False
5250
cls._signature_groups = []
5351
cls._compiled_regexes = {}
5452
cls._yaml_path = None
5553
logger.debug("SecretFinder cache reset")
5654

5755
@classmethod
5856
def initialize(cls, yaml_path: str) -> None:
59-
"""Initialize the SecretFinder with a YAML file path."""
57+
"""Initialize the SecretFinder with a YAML file path and load signatures."""
6058
if not Path(yaml_path).exists():
6159
raise FileNotFoundError(f"Signatures file not found: {yaml_path}")
6260

6361
with cls._instance_lock:
6462
# Only initialize if not already initialized with this path
6563
if cls._yaml_path != yaml_path:
6664
cls._yaml_path = yaml_path
67-
cls._signatures_loaded = False
65+
cls._load_signatures()
6866
logger.debug(f"SecretFinder initialized with {yaml_path}")
6967

7068
@classmethod
@@ -174,6 +172,10 @@ def _add_signature_group(cls, name: str, patterns: Dict[str, str]) -> None:
174172
def _load_signatures(cls) -> None:
175173
"""Load signature patterns from the YAML file."""
176174
try:
175+
# Clear existing signatures before loading new ones
176+
cls._signature_groups = []
177+
cls._compiled_regexes = {}
178+
177179
yaml_data = cls._load_yaml(cls._yaml_path)
178180

179181
# Add custom GitHub token patterns
@@ -207,32 +209,14 @@ def _load_signatures(cls) -> None:
207209
logger.error(f"Error loading signatures: {e}")
208210
raise
209211

210-
@classmethod
211-
def _ensure_signatures_loaded(cls) -> None:
212-
"""Ensure signatures are loaded before use."""
213-
if not cls._signatures_loaded:
214-
with cls._instance_lock:
215-
if not cls._signatures_loaded:
216-
if not cls._yaml_path:
217-
raise RuntimeError("SecretFinder not initialized. Call initialize() first.")
218-
try:
219-
cls._load_signatures()
220-
cls._signatures_loaded = True
221-
except Exception as e:
222-
logger.error(f"Failed to load signatures: {e}")
223-
raise
224-
225212
@classmethod
226213
def find_in_string(cls, text: str) -> List[Match]:
227214
"""Search for secrets in the provided string."""
228215
if not text:
229216
return []
230217

231-
try:
232-
cls._ensure_signatures_loaded()
233-
except Exception as e:
234-
logger.error(f"Failed to load signatures: {e}")
235-
return []
218+
if not cls._yaml_path:
219+
raise RuntimeError("SecretFinder not initialized. Call initialize() first.")
236220

237221
matches = []
238222
lines = text.splitlines()

tests/pipeline/secrets/test_signatures.py

Lines changed: 0 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -69,14 +69,12 @@ def test_initialize_and_reset(temp_yaml_file):
6969

7070
CodegateSignatures.reset()
7171
assert CodegateSignatures._yaml_path is None
72-
assert not CodegateSignatures._signatures_loaded
7372
assert not CodegateSignatures._signature_groups
7473
assert not CodegateSignatures._compiled_regexes
7574

7675

7776
def test_find_in_string_with_aws_credentials(temp_yaml_file):
7877
CodegateSignatures.initialize(temp_yaml_file)
79-
CodegateSignatures._signatures_loaded = False # Force reload of signatures
8078

8179
test_string = """
8280
aws_access_key = 'AKIAIOSFODNN7EXAMPLE'
@@ -100,7 +98,6 @@ def test_find_in_string_with_github_token():
10098

10199
try:
102100
CodegateSignatures.initialize(f.name)
103-
CodegateSignatures._signatures_loaded = False # Force reload of signatures
104101

105102
test_string = "github_token = 'ghp_1234567890abcdef1234567890abcdef123456'"
106103
matches = CodegateSignatures.find_in_string(test_string)
@@ -121,7 +118,6 @@ def test_find_in_string_with_github_token():
121118

122119
def test_find_in_string_with_no_matches(temp_yaml_file):
123120
CodegateSignatures.initialize(temp_yaml_file)
124-
CodegateSignatures._signatures_loaded = False # Force reload of signatures
125121

126122
test_string = "No secrets here!"
127123
matches = CodegateSignatures.find_in_string(test_string)
@@ -158,7 +154,6 @@ def test_duplicate_patterns():
158154

159155
try:
160156
CodegateSignatures.initialize(f.name)
161-
CodegateSignatures._signatures_loaded = False # Force reload of signatures
162157

163158
test_string = "aws_key = 'AKIAIOSFODNN7EXAMPLE'"
164159
matches = CodegateSignatures.find_in_string(test_string)

tests/test_cli.py

Lines changed: 0 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -105,13 +105,6 @@ def test_provider_registration(mock_registry, mock_secrets_mgr, mock_pipeline_fa
105105
assert "ollama" in provider_names
106106

107107

108-
@patch("codegate.server.CodegateSignatures")
109-
def test_signatures_initialization(mock_signatures, mock_pipeline_factory) -> None:
110-
"""Test that signatures are initialized correctly."""
111-
init_app(mock_pipeline_factory)
112-
mock_signatures.initialize.assert_called_once_with("signatures.yaml")
113-
114-
115108
def test_pipeline_initialization(mock_pipeline_factory) -> None:
116109
"""Test that pipelines are initialized correctly."""
117110
app = init_app(mock_pipeline_factory)

tests/test_server.py

Lines changed: 0 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -99,13 +99,6 @@ def test_provider_registration(mock_registry, mock_secrets_mgr, mock_pipeline_fa
9999
assert "ollama" in provider_names
100100

101101

102-
@patch("codegate.server.CodegateSignatures")
103-
def test_signatures_initialization(mock_signatures, mock_pipeline_factory) -> None:
104-
"""Test that signatures are initialized correctly."""
105-
init_app(mock_pipeline_factory)
106-
mock_signatures.initialize.assert_called_once_with("signatures.yaml")
107-
108-
109102
def test_pipeline_initialization(mock_pipeline_factory) -> None:
110103
"""Test that pipelines are initialized correctly."""
111104
app = init_app(mock_pipeline_factory)

0 commit comments

Comments
 (0)