Support China region (#164)

Author: nirLumigo

requirements.txt

@@ -9,3 +9,4 @@ requests==2.24.0
 pymongo==3.11.0
 redispy==3.0.0
 sqlalchemy==1.3.20
+moto==1.3.16

src/lumigo_tracer/lumigo_utils.py

@@ -9,11 +9,17 @@
 import time
 import http.client
 from collections import OrderedDict
+import random
 from typing import Union, List, Optional, Dict, Any, Tuple, Pattern
 from contextlib import contextmanager
 from base64 import b64encode
 import inspect
 
+try:
+    import boto3
+except Exception:
+    boto3 = None
+
 EXECUTION_TAGS_KEY = "lumigo_execution_tags_no_scrub"
 EDGE_HOST = "{region}.lumigo-tracer-edge.golumigo.com"
 EDGE_PATH = "/api/spans"
@@ -56,6 +62,8 @@
 LUMIGO_TOKEN_KEY = "LUMIGO_TRACER_TOKEN"
 KILL_SWITCH = "LUMIGO_SWITCH_OFF"
 ERROR_SIZE_LIMIT_MULTIPLIER = 2
+CHINA_REGION = "cn-northwest-1"
+EDGE_KINESIS_STREAM_NAME = "prod_trc-inges-edge_edge-kinesis-stream"
 
 _logger: Dict[str, logging.Logger] = {}
 
@@ -84,6 +92,9 @@ class Configuration:
     domains_scrubber: Optional[List] = None
     max_entry_size: int = DEFAULT_MAX_ENTRY_SIZE
     get_key_depth: int = DEFAULT_KEY_DEPTH
+    edge_kinesis_stream_name: str = EDGE_KINESIS_STREAM_NAME
+    edge_kinesis_aws_access_key_id: Optional[str] = None
+    edge_kinesis_aws_secret_access_key: Optional[str] = None
 
     @staticmethod
     def get_max_entry_size(has_error: bool = False) -> int:
@@ -104,6 +115,9 @@ def config(
     domains_scrubber: Optional[List[str]] = None,
     max_entry_size: int = DEFAULT_MAX_ENTRY_SIZE,
     get_key_depth: int = None,
+    edge_kinesis_stream_name: Optional[str] = None,
+    edge_kinesis_aws_access_key_id: Optional[str] = None,
+    edge_kinesis_aws_secret_access_key: Optional[str] = None,
 ) -> None:
     """
     This function configure the lumigo wrapper.
@@ -120,6 +134,9 @@ def config(
     :param domains_scrubber: List of regexes. We will not collect data of requests with hosts that match it.
     :param max_entry_size: The maximum size of each entry when sending back the events.
     :param get_key_depth: Max depth to search the lumigo key in the event (relevant to step functions). default 4.
+    :param edge_kinesis_stream_name: The name of the Kinesis to push the spans in China region
+    :param edge_kinesis_aws_access_key_id: The credentials to push to the Kinesis in China region
+    :param edge_kinesis_aws_secret_access_key: The credentials to push to the Kinesis in China region
     """
     if should_report is not None:
         Configuration.should_report = should_report
@@ -162,6 +179,18 @@ def config(
         domains_scrubber_regex = DOMAIN_SCRUBBER_REGEXES
     Configuration.domains_scrubber = [re.compile(r, re.IGNORECASE) for r in domains_scrubber_regex]
     Configuration.max_entry_size = int(os.environ.get("LUMIGO_MAX_ENTRY_SIZE", max_entry_size))
+    Configuration.edge_kinesis_stream_name = (
+        edge_kinesis_stream_name
+        or os.environ.get("LUMIGO_EDGE_KINESIS_STREAM_NAME")  # noqa`
+        or EDGE_KINESIS_STREAM_NAME  # noqa
+    )
+    Configuration.edge_kinesis_aws_access_key_id = edge_kinesis_aws_access_key_id or os.environ.get(
+        "LUMIGO_EDGE_KINESIS_AWS_ACCESS_KEY_ID"
+    )
+    Configuration.edge_kinesis_aws_secret_access_key = (
+        edge_kinesis_aws_secret_access_key
+        or os.environ.get("LUMIGO_EDGE_KINESIS_AWS_SECRET_ACCESS_KEY")  # noqa
+    )
 
 
 def _is_span_has_error(span: dict) -> bool:
@@ -235,9 +264,19 @@ def report_json(region: Union[None, str], msgs: List[dict], should_retry: bool =
     :return: The duration of reporting (in milliseconds),
                 or 0 if we didn't send (due to configuration or fail).
     """
-    global edge_connection
+    if not Configuration.should_report:
+        return 0
     get_logger().info(f"reporting the messages: {msgs[:10]}")
+    try:
+        prune_trace: bool = not os.environ.get("LUMIGO_PRUNE_TRACE_OFF", "").lower() == "true"
+        to_send = _create_request_body(msgs, prune_trace).encode()
+    except Exception as e:
+        get_logger().exception("Failed to create request: A span was lost.", exc_info=e)
+        return 0
+    if region == CHINA_REGION:
+        return _publish_spans_to_kinesis(to_send, CHINA_REGION)
     host = None
+    global edge_connection
     with lumigo_safe_execute("report json: establish connection"):
         host = prepare_host(Configuration.host or EDGE_HOST.format(region=region))
         duration = 0
@@ -246,28 +285,67 @@ def report_json(region: Union[None, str], msgs: List[dict], should_retry: bool =
             if not edge_connection:
                 get_logger().warning("Can not establish connection. Skip sending span.")
                 return duration
-    if Configuration.should_report:
-        try:
-            prune_trace: bool = not os.environ.get("LUMIGO_PRUNE_TRACE_OFF", "").lower() == "true"
-            to_send = _create_request_body(msgs, prune_trace).encode()
-            start_time = time.time()
-            edge_connection.request(
-                "POST", EDGE_PATH, to_send, headers={"Content-Type": "application/json"}
-            )
-            response = edge_connection.getresponse()
-            response.read()  # We most read the response to keep the connection available
-            duration = int((time.time() - start_time) * 1000)
-            get_logger().info(f"successful reporting, code: {getattr(response, 'code', 'unknown')}")
-        except Exception as e:
-            if should_retry:
-                get_logger().exception(f"Could not report to {host}. Retrying.", exc_info=e)
-                edge_connection = establish_connection(host)
-                report_json(region, msgs, should_retry=False)
-            else:
-                get_logger().exception("Could not report: A span was lost.", exc_info=e)
+    try:
+        start_time = time.time()
+        edge_connection.request(
+            "POST", EDGE_PATH, to_send, headers={"Content-Type": "application/json"}
+        )
+        response = edge_connection.getresponse()
+        response.read()  # We most read the response to keep the connection available
+        duration = int((time.time() - start_time) * 1000)
+        get_logger().info(f"successful reporting, code: {getattr(response, 'code', 'unknown')}")
+    except Exception as e:
+        if should_retry:
+            get_logger().exception(f"Could not report to {host}. Retrying.", exc_info=e)
+            edge_connection = establish_connection(host)
+            report_json(region, msgs, should_retry=False)
+        else:
+            get_logger().exception("Could not report: A span was lost.", exc_info=e)
     return duration
 
 
+def _publish_spans_to_kinesis(to_send: bytes, region: str) -> int:
+    start_time = time.time()
+    with lumigo_safe_execute("report json: publish to kinesis"):
+        get_logger().info("Sending spans to Kinesis")
+        if not Configuration.edge_kinesis_aws_access_key_id:
+            get_logger().error("Missing edge_kinesis_aws_access_key_id, can't publish the spans")
+            return 0
+        if not Configuration.edge_kinesis_aws_secret_access_key:
+            get_logger().error(
+                "Missing edge_kinesis_aws_secret_access_key, can't publish the spans"
+            )
+            return 0
+        _send_data_to_kinesis(
+            stream_name=Configuration.edge_kinesis_stream_name,
+            to_send=to_send,
+            region=region,
+            aws_access_key_id=Configuration.edge_kinesis_aws_access_key_id,
+            aws_secret_access_key=Configuration.edge_kinesis_aws_secret_access_key,
+        )
+    return int((time.time() - start_time) * 1000)
+
+
+def _send_data_to_kinesis(
+    stream_name: str,
+    to_send: bytes,
+    region: str,
+    aws_access_key_id: str,
+    aws_secret_access_key: str,
+):
+    if not boto3:
+        get_logger().error("boto3 is missing. Unable to send to Kinesis.")
+        return None
+    client = boto3.client(
+        "kinesis",
+        region_name=region,
+        aws_access_key_id=aws_access_key_id,
+        aws_secret_access_key=aws_secret_access_key,
+    )
+    client.put_record(Data=to_send, StreamName=stream_name, PartitionKey=str(random.random()))
+    get_logger().info("Successful sending to Kinesis")
+
+
 def get_logger(logger_name="lumigo"):
     """
     This function returns lumigo's logger.

src/test/unit/test_lumigo_utils.py

@@ -6,6 +6,7 @@
 from mock import Mock
 
 import pytest
+from lumigo_tracer import lumigo_utils
 from lumigo_tracer.lumigo_utils import (
     _create_request_body,
     _is_span_has_error,
@@ -34,6 +35,7 @@
     is_error_code,
     get_size_upper_bound,
     is_aws_arn,
+    CHINA_REGION,
 )
 import json
 
@@ -387,6 +389,43 @@ def test_report_json_retry(monkeypatch, reporter_mock, caplog, errors, final_log
     assert caplog.records[-1].levelname == final_log
 
 
+def test_report_json_china_missing_access_key_id(monkeypatch, reporter_mock, caplog):
+    monkeypatch.setattr(Configuration, "should_report", True)
+    reporter_mock.side_effect = report_json
+    assert report_json(CHINA_REGION, [{"a": "b"}]) == 0
+    assert any(
+        "edge_kinesis_aws_access_key_id" in record.message and record.levelname == "ERROR"
+        for record in caplog.records
+    )
+
+
+def test_report_json_china_missing_secret_access_key(monkeypatch, reporter_mock, caplog):
+    monkeypatch.setattr(Configuration, "should_report", True)
+    monkeypatch.setattr(Configuration, "edge_kinesis_aws_access_key_id", "my_value")
+    reporter_mock.side_effect = report_json
+    assert report_json(CHINA_REGION, [{"a": "b"}]) == 0
+    assert any(
+        "edge_kinesis_aws_secret_access_key" in record.message and record.levelname == "ERROR"
+        for record in caplog.records
+    )
+
+
+def test_report_json_china_no_boto(monkeypatch, reporter_mock, caplog):
+    reporter_mock.side_effect = report_json
+    monkeypatch.setattr(Configuration, "should_report", True)
+    monkeypatch.setattr(Configuration, "edge_kinesis_aws_access_key_id", "my_value")
+    monkeypatch.setattr(Configuration, "edge_kinesis_aws_secret_access_key", "my_value")
+    monkeypatch.setattr(lumigo_utils, "boto3", None)
+
+    report_json(CHINA_REGION, [{"a": "b"}])
+
+    assert any(
+        "boto3 is missing. Unable to send to Kinesis" in record.message
+        and record.levelname == "ERROR"  # noqa
+        for record in caplog.records
+    )
+
+
 @pytest.mark.parametrize("env, expected", [("True", True), ("other", False), ("123", False)])
 def test_is_kill_switch_on(monkeypatch, env, expected):
     monkeypatch.setenv(KILL_SWITCH, env)

src/test/unit/test_tracer.py

@@ -6,20 +6,26 @@
 import sys
 from functools import wraps
 import logging
+from unittest.mock import MagicMock
 
+import boto3
 import pytest
 from capturer import CaptureOutput
 
 from lumigo_tracer import lumigo_tracer, LumigoChalice, add_execution_tag
+from lumigo_tracer import lumigo_utils
 from lumigo_tracer.lumigo_utils import (
     Configuration,
     STEP_FUNCTION_UID_KEY,
     LUMIGO_EVENT_KEY,
     _create_request_body,
     EXECUTION_TAGS_KEY,
+    report_json,
+    EDGE_KINESIS_STREAM_NAME,
 )
 
 from lumigo_tracer.spans_container import SpansContainer
+from moto import mock_kinesis
 
 
 def test_lambda_wrapper_basic_events(reporter_mock, context):
@@ -364,3 +370,58 @@ def lambda_test_function(event, context):
     # following python's runtime: runtime/lambda_runtime_marshaller.py:27
     expected_message = 'The lambda will probably fail due to bad return value. Original message: "Object of type datetime is not JSON serializable"'
     assert function_span["error"]["message"] == expected_message
+
+
+@mock_kinesis
+def test_china(context, reporter_mock, monkeypatch):
+    china_region_for_test = "ap-east-1"  # Moto doesn't work for China
+    monkeypatch.setattr(lumigo_utils, "CHINA_REGION", china_region_for_test)
+    monkeypatch.setenv("AWS_REGION", china_region_for_test)
+    reporter_mock.side_effect = report_json  # Override the conftest's monkeypatch
+    access_key_id = "my_access_key_id"
+    secret_access_key = "my_secret_access_key"
+    # Create edge Kinesis
+    client = boto3.client(
+        "kinesis",
+        region_name=china_region_for_test,
+        aws_access_key_id=access_key_id,
+        aws_secret_access_key=secret_access_key,
+    )
+    client.create_stream(StreamName=EDGE_KINESIS_STREAM_NAME, ShardCount=1)
+    shard_id = client.describe_stream(StreamName=EDGE_KINESIS_STREAM_NAME)["StreamDescription"][
+        "Shards"
+    ][0]["ShardId"]
+    shard_iterator = client.get_shard_iterator(
+        StreamName=EDGE_KINESIS_STREAM_NAME,
+        ShardId=shard_id,
+        ShardIteratorType="AT_TIMESTAMP",
+        Timestamp=datetime.datetime.utcnow(),
+    )["ShardIterator"]
+
+    original_get_boto_client = boto3.client
+    monkeypatch.setattr(boto3, "client", MagicMock(side_effect=original_get_boto_client))
+
+    @lumigo_tracer(
+        edge_kinesis_aws_access_key_id=access_key_id,
+        edge_kinesis_aws_secret_access_key=secret_access_key,
+        should_report=True,
+    )
+    def lambda_test_function(event, context):
+        return "ret_value"
+
+    event = {"k": "v"}
+    result = lambda_test_function(event, context)
+
+    assert result == "ret_value"
+    # Spans sent to Kinesis
+    records = client.get_records(ShardIterator=shard_iterator)["Records"]
+    assert len(records) == 2  # Start span and end span
+    span_sent = json.loads(records[1]["Data"].decode())[0]
+    assert span_sent["event"] == json.dumps(event)
+    # Used the client from the decorator params
+    boto3.client.assert_called_with(
+        "kinesis",
+        region_name=china_region_for_test,
+        aws_access_key_id=access_key_id,
+        aws_secret_access_key=secret_access_key,
+    )