feat: async stuff!

Author: levydsa

libsql_sqlalchemy/__init__.py

@@ -1,117 +1,10 @@
-import os
-import urllib.parse
-
-from sqlalchemy import util
 from sqlalchemy.dialects import registry as _registry
-from sqlalchemy.dialects.sqlite.pysqlite import SQLiteDialect_pysqlite
-
-from libsql_experimental import Connection
+from .aiolibsql import SQLiteDialect_aiolibsql
+from .libsql import SQLiteDialect_libsql
 
 __version__ = "0.1.0-pre"
 
-_registry.register("sqlite.libsql", "libsql_sqlalchemy", "SQLiteDialect_libsql")
-
-
-def _build_connection_url(url, query, secure):
-    # sorting of keys is for unit test support
-    query_str = urllib.parse.urlencode(sorted(query.items()))
-
-    if not url.host:
-        if query_str:
-            return f"{url.database}?{query_str}"
-        return url.database
-    elif secure:  # yes, pop to remove
-        scheme = "https"
-    else:
-        scheme = "http"
-
-    if url.username and url.password:
-        netloc = f"{url.username}:{url.password}@{url.host}"
-    elif url.username:
-        netloc = f"{url.username}@{url.host}"
-    else:
-        netloc = url.host
-
-    if url.port:
-        netloc += f":{url.port}"
-
-    return urllib.parse.urlunsplit(
-        (
-            scheme,
-            netloc,
-            url.database or "",
-            query_str,
-            "",  # fragment
-        )
-    )
-
-
-class SQLiteDialect_libsql(SQLiteDialect_pysqlite):
-    driver = "libsql"
-    # need to be set explicitly
-    supports_statement_cache = SQLiteDialect_pysqlite.supports_statement_cache
-
-    @classmethod
-    def import_dbapi(cls):
-        import libsql_experimental as libsql
-
-        return libsql
-
-    def on_connect(self):
-        import libsql_experimental as libsql
-
-        sqlite3_connect = super().on_connect()
-
-        def connect(conn):
-            # LibSQL: there is no support for create_function()
-            if isinstance(conn, Connection):
-                return
-            return sqlite3_connect(conn)
-
-        return connect
-
-    def create_connect_args(self, url):
-        pysqlite_args = (
-            ("uri", bool),
-            ("timeout", float),
-            ("isolation_level", str),
-            ("detect_types", int),
-            ("check_same_thread", bool),
-            ("cached_statements", int),
-            ("secure", bool),  # LibSQL extra, selects between ws and wss
-        )
-        opts = url.query
-        libsql_opts = {}
-        for key, type_ in pysqlite_args:
-            util.coerce_kw_type(opts, key, type_, dest=libsql_opts)
-
-        if url.host:
-            libsql_opts["uri"] = True
-
-        if libsql_opts.get("uri", False):
-            uri_opts = dict(opts)
-            # here, we are actually separating the parameters that go to
-            # sqlite3/pysqlite vs. those that go the SQLite URI.  What if
-            # two names conflict?  again, this seems to be not the case right
-            # now, and in the case that new names are added to
-            # either side which overlap, again the sqlite3/pysqlite parameters
-            # can be passed through connect_args instead of in the URL.
-            # If SQLite native URIs add a parameter like "timeout" that
-            # we already have listed here for the python driver, then we need
-            # to adjust for that here.
-            for key, type_ in pysqlite_args:
-                uri_opts.pop(key, None)
-
-            secure = libsql_opts.pop("secure", False)
-            connect_url = _build_connection_url(url, uri_opts, secure)
-        else:
-            connect_url = url.database or ":memory:"
-            if connect_url != ":memory:":
-                connect_url = os.path.abspath(connect_url)
-
-        libsql_opts.setdefault("check_same_thread", not self._is_url_file_db(url))
-
-        return ([connect_url], libsql_opts)
-
-
-dialect = SQLiteDialect_libsql
+_registry.register("sqlite.libsql", "libsql_sqlalchemy.libsql", "SQLiteDialect_libsql")
+_registry.register(
+    "sqlite.aiolibsql", "libsql_sqlalchemy.aiolibsql", "SQLiteDialect_aiolibsql"
+)

