Add option to pass attributes in connection string in SQLAlchemy

Author: hovaesco

README.md

@@ -90,6 +90,7 @@ rows = connection.execute(select(nodes)).fetchall()
 ```
 
 In order to pass additional connection attributes use [connect_args](https://docs.sqlalchemy.org/en/14/core/engines.html#sqlalchemy.create_engine.params.connect_args) method.
+Attributes can be also passed in connection string.
 
 ```python
 from sqlalchemy import create_engine
@@ -101,6 +102,13 @@ engine = create_engine(
       "client_tags": ["tag1", "tag2"]
     }
 )
+
+# or in connection string
+engine = create_engine(
+    'trino://user@localhost:8080/system?'
+    'session_properties={"query_max_run_time": "1d"}'
+    '&client_tags=["tag1", "tag2"]',
+)
 ```
 
 ## Authentications

tests/unit/sqlalchemy/test_dialect.py

@@ -41,6 +41,26 @@ def setup(self):
                     source="trino-rulez"
                 ),
             ),
+            (
+                make_url(
+                    'trino://user@localhost:8080?'
+                    'session_properties={"query_max_run_time": "1d"}'
+                    '&http_headers={"trino": 1}'
+                    '&extra_credential=[("a", "b"), ("c", "d")]'
+                    '&client_tags=[1, "sql"]'),
+                list(),
+                dict(
+                    host="localhost",
+                    port=8080,
+                    catalog="system",
+                    user="user",
+                    source="trino-sqlalchemy",
+                    session_properties={"query_max_run_time": "1d"},
+                    http_headers={"trino": 1},
+                    extra_credential=[("a", "b"), ("c", "d")],
+                    client_tags=[1, "sql"]
+                ),
+            ),
         ],
     )
     def test_create_connect_args(self, url: URL, expected_args: List[Any], expected_kwargs: Dict[str, Any]):

trino/sqlalchemy/dialect.py

@@ -9,6 +9,8 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import json
+from ast import literal_eval
 from textwrap import dedent
 from typing import Any, Dict, List, Mapping, Optional, Sequence, Tuple
 
@@ -105,6 +107,18 @@ def create_connect_args(self, url: URL) -> Tuple[Sequence[Any], Mapping[str, Any
         else:
             kwargs["source"] = "trino-sqlalchemy"
 
+        if "session_properties" in url.query:
+            kwargs["session_properties"] = json.loads(url.query["session_properties"])
+
+        if "http_headers" in url.query:
+            kwargs["http_headers"] = json.loads(url.query["http_headers"])
+
+        if "extra_credential" in url.query:
+            kwargs["extra_credential"] = literal_eval(url.query["extra_credential"])
+
+        if "client_tags" in url.query:
+            kwargs["client_tags"] = json.loads(url.query["client_tags"])
+
         return args, kwargs
 
     def get_columns(self, connection: Connection, table_name: str, schema: str = None, **kw) -> List[Dict[str, Any]]: