Repository: elastic/apm-agent-python

Issue:
APM crashes when postgres is configured with multiple hosts / ports in Django
Hello.

I'm using multiple postgresql hosts (https://www.postgresql.org/docs/current/libpq-connect.html#LIBPQ-CONNSTRING section 34.1.1.3).

The APM crashes trying to convert comma separated list of ports in a int:

```
  File "/usr/local/lib/python3.7/site-packages/elasticapm/instrumentation/packages/psycopg2.py", line 86, in call
    if int(port) != default_ports.get("postgresql"):
ValueError: invalid literal for int() with base 10: '5432,5433'
```

**To Reproduce**

1. Configure multiple hosts in postgres config (more than 1 host and port)
2. Start the app
3. crashes with the exception:
```
  File "/usr/local/lib/python3.7/site-packages/elasticapm/instrumentation/packages/psycopg2.py", line 86, in call
    if int(port) != default_ports.get("postgresql"):
ValueError: invalid literal for int() with base 10: '5432,5433'
```

**Environment (please complete the following information)**
- OS: Linux containerd 
- Python version: 3.7
- Framework and version [e.g. Django 2.1]: Django 2.1
- APM Server version: 7.12.1
- Agent version: 6.6.1


**Additional context**

https://github.com/elastic/apm-agent-python/blob/master/elasticapm/instrumentation/packages/psycopg2.py#L86




Discussion:
Hi @azlev. Thanks for opening the issue! This is indeed a bug. The fix looks relatively easy, so we'll try and get out a quick fix until tomorrow.

Resolving patch:
diff --git a/elasticapm/base.py b/elasticapm/base.py
--- a/elasticapm/base.py
+++ b/elasticapm/base.py
@@ -653,11 +653,11 @@ def send(self, url, **kwargs):
         return None
 
 
-def get_client():
+def get_client() -> Client:
     return CLIENT_SINGLETON
 
 
-def set_client(client):
+def set_client(client: Client):
     global CLIENT_SINGLETON
     if CLIENT_SINGLETON:
         logger = get_logger("elasticapm")
diff --git a/elasticapm/instrumentation/packages/psycopg2.py b/elasticapm/instrumentation/packages/psycopg2.py
--- a/elasticapm/instrumentation/packages/psycopg2.py
+++ b/elasticapm/instrumentation/packages/psycopg2.py
@@ -29,6 +29,8 @@
 #  OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 from __future__ import absolute_import
 