libsql_sqlalchemy/aiolibsql.py

@@ -0,0 +1,10 @@
+from libsql_sqlalchemy.libsql import SQLiteDialect_libsql
+
+
+class SQLiteDialect_aiolibsql(SQLiteDialect_libsql):
+    driver = "libsql"
+    supports_statement_cache = SQLiteDialect_libsql.supports_statement_cache
+    is_async = True
+
+
+dialect = SQLiteDialect_aiolibsql

libsql_sqlalchemy/libsql.py

@@ -0,0 +1,111 @@
+import os
+import urllib.parse
+
+from sqlalchemy import util
+from sqlalchemy.dialects.sqlite.pysqlite import SQLiteDialect_pysqlite
+from libsql_experimental import Connection
+
+
+def _build_connection_url(url, query, secure):
+    # sorting of keys is for unit test support
+    query_str = urllib.parse.urlencode(sorted(query.items()))
+
+    if not url.host:
+        if query_str:
+            return f"{url.database}?{query_str}"
+        return url.database
+    elif secure:  # yes, pop to remove
+        scheme = "https"
+    else:
+        scheme = "http"
+
+    if url.username and url.password:
+        netloc = f"{url.username}:{url.password}@{url.host}"
+    elif url.username:
+        netloc = f"{url.username}@{url.host}"
+    else:
+        netloc = url.host
+
+    if url.port:
+        netloc += f":{url.port}"
+
+    return urllib.parse.urlunsplit(
+        (
+            scheme,
+            netloc,
+            url.database or "",
+            query_str,
+            "",  # fragment
+        )
+    )
+
+
+class SQLiteDialect_libsql(SQLiteDialect_pysqlite):
+    driver = "libsql"
+    # need to be set explicitly
+    supports_statement_cache = SQLiteDialect_pysqlite.supports_statement_cache
+
+    @classmethod
+    def import_dbapi(cls):
+        import libsql_experimental as libsql
+
+        return libsql
+
+    def on_connect(self):
+        import libsql_experimental as libsql
+
+        sqlite3_connect = super().on_connect()
+
+        def connect(conn):
+            # LibSQL: there is no support for create_function()
+            if isinstance(conn, Connection):
+                return
+            return sqlite3_connect(conn)
+
+        return connect
+
+    def create_connect_args(self, url):
+        pysqlite_args = (
+            ("uri", bool),
+            ("timeout", float),
+            ("isolation_level", str),
+            ("detect_types", int),
+            ("check_same_thread", bool),
+            ("cached_statements", int),
+            ("secure", bool),  # LibSQL extra, selects between ws and wss
+        )
+        opts = url.query
+        libsql_opts = {}
+        for key, type_ in pysqlite_args:
+            util.coerce_kw_type(opts, key, type_, dest=libsql_opts)
+
+        if url.host:
+            libsql_opts["uri"] = True
+
+        if libsql_opts.get("uri", False):
+            uri_opts = dict(opts)
+            # here, we are actually separating the parameters that go to
+            # sqlite3/pysqlite vs. those that go the SQLite URI.  What if
+            # two names conflict?  again, this seems to be not the case right
+            # now, and in the case that new names are added to
+            # either side which overlap, again the sqlite3/pysqlite parameters
+            # can be passed through connect_args instead of in the URL.
+            # If SQLite native URIs add a parameter like "timeout" that
+            # we already have listed here for the python driver, then we need
+            # to adjust for that here.
+            for key, type_ in pysqlite_args:
+                uri_opts.pop(key, None)
+
+            secure = libsql_opts.pop("secure", False)
+            connect_url = _build_connection_url(url, uri_opts, secure)
+        else:
+            connect_url = url.database or ":memory:"
+            if connect_url != ":memory:":
+                connect_url = os.path.abspath(connect_url)
+
+        libsql_opts.setdefault("check_same_thread", not self._is_url_file_db(url))
+
+        return ([connect_url], libsql_opts)
+
+
+dialect = SQLiteDialect_libsql

