Changes to accommodate for environment changes

- Adding `CoverageSpanFilter` for different types of filter on what is going to be covered
- Adding `UnableToStartProcessorException` for when we are unable to start the tracking
- Adding new flow where we post to create a version ahead of time, and then we can just use the external_id on future calls

Author: ThiagoCodecov

.gitignore

@@ -85,7 +85,7 @@ ipython_config.py
 # pyenv
 #   For a library or package, you might want to ignore these files since the code is
 #   intended to run in multiple environments; otherwise, check them in:
-# .python-version
+.python-version
 
 # pipenv
 #   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.

codecovopentelem/__init__.py

@@ -4,8 +4,9 @@
 import urllib.parse
 from base64 import b64encode
 from decimal import Decimal
+from enum import Enum
 from io import StringIO
-from typing import Optional, Tuple, Pattern
+from typing import Dict, Optional, Pattern, Tuple
 
 import coverage
 import requests
@@ -16,65 +17,81 @@
 log = logging.getLogger("codecovopentelem")
 
 
+class CoverageSpanFilter(Enum):
+    regex_name_filter = "name_regex"
+    span_kind_filter = "span_kind"
+
+
+class UnableToStartProcessorException(Exception):
+    pass
+
+
 class CodecovCoverageStorageManager(object):
-    def __init__(self, writeable_folder: str):
+    def __init__(self, writeable_folder: str, filters: Dict):
         if writeable_folder is None:
             writeable_folder = "/home/codecov"
         self._writeable_folder = writeable_folder
         self.inner = {}
+        self._filters = filters
 
-    def start_cov_for_span(self, span_id):
+    def possibly_start_cov_for_span(self, span) -> bool:
+        span_id = span.context.span_id
+        if self._filters.get(
+            CoverageSpanFilter.regex_name_filter
+        ) and not self._filters.get(CoverageSpanFilter.regex_name_filter).match(
+            span.name
+        ):
+            return False
+        if self._filters.get(
+            CoverageSpanFilter.span_kind_filter
+        ) and span.kind not in self._filters.get(CoverageSpanFilter.span_kind_filter):
+            return False
         cov = coverage.Coverage(data_file=f"{self._writeable_folder}/.{span_id}file")
         self.inner[span_id] = cov
         cov.start()
+        return True
 
-    def stop_cov_for_span(self, span_id):
+    def stop_cov_for_span(self, span):
+        span_id = span.context.span_id
         cov = self.inner.get(span_id)
         if cov is not None:
             cov.stop()
 
-    def pop_cov_for_span(self, span_id):
+    def pop_cov_for_span(self, span):
+        span_id = span.context.span_id
         return self.inner.pop(span_id, None)
 
 
 class CodecovCoverageGenerator(SpanProcessor):
     def __init__(
-        self,
-        cov_storage: CodecovCoverageStorageManager,
-        sample_rate: Decimal,
-        name_regex: Pattern = None,
+        self, cov_storage: CodecovCoverageStorageManager, sample_rate: Decimal,
     ):
         self._cov_storage = cov_storage
         self._sample_rate = sample_rate
-        self._name_regex = name_regex
 
     def _should_profile_span(self, span, parent_context):
-        return random.random() < self._sample_rate and (
-            self._name_regex is None or self._name_regex.match(span.name)
-        )
+        return random.random() < self._sample_rate
 
     def on_start(self, span, parent_context=None):
         if self._should_profile_span(span, parent_context):
-            span_id = span.context.span_id
-            self._cov_storage.start_cov_for_span(span_id)
+            self._cov_storage.possibly_start_cov_for_span(span)
 
     def on_end(self, span):
-        span_id = span.context.span_id
-        self._cov_storage.stop_cov_for_span(span_id)
+        self._cov_storage.stop_cov_for_span(span)
 
 
 class CoverageExporter(SpanExporter):
     def __init__(
         self,
         cov_storage: CodecovCoverageStorageManager,
         repository_token: str,
-        profiling_identifier: str,
+        profiling_id: str,
         codecov_endpoint: str,
         untracked_export_rate: float,
     ):
         self._cov_storage = cov_storage
         self._repository_token = repository_token
-        self._profiling_identifier = profiling_identifier
+        self._profiling_id = profiling_id
         self._codecov_endpoint = codecov_endpoint
         self._untracked_export_rate = untracked_export_rate
 
@@ -96,71 +113,98 @@ def export(self, spans):
         tracked_spans = []
         untracked_spans = []
         for span in spans:
-            span_id = span.context.span_id
-            cov = self._cov_storage.pop_cov_for_span(span_id)
+            cov = self._cov_storage.pop_cov_for_span(span)
             s = json.loads(span.to_json())
             if cov is not None:
                 s["codecov"] = self._load_codecov_dict(span, cov)
                 tracked_spans.append(s)
             else:
                 if random.random() < self._untracked_export_rate:
                     untracked_spans.append(s)
-        if not tracked_spans:
+        if not tracked_spans and not untracked_spans:
             return SpanExportResult.SUCCESS
         url = urllib.parse.urljoin(self._codecov_endpoint, "/profiling/uploads")
-        res = requests.post(
-            url,
-            headers={"Authorization": f"repotoken {self._repository_token}"},
-            json={"profiling": self._profiling_identifier},
-        )
         try:
+            res = requests.post(
+                url,
+                headers={"Authorization": f"repotoken {self._repository_token}"},
+                json={"profiling": self._profiling_id},
+            )
             res.raise_for_status()