+from typing import Optional, Union
+
 from elasticapm.instrumentation.packages.dbapi2 import (
     ConnectionProxy,
     CursorProxy,
@@ -76,22 +78,11 @@ class Psycopg2Instrumentation(DbApi2Instrumentation):
     def call(self, module, method, wrapped, instance, args, kwargs):
         signature = "psycopg2.connect"
 
-        host = kwargs.get("host")
-        if host:
-            signature += " " + compat.text_type(host)
-
-            port = kwargs.get("port")
-            if port:
-                port = str(port)
-                if int(port) != default_ports.get("postgresql"):
-                    host += ":" + port
-            signature += " " + compat.text_type(host)
-        else:
-            # Parse connection string and extract host/port
-            pass
+        host, port = get_destination_info(kwargs.get("host"), kwargs.get("port"))
+        signature = f"{signature} {host}:{port}"
         destination_info = {
-            "address": kwargs.get("host", "localhost"),
-            "port": int(kwargs.get("port", default_ports.get("postgresql"))),
+            "address": host,
+            "port": port,
         }
         with capture_span(
             signature,
@@ -140,3 +131,19 @@ def call(self, module, method, wrapped, instance, args, kwargs):
                 kwargs["scope"] = kwargs["scope"].__wrapped__
 
         return wrapped(*args, **kwargs)
+
+
+def get_destination_info(host: Optional[str], port: Union[None, str, int]) -> tuple:
+    if host:
+        if "," in host:  # multiple hosts defined, take first
+            host = host.split(",")[0]
+    else:
+        host = "localhost"
+    if port:
+        port = str(port)
+        if "," in port:  # multiple ports defined, take first
+            port = port.split(",")[0]
+        port = int(port)
+    else:
+        port = default_ports.get("postgresql")
+    return host, port


Test patch:
diff --git a/tests/fixtures.py b/tests/fixtures.py
--- a/tests/fixtures.py
+++ b/tests/fixtures.py
@@ -201,6 +201,16 @@ def elasticapm_client(request):
     assert not client._transport.validation_errors
 
 
+@pytest.fixture()
+def elasticapm_transaction(elasticapm_client):
+    """
+    Useful fixture if spans from other fixtures should be captured.
+    This can be achieved by listing this fixture first.
+    """
+    transaction = elasticapm_client.begin_transaction("test")
+    yield transaction
+
+
 @pytest.fixture()
 def elasticapm_client_log_file(request):
     original_exceptionhook = sys.excepthook
diff --git a/tests/instrumentation/psycopg2_tests.py b/tests/instrumentation/psycopg2_tests.py
--- a/tests/instrumentation/psycopg2_tests.py
+++ b/tests/instrumentation/psycopg2_tests.py
@@ -31,12 +31,15 @@
 #  OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
 import os
+from typing import cast
 
 import pytest
 
-from elasticapm.conf.constants import TRANSACTION
-from elasticapm.instrumentation.packages.psycopg2 import PGCursorProxy, extract_signature
+from elasticapm import get_client
+from elasticapm.conf.constants import SPAN, TRANSACTION
+from elasticapm.instrumentation.packages.psycopg2 import PGCursorProxy, extract_signature, get_destination_info
 from elasticapm.utils import default_ports
+from tests.fixtures import TempStoreClient
 
 psycopg2 = pytest.importorskip("psycopg2")
 
@@ -162,10 +165,10 @@ def test_select_with_dollar_quotes_custom_token():
 
 
 def test_select_with_difficult_table_name():
-    sql_statement = u"""SELECT id FROM "myta\n-æøåble" WHERE id = 2323"""
+    sql_statement = """SELECT id FROM "myta\n-æøåble" WHERE id = 2323"""
     actual = extract_signature(sql_statement)
 
-    assert u"SELECT FROM myta\n-æøåble" == actual
+    assert "SELECT FROM myta\n-æøåble" == actual
 
 
 def test_select_subselect():
@@ -507,3 +510,30 @@ def test_psycopg2_execute_values(instrument, postgres_connection, elasticapm_cli
         spans = elasticapm_client.spans_for_transaction(transactions[0])
         assert spans[0]["name"] == "INSERT INTO test"
         assert len(spans[0]["context"]["db"]["statement"]) == 10000, spans[0]["context"]["db"]["statement"]
+
+
+@pytest.mark.integrationtest
+def test_psycopg2_connection(instrument, elasticapm_transaction, postgres_connection):
+    # elastciapm_client.events is only available on `TempStoreClient`, this keeps the type checkers happy
+    elasticapm_client = cast(TempStoreClient, get_client())
+    elasticapm_client.end_transaction("test", "success")
+    span = elasticapm_client.events[SPAN][0]
+    host = os.environ.get("POSTGRES_HOST", "localhost")
+    assert span["name"] == f"psycopg2.connect {host}:5432"
+    assert span["action"] == "connect"
+
+
+@pytest.mark.parametrize(
+    "host_in,port_in,expected_host_out,expected_port_out",
+    (
+        (None, None, "localhost", 5432),
+        (None, 5432, "localhost", 5432),
+        ("localhost", "5432", "localhost", 5432),
+        ("foo.bar", "5432", "foo.bar", 5432),
+        ("localhost,foo.bar", "5432,1234", "localhost", 5432),  # multiple hosts
+    ),
+)
+def test_get_destination(host_in, port_in, expected_host_out, expected_port_out):
+    host, port = get_destination_info(host_in, port_in)
+    assert host == expected_host_out
+    assert port == expected_port_out