Merge pull request #31 from yassun7010/add_connection_classmethod

wip: add connection classmethods.

Author: ya7010

turu-bigquery/src/turu/bigquery/__init__.py

@@ -1,6 +1,6 @@
 import importlib.metadata
 
-from .connection import Connection, connect
+from .connection import Connection
 from .cursor import Cursor
 from .mock_connection import MockConnection
 from .mock_cursor import MockCursor
@@ -14,3 +14,5 @@
     "MockConnection",
     "MockCursor",
 ]
+
+connect = Connection.connect

turu-bigquery/src/turu/bigquery/connection.py

@@ -1,4 +1,4 @@
-from typing import Optional
+from typing import Any, Optional
 
 import google.api_core.client_info
 import google.api_core.client_options
@@ -9,15 +9,45 @@
 import turu.bigquery.cursor
 import turu.core.connection
 import turu.core.mock
-from typing_extensions import Never, deprecated, override
+from typing_extensions import Never, Self, deprecated, override
 
 from .cursor import Cursor
 
+try:
+    from google.cloud.bigquery_storage import BigQueryReadClient  # type: ignore
+
+except ImportError:
+
+    class BigQueryReadClient:
+        pass
+
 
 class Connection(turu.core.connection.Connection):
     def __init__(self, connection: google.cloud.bigquery.dbapi.Connection):
         self._raw_connection = connection
 
+    @override
+    @classmethod
+    def connect(  # type: ignore[override]
+        cls,
+        client: Optional[google.cloud.bigquery.Client] = None,
+        bqstorage_client: Optional[BigQueryReadClient] = None,
+    ) -> Self:
+        import google.cloud.bigquery
+        import google.cloud.bigquery.dbapi
+
+        return cls(
+            google.cloud.bigquery.dbapi.connect(
+                client=client,
+                bqstorage_client=bqstorage_client,
+            ),
+        )
+
+    @classmethod
+    @override
+    def connect_from_env(cls, *args: Any, **kwargs: Any) -> Self:
+        return cls.connect(*args, **kwargs)
+
     @override
     def close(self) -> None:
         """Close the connection and any cursors created from it."""
@@ -39,28 +69,3 @@ def cursor(self) -> Cursor[Never]:
         """Return a new cursor object."""
 
         return Cursor(self._raw_connection.cursor())
-
-
-try:
-    from google.cloud.bigquery_storage import BigQueryReadClient  # type: ignore
-
-
-except ImportError:
-
-    class BigQueryReadClient:
-        pass
-
-
-def connect(
-    client: Optional[google.cloud.bigquery.Client] = None,
-    bqstorage_client: Optional[BigQueryReadClient] = None,
-) -> Connection:
-    import google.cloud.bigquery
-    import google.cloud.bigquery.dbapi
-
-    return Connection(
-        google.cloud.bigquery.dbapi.connect(
-            client=client,
-            bqstorage_client=bqstorage_client,
-        ),
-    )

turu-bigquery/src/turu/bigquery/mock_connection.py

@@ -7,7 +7,7 @@
 from .mock_cursor import MockCursor
 
 
-class MockConnection(Connection, turu.core.mock.MockConnection):
+class MockConnection(turu.core.mock.MockConnection, Connection):
     def __init__(self, *args, **kwargs):
         turu.core.mock.MockConnection.__init__(self)
 

turu-core/src/turu/core/async_connection.py

@@ -4,10 +4,22 @@
 import turu.core.async_cursor
 from turu.core.protocols.async_connection import AsyncConnectionProtocol
 from turu.core.protocols.async_cursor import Parameters
-from typing_extensions import Never
+from typing_extensions import Never, Self
 
 
 class AsyncConnection(AsyncConnectionProtocol):
+    @classmethod
+    @abstractmethod
+    async def connect(cls, *args: Any, **kwargs: Any) -> Self:
+        """Connect to a database."""
+        ...
+
+    @classmethod
+    @abstractmethod
+    async def connect_from_env(cls, *args: Any, **kwargs: Any) -> Self:
+        """Connect to a database using environment variables."""
+        ...
+
     @abstractmethod
     async def close(self) -> None:
         ...

turu-core/src/turu/core/connection.py

@@ -4,10 +4,22 @@
 import turu.core.cursor
 from turu.core.protocols.connection import ConnectionProtocol
 from turu.core.protocols.cursor import Parameters
-from typing_extensions import Never, override
+from typing_extensions import Never, Self, override
 
 
 class Connection(ConnectionProtocol):
+    @classmethod
+    @abstractmethod
+    def connect(cls, *args: Any, **kwargs: Any) -> Self:
+        """Connect to a database."""
+        ...
+
+    @classmethod
+    @abstractmethod
+    def connect_from_env(cls, *args: Any, **kwargs: Any) -> Self:
+        """Connect to a database using environment variables."""
+        ...
+
     @abstractmethod
     @override
     def close(self) -> None:

turu-core/src/turu/core/mock/async_connection.py

@@ -23,6 +23,14 @@ class MockAsyncConnection(turu.core.async_connection.AsyncConnection):
     def __init__(self, store: Optional[TuruMockStore] = None):
         self._turu_mock_store = store or TuruMockStore()
 
+    @classmethod
+    async def connect(cls, *args: Any, **kwargs: Any) -> Self:
+        return cls()
+
+    @classmethod
+    async def connect_from_env(cls, *args: Any, **kwargs: Any) -> Self:
+        return cls()
+
     def chain(self) -> Self:
         """this method is just for code formatting by black."""
 

turu-core/src/turu/core/mock/connection.py

