Make previous commit backwards compatible persist metadata as it was.

Author: IlyaSukhanov

aws_xray_sdk/ext/psycopg2/patch.py

@@ -16,9 +16,14 @@ def patch():
 def _xray_traced_connect(wrapped, instance, args, kwargs):
 
     conn = wrapped(*args, **kwargs)
+    host = kwargs['host'] if 'host' in kwargs else re.search(r'host=(\S+)\b', args[0]).groups()[0]
+    dbname = kwargs['dbname'] if 'dbname' in kwargs else re.search(r'dbname=(\S+)\b', args[0]).groups()[0]
+    port = kwargs['port'] if 'port' in kwargs else re.search(r'port=(\S+)\b', args[0]).groups()[0]
+    user = kwargs['user'] if 'user' in kwargs else re.search(r'user=(\S+)\b', args[0]).groups()[0]
     meta = {
         'database_type': 'PostgreSQL',
-        'connection_string': conn.dsn,
+        'url': 'postgresql://{}@{}:{}/{}'.format(user, host, port, dbname),
+        'user': user,
         '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]