Implement Contract Test for Psychopg2 (#72)

*Description of changes:*
Implement Contract Test for Psychopg2

By submitting this pull request, I confirm that you can use, modify,
copy, and redistribute this contribution, under the terms of your
choice.

---------

Co-authored-by: ADOT Patch workflow <[email protected]>
Co-authored-by: Thomas Pierce <[email protected]>

Author: 3 people

contract-tests/images/applications/psycopg2/Dockerfile

@@ -0,0 +1,15 @@
+# Meant to be run from aws-otel-python-instrumentation/contract-tests.
+# Assumes existence of dist/aws_opentelemetry_distro-<pkg_version>-py3-none-any.whl.
+# Assumes filename of aws_opentelemetry_distro-<pkg_version>-py3-none-any.whl is passed in as "DISTRO" arg.
+FROM python:3.10
+WORKDIR /psycopg2
+COPY ./dist/$DISTRO /psycopg2
+COPY ./contract-tests/images/applications/psycopg2 /psycopg2
+
+ENV PIP_ROOT_USER_ACTION=ignore
+ARG DISTRO
+RUN pip install --upgrade pip && pip install -r requirements.txt && pip install ${DISTRO} --force-reinstall
+RUN opentelemetry-bootstrap -a install
+
+# Without `-u`, logs will be buffered and `wait_for_logs` will never return.
+CMD ["opentelemetry-instrument", "python", "-u", "./psycopg2_server.py"]

contract-tests/images/applications/psycopg2/psycopg2_server.py

@@ -0,0 +1,96 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+import atexit
+import os
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+from threading import Thread
+from typing import Tuple
+
+import psycopg2
+from typing_extensions import override
+
+_PORT: int = 8080
+_SUCCESS: str = "success"
+_ERROR: str = "error"
+_FAULT: str = "fault"
+
+_DB_HOST = os.getenv("DB_HOST")
+_DB_USER = os.getenv("DB_USER")
+_DB_PASS = os.getenv("DB_PASS")
+_DB_NAME = os.getenv("DB_NAME")
+
+
+def prepare_database() -> None:
+    conn = psycopg2.connect(dbname=_DB_NAME, user=_DB_USER, password=_DB_PASS, host=_DB_HOST)
+    cur = conn.cursor()
+    cur.execute("DROP TABLE IF EXISTS test_table")
+    cur.execute(
+        """
+        CREATE TABLE test_table (
+            id SERIAL PRIMARY KEY,
+            name TEXT NOT NULL
+            )
+        """
+    )
+
+    cur.execute("INSERT INTO test_table (name) VALUES (%s)", ("Alice",))
+    cur.execute("INSERT INTO test_table (name) VALUES (%s)", ("Bob",))
+
+    conn.commit()
+
+    cur.close()
+    conn.close()
+
+
+class RequestHandler(BaseHTTPRequestHandler):
+    @override
+    # pylint: disable=invalid-name
+    def do_GET(self):
+        status_code: int = 200
+        conn = psycopg2.connect(dbname=_DB_NAME, user=_DB_USER, password=_DB_PASS, host=_DB_HOST)
+        if self.in_path(_SUCCESS):
+            cur = conn.cursor()
+            cur.execute("SELECT id, name FROM test_table")
+            rows = cur.fetchall()
+            cur.close()
+            if len(rows) == 2:
+                status_code = 200
+            else:
+                status_code = 400
+        elif self.in_path(_FAULT):
+            cur = conn.cursor()
+            try:
+                cur.execute("SELECT DISTINCT id, name FROM invalid_table")
+            except psycopg2.ProgrammingError as exception:
+                print("Expected Exception with Invalid SQL occurred:", exception)
+                status_code = 500
+            except Exception as exception:  # pylint: disable=broad-except
+                print("Exception Occurred:", exception)
+            else:
+                status_code = 200
+            finally:
+                cur.close()
+        else:
+            status_code = 404
+        conn.close()
+        self.send_response_only(status_code)
+        self.end_headers()
+
+    def in_path(self, sub_path: str):
+        return sub_path in self.path
+
+
+def main() -> None:
+    prepare_database()
+    server_address: Tuple[str, int] = ("0.0.0.0", _PORT)
+    request_handler_class: type = RequestHandler
+    requests_server: ThreadingHTTPServer = ThreadingHTTPServer(server_address, request_handler_class)
+    atexit.register(requests_server.shutdown)
+    server_thread: Thread = Thread(target=requests_server.serve_forever)
+    server_thread.start()
+    print("Ready")
+    server_thread.join()
+
+
+if __name__ == "__main__":
+    main()

contract-tests/images/applications/psycopg2/pyproject.toml

@@ -0,0 +1,6 @@
+[project]
+name = "psycopg2-server"
+description = "Simple server that relies on psycopg2 library"
+version = "1.0.0"
+license = "Apache-2.0"
+requires-python = ">=3.8"

contract-tests/images/applications/psycopg2/requirements.txt

@@ -0,0 +1,4 @@
+opentelemetry-distro==0.43b0
+opentelemetry-exporter-otlp-proto-grpc==1.22.0
+typing-extensions==4.9.0
+psycopg2==2.9.9

contract-tests/images/applications/requests/Dockerfile

@@ -1,7 +1,7 @@
 # Meant to be run from aws-otel-python-instrumentation/contract-tests.
 # Assumes existence of dist/aws_opentelemetry_distro-<pkg_version>-py3-none-any.whl.
 # Assumes filename of aws_opentelemetry_distro-<pkg_version>-py3-none-any.whl is passed in as "DISTRO" arg.
-FROM public.ecr.aws/docker/library/python:3.11-slim
+FROM python:3.10
 WORKDIR /requests
 COPY ./dist/$DISTRO /requests
 COPY ./contract-tests/images/applications/requests /requests

contract-tests/images/mock-collector/Dockerfile

@@ -1,4 +1,4 @@
-FROM public.ecr.aws/docker/library/python:3.11-slim
+FROM python:3.10
 WORKDIR /mock-collector
 COPY . /mock-collector
 

contract-tests/tests/test/amazon/base/contract_test_base.py

@@ -7,11 +7,15 @@
 from docker import DockerClient
 from docker.models.networks import Network, NetworkCollection
 from docker.types import EndpointConfig
-from mock_collector_client import MockCollectorClient
+from mock_collector_client import MockCollectorClient, ResourceScopeMetric, ResourceScopeSpan
+from requests import Response, request
 from testcontainers.core.container import DockerContainer
 from testcontainers.core.waiting_utils import wait_for_logs
 from typing_extensions import override
 
+from amazon.utils.app_signals_constants import ERROR_METRIC, FAULT_METRIC, LATENCY_METRIC
+from opentelemetry.proto.common.v1.common_pb2 import AnyValue
+
 NETWORK_NAME: str = "aws-appsignals-network"
 
 _logger: Logger = getLogger(__name__)
@@ -115,6 +119,45 @@ def tear_down(self) -> None:
 
         self.mock_collector_client.clear_signals()
 
+    def do_test_requests(
+        self, path: str, method: str, status_code: int, expected_error: int, expected_fault: int, **kwargs
+    ) -> None:
+        address: str = self.application.get_container_host_ip()
+        port: str = self.application.get_exposed_port(self.get_application_port())
+        url: str = f"http://{address}:{port}/{path}"
+        response: Response = request(method, url, timeout=20)
+
+        self.assertEqual(status_code, response.status_code)
+
+        resource_scope_spans: List[ResourceScopeSpan] = self.mock_collector_client.get_traces()
+        self._assert_aws_span_attributes(resource_scope_spans, path, **kwargs)
+        self._assert_semantic_conventions_span_attributes(resource_scope_spans, method, path, status_code, **kwargs)
+
+        metrics: List[ResourceScopeMetric] = self.mock_collector_client.get_metrics(
+            {LATENCY_METRIC, ERROR_METRIC, FAULT_METRIC}
+        )
+        self._assert_metric_attributes(metrics, LATENCY_METRIC, 5000, **kwargs)
+        self._assert_metric_attributes(metrics, ERROR_METRIC, expected_error, **kwargs)
+        self._assert_metric_attributes(metrics, FAULT_METRIC, expected_fault, **kwargs)
+
+    def _assert_str_attribute(self, attributes_dict: Dict[str, AnyValue], key: str, expected_value: str):
+        self.assertIn(key, attributes_dict)
+        actual_value: AnyValue = attributes_dict[key]
+        self.assertIsNotNone(actual_value)
+        self.assertEqual(expected_value, actual_value.string_value)
+
+    def _assert_int_attribute(self, attributes_dict: Dict[str, AnyValue], key: str, expected_value: int) -> None:
+        self.assertIn(key, attributes_dict)
+        actual_value: AnyValue = attributes_dict[key]
+        self.assertIsNotNone(actual_value)
+        self.assertEqual(expected_value, actual_value.int_value)
+
+    def check_sum(self, metric_name: str, actual_sum: float, expected_sum: float) -> None:
+        if metric_name is LATENCY_METRIC:
+            self.assertTrue(0 < actual_sum < expected_sum)
+        else:
+            self.assertEqual(actual_sum, expected_sum)
+
     # pylint: disable=no-self-use
     # Methods that should be overridden in subclasses
     @classmethod
@@ -145,3 +188,16 @@ def get_application_otel_service_name(self) -> str:
 
     def get_application_otel_resource_attributes(self) -> str:
         return "service.name=" + self.get_application_otel_service_name()
+
+    def _assert_aws_span_attributes(self, resource_scope_spans: List[ResourceScopeSpan], path: str, **kwargs):
+        self.fail("Tests must implement this function")
+
+    def _assert_semantic_conventions_span_attributes(
+        self, resource_scope_spans: List[ResourceScopeSpan], method: str, path: str, status_code: int, **kwargs
+    ):
+        self.fail("Tests must implement this function")
+
+    def _assert_metric_attributes(
+        self, resource_scope_metrics: List[ResourceScopeMetric], metric_name: str, expected_sum: int, **kwargs
+    ):
+        self.fail("Tests must implement this function")

contract-tests/tests/test/amazon/psycopg2/psycopg2_test.py

@@ -0,0 +1,152 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+from typing import Dict, List
+
+from mock_collector_client import ResourceScopeMetric, ResourceScopeSpan
+from testcontainers.postgres import PostgresContainer
+from typing_extensions import override
+
+from amazon.base.contract_test_base import NETWORK_NAME, ContractTestBase
+from amazon.utils.app_signals_constants import (
+    AWS_LOCAL_OPERATION,
+    AWS_LOCAL_SERVICE,
+    AWS_REMOTE_OPERATION,
+    AWS_REMOTE_SERVICE,
+    AWS_SPAN_KIND,
+)
+from opentelemetry.proto.common.v1.common_pb2 import AnyValue, KeyValue
+from opentelemetry.proto.metrics.v1.metrics_pb2 import ExponentialHistogramDataPoint, Metric
+from opentelemetry.proto.trace.v1.trace_pb2 import Span
+from opentelemetry.trace import StatusCode
+
+
+class Psycopg2Test(ContractTestBase):
+    @override
+    @classmethod
+    def set_up_dependency_container(cls) -> None:
+        cls.container = (
+            PostgresContainer(user="dbuser", password="example", dbname="postgres")
+            .with_kwargs(network=NETWORK_NAME)
+            .with_name("mydb")
+        )
+        cls.container.start()
+
+    @override
+    @classmethod
+    def tear_down_dependency_container(cls) -> None:
+        cls.container.stop()
+
+    @override
+    def get_application_extra_environment_variables(self) -> Dict[str, str]:
+        return {
+            "DB_HOST": "mydb",
+            "DB_USER": "dbuser",
+            "DB_PASS": "example",
+            "DB_NAME": "postgres",
+        }
+
+    @override
+    def get_application_image_name(self) -> str:
+        return "aws-appsignals-tests-psycopg2-app"
+
+    def test_success(self) -> None:
+        self.mock_collector_client.clear_signals()
+        self.do_test_requests("success", "GET", 200, 0, 0, sql_command="SELECT")
+
+    def test_fault(self) -> None:
+        self.mock_collector_client.clear_signals()
+        self.do_test_requests("fault", "GET", 500, 0, 1, sql_command="SELECT DISTINCT")
+
+    @override
+    def _assert_aws_span_attributes(self, resource_scope_spans: List[ResourceScopeSpan], path: str, **kwargs) -> None:
+        target_spans: List[Span] = []
+        for resource_scope_span in resource_scope_spans:
+            # pylint: disable=no-member
+            if resource_scope_span.span.kind == Span.SPAN_KIND_CLIENT:
+                target_spans.append(resource_scope_span.span)
+
+        self.assertEqual(len(target_spans), 1)
+        self._assert_aws_attributes(target_spans[0].attributes, **kwargs)
+
+    @override
+    def _assert_aws_attributes(self, attributes_list: List[KeyValue], **kwargs) -> None:
+        attributes_dict: Dict[str, AnyValue] = self._get_attributes_dict(attributes_list)
+        self._assert_str_attribute(attributes_dict, AWS_LOCAL_SERVICE, self.get_application_otel_service_name())
+        # InternalOperation as OTEL does not instrument the basic server we are using, so the client span is a local
+        # root.
+        self._assert_str_attribute(attributes_dict, AWS_LOCAL_OPERATION, "InternalOperation")
+        self._assert_str_attribute(attributes_dict, AWS_REMOTE_SERVICE, "postgresql")
+        command: str = kwargs.get("sql_command")
+        self._assert_str_attribute(attributes_dict, AWS_REMOTE_OPERATION, f"{command}")
+        # See comment above AWS_LOCAL_OPERATION
+        self._assert_str_attribute(attributes_dict, AWS_SPAN_KIND, "LOCAL_ROOT")
+
+    def _get_attributes_dict(self, attributes_list: List[KeyValue]) -> Dict[str, AnyValue]:
+        attributes_dict: Dict[str, AnyValue] = {}
+        for attribute in attributes_list:
+            key: str = attribute.key
+            value: AnyValue = attribute.value
+            if key in attributes_dict:
+                old_value: AnyValue = attributes_dict[key]
+                self.fail(f"Attribute {key} unexpectedly duplicated. Value 1: {old_value} Value 2: {value}")
+            attributes_dict[key] = value
+        return attributes_dict
+
+    @override
+    def _assert_semantic_conventions_span_attributes(
+        self, resource_scope_spans: List[ResourceScopeSpan], method: str, path: str, status_code: int, **kwargs
+    ) -> None:
+        target_spans: List[Span] = []
+        for resource_scope_span in resource_scope_spans:
+            # pylint: disable=no-member
+            if resource_scope_span.span.kind == Span.SPAN_KIND_CLIENT:
+                target_spans.append(resource_scope_span.span)
+
+        self.assertEqual(target_spans[0].name, kwargs.get("sql_command").split()[0])
+        if status_code == 200:
+            self.assertEqual(target_spans[0].status.code, StatusCode.UNSET.value)
+        else:
+            self.assertEqual(target_spans[0].status.code, StatusCode.ERROR.value)
+
+        self._assert_semantic_conventions_attributes(target_spans[0].attributes, kwargs.get("sql_command"))
+
+    def _assert_semantic_conventions_attributes(self, attributes_list: List[KeyValue], command: str) -> None:
+        attributes_dict: Dict[str, AnyValue] = self._get_attributes_dict(attributes_list)
+        self.assertTrue(attributes_dict.get("db.statement").string_value.startswith(command))
+        self._assert_str_attribute(attributes_dict, "db.system", "postgresql")
+        self._assert_str_attribute(attributes_dict, "db.name", "postgres")
+        self.assertTrue("db.operation" not in attributes_dict)
+
+    @override
+    def _assert_metric_attributes(
+        self, resource_scope_metrics: List[ResourceScopeMetric], metric_name: str, expected_sum: int, **kwargs
+    ) -> None:
+        target_metrics: List[Metric] = []
+        for resource_scope_metric in resource_scope_metrics:
+            if resource_scope_metric.metric.name.lower() == metric_name.lower():
+                target_metrics.append(resource_scope_metric.metric)
+
+        self.assertEqual(len(target_metrics), 1)
+        target_metric: Metric = target_metrics[0]
+        dp_list: List[ExponentialHistogramDataPoint] = target_metric.exponential_histogram.data_points
+
+        self.assertEqual(len(dp_list), 2)
+        dependency_dp: ExponentialHistogramDataPoint = dp_list[0]
+        service_dp: ExponentialHistogramDataPoint = dp_list[1]
+        if len(dp_list[1].attributes) > len(dp_list[0].attributes):
+            dependency_dp = dp_list[1]
+            service_dp = dp_list[0]
+        attribute_dict: Dict[str, AnyValue] = self._get_attributes_dict(dependency_dp.attributes)
+        self._assert_str_attribute(attribute_dict, AWS_LOCAL_SERVICE, self.get_application_otel_service_name())
+        # See comment on AWS_LOCAL_OPERATION in _assert_aws_attributes
+        self._assert_str_attribute(attribute_dict, AWS_LOCAL_OPERATION, "InternalOperation")
+        self._assert_str_attribute(attribute_dict, AWS_REMOTE_SERVICE, "postgresql")
+        self._assert_str_attribute(attribute_dict, AWS_REMOTE_OPERATION, kwargs.get("sql_command"))
+        self._assert_str_attribute(attribute_dict, AWS_SPAN_KIND, "CLIENT")
+        self.check_sum(metric_name, dependency_dp.sum, expected_sum)
+
+        attribute_dict: Dict[str, AnyValue] = self._get_attributes_dict(service_dp.attributes)
+        # See comment on AWS_LOCAL_OPERATION in _assert_aws_attributes
+        self._assert_str_attribute(attribute_dict, AWS_LOCAL_OPERATION, "InternalOperation")
+        self._assert_str_attribute(attribute_dict, AWS_SPAN_KIND, "LOCAL_ROOT")
+        self.check_sum(metric_name, service_dp.sum, expected_sum)