Fix sqlalchemy_core patch errors for unencoded special characters in db url #416. (#418)

* Fix unit test and integration test
* Fix issues #416

Author: vastin

.github/workflows/UnitTesting.yaml

@@ -34,7 +34,10 @@ jobs:
         run: |
           sudo /etc/init.d/mysql start
           mysql -e 'CREATE DATABASE ${{ env.DB_DATABASE }};' -u${{ env.DB_USER }} -p${{ env.DB_PASSWORD }}
-
+          mysql -e 'CREATE DATABASE test_dburl;' -u${{ env.DB_USER }} -p${{ env.DB_PASSWORD }}
+          mysql -e "CREATE USER test_dburl_user@localhost IDENTIFIED BY 'test]password';" -u${{ env.DB_USER }} -p${{ env.DB_PASSWORD }}
+          mysql -e "GRANT ALL PRIVILEGES ON test_dburl.* TO test_dburl_user@localhost;" -u${{ env.DB_USER }} -p${{ env.DB_PASSWORD }}
+          mysql -e "FLUSH PRIVILEGES;" -u${{ env.DB_USER }} -p${{ env.DB_PASSWORD }}
       - name: Setup Python
         uses: actions/setup-python@v4
         with:

aws_xray_sdk/ext/sqlalchemy_core/patch.py

@@ -1,6 +1,6 @@
 import logging
 import sys
-from urllib.parse import urlparse, uses_netloc
+from urllib.parse import urlparse, uses_netloc, quote_plus
 
 import wrapt
 from sqlalchemy.sql.expression import ClauseElement
@@ -14,18 +14,21 @@
 def _sql_meta(engine_instance, args):
     try:
         metadata = {}
-        url = urlparse(str(engine_instance.engine.url))
+        # Workaround for https://github.com/sqlalchemy/sqlalchemy/issues/10662
+        # sqlalchemy.engine.url.URL's __repr__ does not url encode username nor password.
+        # This will continue to work once sqlalchemy fixes the bug.
+        sa_url = engine_instance.engine.url
+        username = sa_url.username
+        sa_url = sa_url._replace(username=None, password=None)
+        url = urlparse(str(sa_url))
+        name = url.netloc
+        if username:
+            # Restore url encoded username
+            quoted_username = quote_plus(username)
+            url = url._replace(netloc='{}@{}'.format(quoted_username, url.netloc))
         # Add Scheme to uses_netloc or // will be missing from url.
         uses_netloc.append(url.scheme)
-        if url.password is None:
-            metadata['url'] = url.geturl()
-            name = url.netloc
-        else:
-            # Strip password from URL
-            host_info = url.netloc.rpartition('@')[-1]
-            parts = url._replace(netloc='{}@{}'.format(url.username, host_info))
-            metadata['url'] = parts.geturl()
-            name = host_info
+        metadata['url'] = url.geturl()
         metadata['user'] = url.username
         metadata['database_type'] = engine_instance.engine.name
         try:

sample-apps/flask/requirements.txt

@@ -1,10 +1,10 @@
-boto3
-certifi
-chardet
-Flask
-idna
-requests
-urllib3
+boto3==1.34.26
+certifi==2023.11.17
+chardet==5.2.0
+Flask==2.3.3
+idna==3.6
+requests==2.31.0
+urllib3==1.26.18
 Werkzeug==3.0.1
 flask-sqlalchemy==2.5.1
 SQLAlchemy==1.4

terraform/eb.tf

@@ -55,7 +55,7 @@ resource "aws_elastic_beanstalk_application_version" "eb_app_version" {
 resource "aws_elastic_beanstalk_environment" "eb_env" {
   name                = "${var.resource_prefix}-EB-App-Env"
   application         = aws_elastic_beanstalk_application.eb_app.name
-  solution_stack_name = "64bit Amazon Linux 2 v3.1.3 running Python 3.7"
+  solution_stack_name = "64bit Amazon Linux 2 v3.5.10 running Python 3.8"
   tier = "WebServer"
   version_label = aws_elastic_beanstalk_application_version.eb_app_version.name
   cname_prefix = "${var.resource_prefix}-Eb-app-env"

tests/ext/sqlalchemy_core/test_dburl.py

@@ -0,0 +1,48 @@
+from sqlalchemy import create_engine
+import urllib
+import pytest
+
+from aws_xray_sdk.core import xray_recorder, patch
+from aws_xray_sdk.ext.sqlalchemy_core import unpatch
+from aws_xray_sdk.core.context import Context
+
+MYSQL_USER = "test_dburl_user"
+MYSQL_PASSWORD = "test]password"
+MYSQL_HOST = "localhost"
+MYSQL_PORT = 3306
+MYSQL_DB_NAME = "test_dburl"
+
+patch(('sqlalchemy_core',))
+
+@pytest.fixture(autouse=True)
+def construct_ctx():
+    """
+    Clean up context storage on each test run and begin a segment
+    so that later subsegment can be attached. After each test run
+    it cleans up context storage again.
+    """
+    xray_recorder.configure(service='test', sampling=False, context=Context())
+    xray_recorder.clear_trace_entities()
+    xray_recorder.begin_segment('name')
+    yield
+    xray_recorder.clear_trace_entities()
+
+
+def test_db_url_with_special_char():
+    password = urllib.parse.quote_plus(MYSQL_PASSWORD)
+    db_url = f"mysql+pymysql://{MYSQL_USER}:{password}@{MYSQL_HOST}:{MYSQL_PORT}/{MYSQL_DB_NAME}"
+
+    engine = create_engine(db_url)
+
+    conn = engine.connect()
+
+    conn.execute("select 1")
+
+    subsegment = xray_recorder.current_segment().subsegments[-1]
+
+    assert subsegment.name == f"{MYSQL_HOST}:{MYSQL_PORT}"
+    sql = subsegment.sql
+    assert sql['database_type'] == 'mysql'
+    assert sql['user'] == MYSQL_USER
+    assert sql['driver_version'] == 'pymysql'
+    assert sql['database_version']

tox.ini

@@ -56,7 +56,7 @@ deps =
     wrapt
 
     ; Python 3.5+ only deps
-    py{37,38,39,310,311,312}: pytest-asyncio
+    py{37,38,39,310,311,312}: pytest-asyncio == 0.21.1
 
     ; For pkg_resources
     py{37,38,39,310,311,312}: setuptools
@@ -86,13 +86,15 @@ deps =
     ext-sqlalchemy_core: sqlalchemy >=1.0.0,<2.0.0
     ext-sqlalchemy_core: testing.postgresql
     ext-sqlalchemy_core: psycopg2
+    ext-sqlalchemy_core: pymysql >= 1.0.0
+    ext-sqlalchemy_core: cryptography
 
     ext-django-2: Django >=2.0,<3.0
     ext-django-3: Django >=3.0,<4.0
     ext-django-4: Django >=4.0,<5.0
     ext-django: django-fake-model
 
-    py{37,38,39,310,311,312}-ext-pynamodb: pynamodb >=3.3.1
+    py{37,38,39,310,311,312}-ext-pynamodb: pynamodb >=3.3.1,<6.0.0
 
     ext-psycopg2: psycopg2
     ext-psycopg2: testing.postgresql
@@ -101,6 +103,7 @@ deps =
     ext-pg8000: testing.postgresql
 
     py{37,38,39,310,311,312}-ext-pymysql: pymysql >= 1.0.0
+    py{37,38,39,310,311,312}-ext-pymysql: cryptography
 
 setenv =
     DJANGO_SETTINGS_MODULE = tests.ext.django.app.settings