From 9e5fbf7892c139fb3576202ee10b3f0e0c11d750 Mon Sep 17 00:00:00 2001 From: ZStriker19 Date: Tue, 23 Apr 2024 05:49:57 -0400 Subject: [PATCH] fix parsing code and tests --- ddtrace/settings/config.py | 36 ++++++++++++++++--- .../remoteconfig/test_remoteconfig.py | 16 ++++++--- 2 files changed, 43 insertions(+), 9 deletions(-) diff --git a/ddtrace/settings/config.py b/ddtrace/settings/config.py index 5eb004d0d9c..910eb207560 100644 --- a/ddtrace/settings/config.py +++ b/ddtrace/settings/config.py @@ -762,7 +762,7 @@ def _handle_remoteconfig(self, data, test_tracer=None): if trace_sampling_rules: # returns None if we error out trace_sampling_rules = self.convert_rc_trace_sampling_rules(trace_sampling_rules) - if trace_sampling_rules is not None: + if trace_sampling_rules: base_rc_config["_trace_sampling_rules"] = trace_sampling_rules if "log_injection_enabled" in lib_config: @@ -815,8 +815,27 @@ def enable_remote_configuration(self): remoteconfig_poller.register("AGENT_CONFIG", remoteconfig_pubsub) remoteconfig_poller.register("AGENT_TASK", remoteconfig_pubsub) + def _remove_invalid_rules(self, rc_rules: List) -> List: + """Remove invalid sampling rules from the JSON string.""" + # loop through list of dictionaries, if a dictionary doesn't have certain attributes, remove it + for rule in rc_rules: + if ( + ("service" not in rule and "name" not in rule and "resource" not in rule and "tags" not in rule) + or "sample_rate" not in rule + or "provenance" not in rule + ): + log.debug("Invalid sampling rule from remoteconfig found, rule will be removed: %s", rule) + rc_rules.remove(rule) + + return rc_rules + def _tags_to_dict(self, tags): - return {tag["key"]: tag["value_glob"] for tag in tags} + """ + Converts a list of tag dictionaries to a single dictionary. + """ + if isinstance(tags, list): + return {tag["key"]: tag["value_glob"] for tag in tags} + return tags def convert_rc_trace_sampling_rules(self, rc_rules: List[Dict[str, Any]]) -> Optional[str]: """Example of an incoming rule: @@ -843,5 +862,14 @@ def convert_rc_trace_sampling_rules(self, rc_rules: List[Dict[str, Any]]) -> Opt Example of a converted rule: '[{"sample_rate":1.0,"service":"my-service","resource":"*","name":"web.request","tags":{"care_about":"yes","region":"us-*"},provenance":"customer"}]' """ - # Convert JSON to string - return json.dumps(rc_rules, default=lambda o: self._tags_to_dict(o) if isinstance(o, list) else o) + rc_rules = self._remove_invalid_rules(rc_rules) + for rule in rc_rules: + # Convert tags to dictionary + tags = rule.get("tags") + if tags: + rule["tags"] = self._tags_to_dict(tags) + if rc_rules: + # Convert JSON to string + return json.dumps(rc_rules) + else: + return None diff --git a/tests/internal/remoteconfig/test_remoteconfig.py b/tests/internal/remoteconfig/test_remoteconfig.py index f007f77bace..0a645079786 100644 --- a/tests/internal/remoteconfig/test_remoteconfig.py +++ b/tests/internal/remoteconfig/test_remoteconfig.py @@ -447,7 +447,8 @@ def test_rc_default_products_registered(): "tags": [{"key": "care_about", "value_glob": "yes"}, {"key": "region", "value_glob": "us-*"}], } ], - '[{"sample_rate":1.0,"service":"my-service","resource":"*","name":"web.request","tags":{"care_about":"yes","region":"us-*"},"provenance":"dynamic"}]', + '[{"service": "my-service", "name": "web.request", "resource": "*", "provenance": "dynamic",' + ' "sample_rate": 1.0, "tags": {"care_about": "yes", "region": "us-*"}}]', [ SamplingRule( sample_rate=1.0, @@ -469,7 +470,8 @@ def test_rc_default_products_registered(): "tags": [{"key": "care_about", "value_glob": "yes"}, {"key": "region", "value_glob": "us-*"}], } ], - '[{"sample_rate":1.0,"resource":"*","name":"web.request","tags":{"care_about":"yes","region":"us-*"},"provenance":"customer"}]', + '[{"name": "web.request", "resource": "*", "provenance": "customer", "sample_rate": 1.0, "tags": ' + '{"care_about": "yes", "region": "us-*"}}]', [ SamplingRule( sample_rate=1.0, @@ -492,7 +494,8 @@ def test_rc_default_products_registered(): "sample_rate": 1.0, } ], - '[{"sample_rate":1.0,"service":"my-service","resource":"*","name":"web.request","provenance":"customer"}]', + '[{"service": "my-service", "name": "web.request", "resource": "*", "provenance": ' + '"customer", "sample_rate": 1.0}]', [ SamplingRule( sample_rate=1.0, @@ -515,7 +518,8 @@ def test_rc_default_products_registered(): "tags": [{"key": "care_about", "value_glob": "yes"}, {"key": "region", "value_glob": "us-*"}], } ], - '[{"sample_rate":1.0,"service":"my-service","name":"web.request","tags":{"care_about":"yes","region":"us-*"},"provenance":"customer"}]', + '[{"service": "my-service", "name": "web.request", "provenance": "customer", "sample_rate": 1.0, "tags":' + ' {"care_about": "yes", "region": "us-*"}}]', [ SamplingRule( sample_rate=1.0, @@ -538,7 +542,8 @@ def test_rc_default_products_registered(): "tags": [{"key": "care_about", "value_glob": "yes"}, {"key": "region", "value_glob": "us-*"}], } ], - '[{"sample_rate":1.0,"service":"my-service","resource":"*","tags":{"care_about":"yes","region":"us-*"},"provenance":"customer"}]', + '[{"service": "my-service", "resource": "*", "provenance": "customer", "sample_rate": 1.0, "tags":' + ' {"care_about": "yes", "region": "us-*"}}]', [ SamplingRule( sample_rate=1.0, @@ -578,6 +583,7 @@ def test_rc_default_products_registered(): ], ) def test_trace_sampling_rules_conversion(rc_rules, expected_config_rules, expected_sampling_rules): + # import pdb; pdb.set_trace() trace_sampling_rules = config.convert_rc_trace_sampling_rules(rc_rules) assert trace_sampling_rules == expected_config_rules