Merge pull request #201 from octoenergy/add-sqlalchemy-2-support

Add sqlalchemy 2 support

CHANGELOG.txt

@@ -4,6 +4,11 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [1.2.0] - 2023-04-06
+### Added
+  - Add support for SQLAlchemy 2.0 by correctly handling `sqlalchemy.engine.url.URL`s, which are now immutable.
+  Older versions of tentaclio should pin `sqlalchemy < 2.0` to avoid this issue.
+
 ## [1.1.0] - 2023-04-05
 ### Changed 
   - Credential files error reporting to help users identify the credentials issues

Makefile

@@ -41,10 +41,10 @@ functional-sftp:
 	pipenv run pytest tests/functional/sftp
 
 format:
-	black -l 99 src
-	black -l 99 tests
-	isort -rc src
-	isort -rc tests
+	pipenv run black -l 99 src
+	pipenv run black -l 99 tests
+	pipenv run isort src
+	pipenv run isort tests
 
 # Deployment
 

Pipfile

@@ -9,7 +9,6 @@ python_version = "3.9"
 [dev-packages]
 # Symlink to project root
 tentaclio = {editable = true,path = "."}
-
 # Linting
 black = "*"
 isort = "*"
@@ -18,13 +17,11 @@ mypy = "*"
 pydocstyle = "*"
 types-pyyaml = "*"
 types-requests = "*"
-
 # Testing
 moto = "*"
 pytest = "*"
 pytest-cov = "*"
 pytest-mock = "*"
-
 # Releasing
 twine = "*"
 secretstorage = "*"

setup.py

@@ -8,7 +8,7 @@
 from setuptools.command.install import install
 
 
-VERSION = "1.1.0"
+VERSION = "1.2.0"
 
 REPO_ROOT = pathlib.Path(__file__).parent
 

src/tentaclio/clients/ftp_client.py

@@ -41,7 +41,7 @@ def _connect(self) -> ftplib.FTP:
     # Stream methods:
 
     @decorators.check_conn
-    def get(self, writer: protocols.ByteWriter, file_path: str = None) -> None:
+    def get(self, writer: protocols.ByteWriter, file_path: Optional[str] = None) -> None:
         """Write the contents of a remote file into the passed writer.
 
         Arguments:
@@ -162,7 +162,7 @@ def _connect(self) -> pysftp.Connection:
     # Stream methods:
 
     @decorators.check_conn
-    def get(self, writer: protocols.ByteWriter, file_path: str = None) -> None:
+    def get(self, writer: protocols.ByteWriter, file_path: Optional[str] = None) -> None:
         """Write the contents of a remote file into the passed writer.
 
         Arguments:
@@ -180,7 +180,7 @@ def get(self, writer: protocols.ByteWriter, file_path: str = None) -> None:
         self.conn.getfo(remote_path, writer)
 
     @decorators.check_conn
-    def put(self, reader: protocols.ByteReader, file_path: str = None) -> None:
+    def put(self, reader: protocols.ByteReader, file_path: Optional[str] = None) -> None:
         """Write the contents of the reader into the remote file.
 
         Arguments:

src/tentaclio/clients/http_client.py

@@ -38,7 +38,10 @@ class HTTPClient(base_client.BaseClient["HTTPClient"]):
     endpoint: str
 
     def __init__(
-        self, url: str, default_timeout: float = None, default_headers: dict = None
+        self,
+        url: str,
+        default_timeout: Optional[float] = None,
+        default_headers: Optional[dict] = None,
     ) -> None:
         """Create a new http/https client based on the passed url and extra params."""
         # Default connection timeout at 10''
@@ -81,9 +84,9 @@ def _connect(self) -> requests.Session:
     def get(
         self,
         writer: protocols.ByteWriter,
-        endpoint: str = None,
-        params: dict = None,
-        options: dict = None,
+        endpoint: Optional[str] = None,
+        params: Optional[dict] = None,
+        options: Optional[dict] = None,
     ) -> None:
         """Read the contents from the url and write them into the provided writer.
 
