From 314a36be2d3c02ddc5113725d6d8c3579c48a844 Mon Sep 17 00:00:00 2001
From: Ilya Sukhanov <ilya@sukhanov.net>
Date: Thu, 6 Sep 2018 17:03:44 -0600
Subject: [PATCH] In psycopg2 connect accept database name either as dbname or
 database

---
 aws_xray_sdk/ext/psycopg2/patch.py  | 15 +++++++++------
 tests/ext/psycopg2/test_psycopg2.py | 26 ++++++++++++++++++++++++++
 2 files changed, 35 insertions(+), 6 deletions(-)

diff --git a/aws_xray_sdk/ext/psycopg2/patch.py b/aws_xray_sdk/ext/psycopg2/patch.py
index f6f1876d..c6d8d1b8 100644
--- a/aws_xray_sdk/ext/psycopg2/patch.py
+++ b/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,14 +17,16 @@ 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]
+    parameterized_dsn = { c[0]: c[-1] for c in map(methodcaller('split', '='), conn.dsn.split(' '))}
     meta = {
         'database_type': 'PostgreSQL',
-        'url': 'postgresql://{}@{}:{}/{}'.format(user, host, port, dbname),
-        'user': user,
+        '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'
     }
diff --git a/tests/ext/psycopg2/test_psycopg2.py b/tests/ext/psycopg2/test_psycopg2.py
index dd394990..941d8675 100644
--- a/tests/ext/psycopg2/test_psycopg2.py
+++ b/tests/ext/psycopg2/test_psycopg2.py
@@ -47,6 +47,32 @@ def test_execute_dsn_kwargs():
     assert sql['database_version']
 
 
+def test_execute_dsn_kwargs_alt_dbname():
+    """
+    Psycopg supports database to be passed as `database` or `dbname`
+    """
+    q = 'SELECT 1'
+
+    with testing.postgresql.Postgresql() as postgresql:
+        url = postgresql.url()
+        dsn = postgresql.dsn()
+        conn = psycopg2.connect(database=dsn['database'],
+                                user=dsn['user'],
+                                password='',
+                                host=dsn['host'],
+                                port=dsn['port'])
+        cur = conn.cursor()
+        cur.execute(q)
+
+    subsegment = xray_recorder.current_segment().subsegments[0]
+    assert subsegment.name == 'execute'
+    sql = subsegment.sql
+    assert sql['database_type'] == 'PostgreSQL'
+    assert sql['user'] == dsn['user']
+    assert sql['url'] == url
+    assert sql['database_version']
+
+
 def test_execute_dsn_string():
     q = 'SELECT 1'
     with testing.postgresql.Postgresql() as postgresql: