Make previous commit backwards compatible persist metadata as it was.

Author: IlyaSukhanov

aws_xray_sdk/ext/psycopg2/patch.py

@@ -1,5 +1,6 @@
 import re
 import wrapt
+from operator import methodcaller
 
 from aws_xray_sdk.ext.dbapi2 import XRayTracedConn
 
@@ -16,9 +17,16 @@ def patch():
 def _xray_traced_connect(wrapped, instance, args, kwargs):
 
     conn = wrapped(*args, **kwargs)
+    parameterized_dsn = { c[0]: c[-1] for c in map(methodcaller('split', '='), conn.dsn.split(' '))}
     meta = {
         'database_type': 'PostgreSQL',
-        'connection_string': conn.dsn,
+        'url': 'postgresql://{}@{}:{}/{}'.format(
+            parameterized_dsn.get('user', 'unknown'),
+            parameterized_dsn.get('host', 'unknown'),
+            parameterized_dsn.get('port', 'unknown'),
+            parameterized_dsn.get('dbname', 'unknown'),
+        ),
+        'user': parameterized_dsn.get('user', 'unknown'),
         'database_version': str(conn.server_version),
         'driver_version': 'Psycopg 2'
     }

tests/ext/psycopg2/test_psycopg2.py

@@ -42,11 +42,8 @@ def test_execute_dsn_kwargs():
     assert subsegment.name == 'execute'
     sql = subsegment.sql
     assert sql['database_type'] == 'PostgreSQL'
-    assert 'user={}'.format(dsn['user']) in sql['connection_string']
-    assert 'host={}'.format(dsn['host']) in sql['connection_string']
-    assert 'dbname={}'.format(dsn['database']) in sql['connection_string']
-    assert 'port={}'.format(dsn['port']) in sql['connection_string']
-    assert 'password=xxx' in sql['connection_string']
+    assert sql['user'] == dsn['user']
+    assert sql['url'] == url
     assert sql['database_version']
 
 
@@ -67,11 +64,8 @@ def test_execute_dsn_string():
     assert subsegment.name == 'execute'
     sql = subsegment.sql
     assert sql['database_type'] == 'PostgreSQL'
-    assert 'user={}'.format(dsn['user']) in sql['connection_string']
-    assert 'host={}'.format(dsn['host']) in sql['connection_string']
-    assert 'dbname={}'.format(dsn['database']) in sql['connection_string']
-    assert 'port={}'.format(dsn['port']) in sql['connection_string']
-    assert 'password=xxx' in sql['connection_string']
+    assert sql['user'] == dsn['user']
+    assert sql['url'] == url
     assert sql['database_version']
 
 
@@ -93,11 +87,8 @@ def test_execute_in_pool():
     assert subsegment.name == 'execute'
     sql = subsegment.sql
     assert sql['database_type'] == 'PostgreSQL'
-    assert 'user={}'.format(dsn['user']) in sql['connection_string']
-    assert 'host={}'.format(dsn['host']) in sql['connection_string']
-    assert 'dbname={}'.format(dsn['database']) in sql['connection_string']
-    assert 'port={}'.format(dsn['port']) in sql['connection_string']
-    assert 'password=xxx' in sql['connection_string']
+    assert sql['user'] == dsn['user']
+    assert sql['url'] == url
     assert sql['database_version']
 
 
@@ -121,11 +112,8 @@ def test_execute_bad_query():
     assert subsegment.name == 'execute'
     sql = subsegment.sql
     assert sql['database_type'] == 'PostgreSQL'
-    assert 'user={}'.format(dsn['user']) in sql['connection_string']
-    assert 'host={}'.format(dsn['host']) in sql['connection_string']
-    assert 'dbname={}'.format(dsn['database']) in sql['connection_string']
-    assert 'port={}'.format(dsn['port']) in sql['connection_string']
-    assert 'password=xxx' in sql['connection_string']
+    assert sql['user'] == dsn['user']
+    assert sql['url'] == url
     assert sql['database_version']
 
     exception = subsegment.cause['exceptions'][0]