pyproject.toml

@@ -20,12 +20,14 @@ classifiers = [
 dependencies = [
     "libsql-experimental>=0.0.47",
     "sqlalchemy>=2.0.0",
+    "greenlet>=3.0.3",
 ]
 
 [project.optional-dependencies]
 dev = [
     "pytest>=8.3.5",
     "pytest-cov>=4.1.0",
+    "pytest-asyncio>=0.23.5",
 ]
 
 [project.urls]
@@ -35,4 +37,5 @@ Repository = "https://github.com/tursodatabase/libsql-sqlalchemy.git"
 Issues = "https://github.com/tursodatabase/libsql-sqlalchemy/issues"
 
 [project.entry-points."sqlalchemy.dialects"]
-"sqlite.libsql" = "libsql_sqlalchemy:SQLiteDialect_libsql"
+"sqlite.libsql" = "libsql_sqlalchemy:SQLiteDialect_libsql"
+"sqlite.aiolibsql" = "libsql_sqlalchemy:SQLiteDialect_aiolibsql"

tests/test_async.py

@@ -0,0 +1,98 @@
+import pytest
+import pytest_asyncio
+import asyncio
+from sqlalchemy import Column, Integer, String, text, select
+from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker
+from sqlalchemy.orm import declarative_base
+from sqlalchemy.pool import AsyncAdaptedQueuePool
+
+Base = declarative_base()
+
+
+class User(Base):
+    __tablename__ = "users"
+
+    id = Column(Integer, primary_key=True)
+    name = Column(String)
+    email = Column(String)
+
+
+@pytest_asyncio.fixture
+async def engine():
+    engine = create_async_engine(
+        "sqlite+aiolibsql://",
+        poolclass=AsyncAdaptedQueuePool,
+    )
+    yield engine
+    await engine.dispose()
+
+
+@pytest_asyncio.fixture
+async def session(engine):
+    async with engine.begin() as conn:
+        await conn.run_sync(Base.metadata.create_all)
+
+    async_session = async_sessionmaker(
+        engine, class_=AsyncSession, expire_on_commit=False
+    )
+    async with async_session() as session:
+        yield session
+
+    async with engine.begin() as conn:
+        await conn.run_sync(Base.metadata.drop_all)
+
+
+@pytest.mark.asyncio
+async def test_connection(session: AsyncSession):
+    result = await session.execute(text("SELECT 1"))
+    assert result.scalar() == 1
+
+
+@pytest.mark.asyncio
+async def test_create_table(session):
+    result = await session.execute(
+        text("SELECT name FROM sqlite_master WHERE type='table' AND name='users'")
+    )
+    assert result.scalar() == "users"
+
+
+@pytest.mark.asyncio
+async def test_insert_and_query(session):
+    user = User(name="Test User", email="[email protected]")
+    session.add(user)
+    await session.commit()
+
+    stmt = select(User)
+    result = await session.execute(stmt)
+    queried_user = result.scalars().first()
+    assert queried_user.name == "Test User"
+    assert queried_user.email == "[email protected]"
+
+
+@pytest.mark.asyncio
+async def test_update(session):
+    user = User(name="Test User", email="[email protected]")
+    session.add(user)
+    await session.commit()
+
+    user.name = "Updated User"
+    await session.commit()
+
+    stmt = select(User)
+    result = await session.execute(stmt)
+    updated_user = result.scalars().first()
+    assert updated_user.name == "Updated User"
+
+
+@pytest.mark.asyncio
+async def test_delete(session):
+    user = User(name="Test User", email="[email protected]")
+    session.add(user)
+    await session.commit()
+
+    await session.delete(user)
+    await session.commit()
+
+    stmt = select(User)
+    result = await session.execute(stmt)
+    assert result.scalars().first() is None