-        except requests.HTTPError:
+        except requests.RequestException:
             log.warning("Unable to send profiling data to codecov")
             return SpanExportResult.FAILURE
         location = res.json()["raw_upload_location"]
         requests.put(
             location,
             headers={"Content-Type": "application/txt"},
-            data=json.dumps({"spans": tracked_spans, "untracked": untracked_spans}).encode(),
+            data=json.dumps(
+                {"spans": tracked_spans, "untracked": untracked_spans}
+            ).encode(),
         )
         return SpanExportResult.SUCCESS
 
 
 def get_codecov_opentelemetry_instances(
     repository_token: str,
-    profiling_identifier: str,
     sample_rate: float,
-    name_regex: Optional[Pattern],
+    untracked_export_rate: float,
+    filters: Optional[Dict] = None,
+    profiling_identifier: Optional[str] = None,
+    environment: Optional[str] = None,
+    profiling_id: Optional[str] = None,
     codecov_endpoint: str = None,
     writeable_folder: str = None,
 ) -> Tuple[CodecovCoverageGenerator, CoverageExporter]:
     """
     Entrypoint for getting a span processor/span exporter
         pair for getting profiling data into codecov
 
+    Notice that either `profiling_id` or `profiling_identifier` and `environment` need to be set.
+        If `profiling_id` is set, we just use it directly on the exporter. If not, we will use
+        `profiling_identifier` and `environment` to generate fetch a `profiling_id` from the
+        database
+
     Args:
         repository_token (str): The profiling-capable authentication token
-        profiling_identifier (str): The identifier for what profiling one is doing
         sample_rate (float): The sampling rate for codecov
-        name_regex (Optional[Pattern]): A regex to filter which spans should be
-            sampled
+        untracked_export_rate (float): Description
+        filters (Optional[Dict], optional): A dictionary of filters for determining which
+            spans should have its coverage tracked
+        profiling_identifier (Optional[str], optional): The identifier for what profiling one is doing
+        environment (Optional[str], optional): Which environment this profiling is running on
+        profiling_id (Optional[str], optional): Description
         codecov_endpoint (str, optional): For configuring the endpoint in case
             the user is in enterprise (not supported yet). Default is "https://api.codecov.io/"
         writeable_folder (str, optional): A folder that is guaranteed to be write-able
             in the system. It's only used for temporary files, and nothing is expected
             to live very long in there.
     """
-    if codecov_endpoint is None:
-        codecov_endpoint = "https://api.codecov.io"
-    manager = CodecovCoverageStorageManager(writeable_folder)
-    generator = CodecovCoverageGenerator(manager, sample_rate, name_regex)
-    # untracked rate set to make it so we export roughly as many tracked and untracked spans
-    untracked_export_rate = sample_rate / (1 - sample_rate) if sample_rate < 1 else 0
+    codecov_endpoint = codecov_endpoint or "https://api.codecov.io"
+    if profiling_id is None:
+        if profiling_identifier is None or environment is None:
+            raise UnableToStartProcessorException(
+                "Codecov profiling needs either the id or identifier + environment"
+            )
+        response = requests.post(
+            urllib.parse.urljoin(codecov_endpoint, "/profiling/versions"),
+            json={
+                "version_identifier": profiling_identifier,
+                "environment": environment,
+            },
+            headers={"Authorization": f"repotoken {repository_token}"},
+        )
+        try:
+            response.raise_for_status()
+        except requests.HTTPError:
+            raise UnableToStartProcessorException()
+        profiling_id = response.json()["external_id"]
+    manager = CodecovCoverageStorageManager(writeable_folder, filters or {})
+    generator = CodecovCoverageGenerator(manager, sample_rate)
     exporter = CoverageExporter(
         manager,
         repository_token,
-        profiling_identifier,
+        profiling_id,
         codecov_endpoint,
         untracked_export_rate,
     )

tests/test_interface.py

@@ -0,0 +1,64 @@
+import json
+from uuid import uuid4
+
+import pytest
+import responses
+from coverage import Coverage
+from coverage.xmlreport import XmlReporter
+
+from codecovopentelem import (
+    CodecovCoverageGenerator,
+    CoverageExporter,
+    get_codecov_opentelemetry_instances,
+    UnableToStartProcessorException,
+)
+
+
+@pytest.fixture
+def mocked_responses():
+    with responses.RequestsMock() as rsps:
+        yield rsps
+
+
+def test_get_codecov_opentelemetry_instances_nothing_set(mocker, mocked_responses):
+    with pytest.raises(UnableToStartProcessorException) as exc:
+        get_codecov_opentelemetry_instances(
+            repository_token="repository_token",
+            sample_rate=0.1,
+            untracked_export_rate=0.1,
+        )
+    assert exc.value.args == (
+        "Codecov profiling needs either the id or identifier + environment",
+    )
+
+
+def test_get_codecov_opentelemetry_instances_nothing_set_env_and_version(
+    mocker, mocked_responses
+):
+    uuid = uuid4().hex
+    mocked_responses.add(
+        responses.POST,
+        "https://api.codecov.io/profiling/versions",
+        json={"external_id": uuid},
+        status=200,
+        content_type="application/json",
+        match=[
+            responses.matchers.json_params_matcher(
+                {
+                    "version_identifier": "profiling_identifier",
+                    "environment": "production",
+                }
+            )
+        ],
+    )
+    res = get_codecov_opentelemetry_instances(
+        repository_token="repository_token",
+        sample_rate=0.1,
+        untracked_export_rate=0.1,
+        profiling_identifier="profiling_identifier",
+        environment="production",
+    )
+    assert len(res) == 2
+    generator, exporter = res
+    assert isinstance(generator, CodecovCoverageGenerator)
+    assert isinstance(exporter, CoverageExporter)