@@ -27,6 +27,14 @@ class MockConnection(turu.core.connection.Connection):
     def __init__(self, store: Optional[TuruMockStore] = None):
         self._turu_mock_store = store or TuruMockStore()
 
+    @classmethod
+    def connect(cls, *args: Any, **kwargs: Any) -> Self:
+        return cls()
+
+    @classmethod
+    def connect_from_env(cls, *args: Any, **kwargs: Any) -> Self:
+        return cls()
+
     def chain(self) -> Self:
         """this method is just for code formatting by black."""
 

turu-core/tests/turu/core/test_record.py

@@ -8,7 +8,7 @@
 from pydantic import BaseModel
 from turu.core.exception import TuruRowTypeNotSupportedError
 from turu.core.record import RecordCursor, record_to_csv
-from typing_extensions import Never
+from typing_extensions import Never, Self
 
 
 class RowPydantic(BaseModel):
@@ -197,6 +197,14 @@ def custom_method(self, value: int) -> None:
                 pass
 
         class CustomConnection(turu.core.mock.MockConnection):
+            @classmethod
+            def connect(cls, *args, **kwargs) -> Self:
+                return cls()
+
+            @classmethod
+            def connect_from_env(cls, *args, **kwargs) -> Self:
+                return cls.connect()
+
             def cursor(self) -> CustomCursor:
                 return CustomCursor(self._turu_mock_store)
 

turu-mysql/src/turu/mysql/__init__.py

@@ -2,14 +2,10 @@
 
 from .async_connection import (
     AsyncConnection,
-    connect_async,
-    connect_async_from_env,
 )
 from .async_cursor import AsyncCursor
 from .connection import (
     Connection,
-    connect,
-    connect_from_env,
 )
 from .cursor import Cursor
 from .mock_async_connection import MockAsyncConnection
@@ -19,7 +15,6 @@
 
 __version__ = importlib.metadata.version("turu-mysql")
 
-
 __all__ = [
     "AsyncConnection",
     "AsyncCursor",
@@ -34,3 +29,8 @@
     "MockConnection",
     "MockCursor",
 ]
+
+connect = Connection.connect
+connect_from_env = Connection.connect_from_env
+connect_async = AsyncConnection.connect
+connect_async_from_env = AsyncConnection.connect_from_env

turu-mysql/src/turu/mysql/async_connection.py

@@ -9,7 +9,7 @@
 import turu.core.mock
 import turu.mysql.async_cursor
 import turu.mysql.cursor
-from typing_extensions import Never, Unpack, override
+from typing_extensions import Never, Self, Unpack, override
 
 from .async_cursor import AsyncCursor
 
@@ -18,6 +18,48 @@ class AsyncConnection(turu.core.async_connection.AsyncConnection):
     def __init__(self, connection: aiomysql.Connection):
         self._raw_connection = connection
 
+    @override
+    @classmethod
+    async def connect(  # type: ignore[override]
+        cls,
+        user: Optional[str] = None,
+        password: str = "",
+        host: str = "localhost",
+        database: Optional[str] = None,
+        port: int = 0,
+        **kwargs: Unpack["_ConnectParams"],
+    ) -> Self:
+        return cls(
+            await aiomysql.connection._connect(
+                user=user,
+                password=password,
+                host=host,
+                db=database,
+                port=port,
+                **kwargs,
+            )
+        )
+
+    @override
+    @classmethod
+    async def connect_from_env(  # type: ignore[override]
+        cls,
+        user_envname: str = "MYSQL_USER",
+        password_envname: str = "MYSQL_PASSWORD",
+        host_envname: str = "MYSQL_HOST",
+        database_envname: str = "MYSQL_DATABASE",
+        port_envname: str = "MYSQL_PORT",
+        **kwargs: Unpack["_ConnectParams"],
+    ) -> Self:
+        return await cls.connect(
+            user=os.environ.get(user_envname),
+            password=os.environ.get(password_envname, ""),
+            host=os.environ.get(host_envname, "localhost"),
+            database=os.environ.get(database_envname),
+            port=int(os.environ.get(port_envname, 0)),
+            **kwargs,
+        )
+
     @override
     async def close(self) -> None:
         await self._raw_connection.ensure_closed()
@@ -60,41 +102,3 @@ class _ConnectParams(TypedDict, total=False):
     auth_plugin: str
     program_name: str
     server_public_key: Optional[Any]
-
-
-async def connect_async(
-    user: Optional[str] = None,
-    password: str = "",
-    host: str = "localhost",
-    database: Optional[str] = None,
-    port: int = 0,
-    **kwargs: Unpack[_ConnectParams],
-) -> AsyncConnection:
-    return AsyncConnection(
-        await aiomysql.connection._connect(
-            user=user,
-            password=password,
-            host=host,
-            db=database,
-            port=port,
-            **kwargs,
-        )
-    )
-
-
-async def connect_async_from_env(
-    user_envname: str = "MYSQL_USER",
-    password_envname: str = "MYSQL_PASSWORD",
-    host_envname: str = "MYSQL_HOST",
-    database_envname: str = "MYSQL_DATABASE",
-    port_envname: str = "MYSQL_PORT",
-    **kwargs: Unpack[_ConnectParams],
-) -> AsyncConnection:
-    return await connect_async(
-        user=os.environ.get(user_envname),
-        password=os.environ.get(password_envname, ""),
-        host=os.environ.get(host_envname, "localhost"),
-        database=os.environ.get(database_envname),
-        port=int(os.environ.get(port_envname, 0)),
-        **kwargs,
-    )