Fix mysqljdbc edge case (#61)

koxudaxi · Apr 30, 2020 · a0931a4 · a0931a4
1 parent 5e6ee17
commit a0931a4
Showing 1 changed file with 7 additions and 5 deletions.
diff --git a/local_data_api/resources/resource.py b/local_data_api/resources/resource.py
@@ -170,8 +170,10 @@ def create_connection(
     connection = RESOURCE_METAS[resource_arn].connection_maker(  # type: ignore
         database, **connection_kwargs
     )
-    if hasattr(connection, 'database'):  # pragma: no cover
+    try:
         connection.database = database
+    except AttributeError:  # pragma: no cover
+        pass  # for psycopg2
     return connection
 
 
@@ -210,12 +212,12 @@ def get_resource(
     else:
         connection = get_connection(transaction_id)
         if database:
-            if hasattr(connection, 'database'):  # pragma: no cover
+            try:
                 connected_database: Optional[str] = connection.database
-            else:
-                connected_database = connection.get_dsn_parameters()[
+            except AttributeError:  # pragma: no cover
+                connected_database = connection.get_dsn_parameters()[  # for psycopg2
                     'dbname'
-                ]  # pragma: no cover
+                ]
             if database != connected_database:  # pragma: no cover
                 raise BadRequestException(
                     'Database name is not the same as when transaction was created'