@@ -103,9 +106,9 @@ def get(
     def put(
         self,
         reader: protocols.ByteReader,
-        endpoint: str = None,
-        params: dict = None,
-        options: dict = None,
+        endpoint: Optional[str] = None,
+        params: Optional[dict] = None,
+        options: Optional[dict] = None,
     ) -> None:
         """Write the contents of the provided reader into the url using POST.
 
@@ -132,8 +135,8 @@ def _build_request(
         self,
         method: str,
         url: str,
-        default_data: protocols.Reader = None,
-        default_params: dict = None,
+        default_data: Optional[protocols.Reader] = None,
+        default_params: Optional[dict] = None,
     ):
         data: Union[protocols.Reader, list] = default_data or []
         params = default_params or {}
@@ -149,7 +152,9 @@ def _build_request(
 
         return self.conn.prepare_request(request)
 
-    def _send_request(self, request: requests.PreparedRequest, default_options: dict = None):
+    def _send_request(
+        self, request: requests.PreparedRequest, default_options: Optional[dict] = None
+    ):
         options = default_options or {}
 
         response = self.conn.send(

src/tentaclio/clients/sqla_client.py

@@ -7,8 +7,9 @@
 from typing import Container, Generator, Optional, Union
 
 import pandas as pd
+from sqlalchemy import text
 from sqlalchemy.engine import Connection, CursorResult, Engine, create_engine
-from sqlalchemy.engine import url as sqla_url
+from sqlalchemy.engine.url import URL as sqla_url
 from sqlalchemy.orm import session, sessionmaker
 from sqlalchemy.sql.schema import MetaData
 
@@ -20,7 +21,7 @@
 __all__ = ["SQLAlchemyClient", "bound_session", "atomic_session"]
 
 
-SessionGenerator = Generator[None, session.Session, None]
+SessionGenerator = Generator[session.Session, None, None]
 
 
 class _TrueContainer(Container[str]):
@@ -57,7 +58,10 @@ class SQLAlchemyClient(base_client.BaseClient["SQLAlchemyClient"]):
     port: Optional[int]
 
     def __init__(
-        self, url: Union[str, urls.URL], execution_options: dict = None, connect_args: dict = None
+        self,
+        url: Union[str, urls.URL],
+        execution_options: Optional[dict] = None,
+        connect_args: Optional[dict] = None,
     ) -> None:
         """Create sqlalchemy client based on the passed url.
 
@@ -84,16 +88,16 @@ def _extract_url_params(self) -> None:
     # Connection methods:
 
     def _connect(self) -> Connection:
-
-        parsed_url = sqla_url.URL(
+        parsed_url = sqla_url.create(
             drivername=self.drivername,
             username=self.username,
             password=self.password,
             host=self.host,
             port=self.port,
             database=self.database,
-            query=self.url_query,
         )
+        if self.url.query_string:
+            parsed_url.update_query_string(self.url.query_string)
         if self.engine is None:
             self.engine = create_engine(
                 parsed_url,
@@ -124,14 +128,14 @@ def query(self, sql_query: str, **kwargs) -> CursorResult:
 
         This will not commit any changes to the database.
         """
-        return self.conn.execute(sql_query, **kwargs)
+        return self.conn.execute(text(sql_query), **kwargs)
 
     @decorators.check_conn
     def execute(self, sql_query: str, **kwargs) -> None:
         """Execute a raw SQL query command."""
         trans = self.conn.begin()
         try:
-            self.conn.execute(sql_query, **kwargs)
+            self.conn.execute(text(sql_query), **kwargs)
         except Exception:
             trans.rollback()
             raise
@@ -141,7 +145,7 @@ def execute(self, sql_query: str, **kwargs) -> None:
     # Dataframe methods:
 
     @decorators.check_conn
-    def get_df(self, sql_query: str, params: dict = None, **kwargs) -> pd.DataFrame:
+    def get_df(self, sql_query: str, params: Optional[dict] = None, **kwargs) -> pd.DataFrame:
         """Run a raw SQL query and return a data frame."""
         return pd.read_sql(sql_query, self.conn, params=params, **kwargs)
 

src/tentaclio/hooks/slack_hook.py

@@ -1,6 +1,7 @@
 """Slack http hook."""
 import io
 import json
+from typing import Optional
 
 from tentaclio.clients import http_client
 
@@ -16,7 +17,11 @@ def __init__(self, url: str) -> None:
         self.url = url
 
     def notify(
-        self, message: str, channel: str = None, username: str = None, icon_emoji: str = None
+        self,
+        message: str,
+        channel: Optional[str] = None,
+        username: Optional[str] = None,
+        icon_emoji: Optional[str] = None,
     ) -> None:
         """Send a notification to slack."""
         body = self._build_request_body(
@@ -31,7 +36,10 @@ def notify(
 
     @staticmethod
     def _build_request_body(
-        message: str, channel: str = None, username: str = None, icon_emoji: str = None
+        message: str,
+        channel: Optional[str] = None,
+        username: Optional[str] = None,
+        icon_emoji: Optional[str] = None,
     ) -> str:
         # Fetch message payload
         payload = dict(text=message)

src/tentaclio/streams/api.py

@@ -1,5 +1,5 @@
 """Main entry points to tentaclio-io."""
-from typing import ContextManager, Union
+from typing import ContextManager, Optional, Union
 
 from tentaclio import protocols
 from tentaclio.credentials import authenticate
@@ -17,7 +17,7 @@
 ]
 
 
-def open(url: str, mode: str = None, **kwargs) -> AnyContextStreamerReaderWriter:
+def open(url: str, mode: Optional[str] = None, **kwargs) -> AnyContextStreamerReaderWriter:
     """Open a url and return a reader or writer depending on mode.
 
     Arguments:

src/tentaclio/urls.py

@@ -38,6 +38,7 @@ class URL:
     _port: Optional[int] = None
     _path: str
     _query: Optional[Dict[str, str]] = None
+    _query_string: Optional[str] = None
 
     def __init__(self, url: str) -> None:
         """Create a url by parsing the parametre."""
@@ -57,6 +58,7 @@ def _parse_url(self) -> None:
         self._hostname = parsed_url.hostname
         self._port = parsed_url.port
         self._path = parsed_url.path
+        self._query_string = parsed_url.query
 
         # Replace %xx escapes - ONLY for username & password
         if parsed_url.username and self._username:
@@ -74,13 +76,13 @@ def _parse_url(self) -> None:
 
     def copy(
         self,
-        scheme: str = None,
-        username: str = None,
-        password: str = None,
-        hostname: str = None,
-        port: int = None,
-        path: str = None,
-        query: str = None,
+        scheme: Optional[str] = None,
+        username: Optional[str] = None,
+        password: Optional[str] = None,
+        hostname: Optional[str] = None,
+        port: Optional[int] = None,
+        path: Optional[str] = None,
+        query: Optional[str] = None,
     ) -> "URL":
         """Copy this url optionally overwriting the provided components."""
         return URL.from_components(
@@ -168,6 +170,11 @@ def query(self) -> Optional[Dict[str, str]]:
         """Access the query."""
         return self._query
 
+    @property
+    def query_string(self) -> Optional[str]:
+        """Access the query string."""
+        return self._query_string
+
     @property
     def url(self) -> str:
         """Return the original url."""

tests/unit/clients/test_decorators.py

@@ -22,7 +22,6 @@ def test_inactive_client_connection(self):
         url = "file:///path"
 
         class TestClient(base_client.BaseClient):
-
             allowed_schemes = ["file"]
 
             def connect(self):

tests/unit/clients/test_ftp_client.py

@@ -49,7 +49,6 @@ def test_parsing_ftp_url(self, url, username, password, hostname, port, path):
     @pytest.mark.parametrize("url,path", [("ftp://:@localhost", None)])
     def test_get_invalid_path(self, url, path, mocked_ftp_conn):
         with ftp_client.FTPClient(url) as client:
-
             with pytest.raises(exceptions.FTPError):
                 client.get(io.StringIO(), file_path=path)
 
@@ -149,7 +148,6 @@ def test_invalid_scheme(self, url):
     @pytest.mark.parametrize("url,path", [("sftp://:@localhost", None)])
     def test_get_invalid_path(self, url, path, mocked_sftp_conn):
         with ftp_client.SFTPClient(url) as client:
-
             with pytest.raises(exceptions.FTPError):
                 client.get(io.StringIO(), file_path=path)
 

tests/unit/clients/test_http_client.py

@@ -53,7 +53,6 @@ def test_get_invalid_endpoint(self, url, path, mocker, mocked_http_conn):
         mocked_writer = mocker.Mock()
 
         with pytest.raises(exceptions.HTTPError):
-
             with http_client.HTTPClient(url) as client:
                 client.get(mocked_writer, endpoint=path)
 
@@ -73,7 +72,6 @@ def test_get_invalid_endpoint(self, url, path, mocker, mocked_http_conn):
     )
     def test_fetching_url_endpoint(self, base_url, endpoint, auth, full_url):
         with http_client.HTTPClient(base_url) as client:
-
             assert client.conn.auth == auth
             assert client._fetch_url(endpoint) == full_url
 

tests/unit/credentials/test_api.py

@@ -2,7 +2,6 @@
 
 
 def test_authenticate(mocker):
-
     injector = credentials.CredentialsInjector()
     injector.register_credentials(urls.URL("ftp://user:[email protected]"))
     mock_cred = mocker.patch("tentaclio.credentials.api.load_credentials_injector")

tests/unit/streams/test_csv_db_stream.py

@@ -22,7 +22,6 @@ def csv_data():
 
 
 def test_dump_csv(csv_data, csv_dumper):
-
     with csv_db_stream.DatabaseCsvWriter(csv_dumper, "my_table") as writer:
         writer.write(csv_data.getvalue())
 

tests/unit/streams/test_stream_handler.py

@@ -1,4 +1,5 @@
 import io
+from typing import Optional
 
 import pytest
 
@@ -9,7 +10,7 @@
 
 class FakeClient(base_client.BaseClient["FakeClient"]):
     # clients only understand bytes
-    def __init__(self, url: URL, message: bytearray = None, *args, **kwargs):
+    def __init__(self, url: URL, message: Optional[bytearray] = None, *args, **kwargs):
         self._writer = io.BytesIO()
         self._message = message or bytes("hello", encoding="utf-8")
         self._closed = False

tests/unit/test_urls.py

@@ -4,7 +4,6 @@
 
 
 class TestURL:
-
     # Generic parsing rules:
     def test_missing_url(self):
         with pytest.raises(urls.URLError):