Repository: dask/dask

Issue:
to_sql sqlalchemy pickle error with 2022.2.0
**What happened**:
to_sql sqlalchemy pickle error with dask 2022.2.0

**What you expected to happen**:
to_sql produces no pickle error with 2022.1.1

**Minimal Complete Verifiable Example**:

```python
import pandas as pd
import dask.dataframe as dd
from dask.distributed import Client

client = Client()

data = dd.from_pandas(pd.DataFrame([1, 2, 3, 4]), chunksize=20000)
data.to_sql(
    "testtest",
    uri=str(sqlengine.url),
    if_exists="replace",
    schema=None,
    index=False,
)
```

<details>
<summary>Traceback</summary>

```
Traceback (most recent call last):
  File "/opt/conda/lib/python3.9/site-packages/distributed/worker.py", line 4442, in dumps_function
    result = cache_dumps[func]
  File "/opt/conda/lib/python3.9/site-packages/distributed/utils.py", line 1360, in __getitem__
    value = super().__getitem__(key)
  File "/opt/conda/lib/python3.9/collections/__init__.py", line 1058, in __getitem__
    raise KeyError(key)
KeyError: <function to_sql.<locals>.make_meta at 0x7f05284b3af0>

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/opt/conda/lib/python3.9/site-packages/distributed/protocol/pickle.py", line 49, in dumps
    result = pickle.dumps(x, **dump_kwargs)
AttributeError: Can't pickle local object 'to_sql.<locals>.make_meta'

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/opt/conda/lib/python3.9/site-packages/dask/dataframe/core.py", line 1580, in to_sql
    return to_sql(
  File "/opt/conda/lib/python3.9/site-packages/dask/dataframe/io/sql.py", line 850, in to_sql
    dask.compute(result)
  File "/opt/conda/lib/python3.9/site-packages/dask/base.py", line 573, in compute
    results = schedule(dsk, keys, **kwargs)
  File "/opt/conda/lib/python3.9/site-packages/distributed/client.py", line 2974, in get
    futures = self._graph_to_futures(
  File "/opt/conda/lib/python3.9/site-packages/distributed/client.py", line 2882, in _graph_to_futures
    dsk = dsk.__dask_distributed_pack__(self, keyset, annotations)
  File "/opt/conda/lib/python3.9/site-packages/dask/highlevelgraph.py", line 1054, in __dask_distributed_pack__
    "state": layer.__dask_distributed_pack__(
  File "/opt/conda/lib/python3.9/site-packages/dask/highlevelgraph.py", line 425, in __dask_distributed_pack__
    dsk = toolz.valmap(dumps_task, dsk)
  File "cytoolz/dicttoolz.pyx", line 178, in cytoolz.dicttoolz.valmap
  File "cytoolz/dicttoolz.pyx", line 203, in cytoolz.dicttoolz.valmap
  File "/opt/conda/lib/python3.9/site-packages/distributed/worker.py", line 4480, in dumps_task
    return {"function": dumps_function(task[0]), "args": warn_dumps(task[1:])}
  File "/opt/conda/lib/python3.9/site-packages/distributed/worker.py", line 4444, in dumps_function
    result = pickle.dumps(func, protocol=4)
  File "/opt/conda/lib/python3.9/site-packages/distributed/protocol/pickle.py", line 60, in dumps
    result = cloudpickle.dumps(x, **dump_kwargs)
  File "/opt/conda/lib/python3.9/site-packages/cloudpickle/cloudpickle_fast.py", line 73, in dumps
    cp.dump(obj)
  File "/opt/conda/lib/python3.9/site-packages/cloudpickle/cloudpickle_fast.py", line 602, in dump
    return Pickler.dump(self, obj)
TypeError: cannot pickle 'sqlalchemy.cprocessors.UnicodeResultProcessor' object
```
</details>

**Anything else we need to know?**:
tested with SQLAlchemy 1.4.31

**Environment**:

- Dask version: 2022.2.0
- Python version: 3.9.10
- Operating System: linux
- Install method (conda, pip, source): conda/mamba

<!-- If you are reporting an issue such as scale stability, cluster deadlock.
Please provide a cluster dump state with this issue, by running client.dump_cluster_state()

https://distributed.dask.org/en/stable/api.html?highlight=dump_cluster_state#distributed.Client.dump_cluster_state

-->


Discussion:
@jsignell I think it might have to do with this change: https://github.com/dask/dask/blob/f634f2b01caeb10dcd33734f780d9bb6056a21d9/dask/dataframe/io/sql.py#L798

the `engine` is created already at this stage and is not picklable (just changed with https://github.com/dask/dask/pull/8609)
Thanks for raising this @mgsnuno! I think you are right in your diagnosis. I guess our options are to revert that #8609 or figure out a way to pickle the engine object. Ping @scharlottej13 and @ian-r-rose for options on the best path forward.
Ah, yeah, we shouldn't be trying to create the `engine` directly at graph construction time for exactly this reason. The `read_sql` path defers engine creation until `_read_sql_chunk`, this `to_sql` be doing something similar.

It should be a reasonably straightforward fix, though perhaps we should revert #8609 until we can get it in.
Yeah it does seem like we could write a similar `_to_sql_chunk`. 

Resolving patch:
diff --git a/dask/dataframe/io/sql.py b/dask/dataframe/io/sql.py
--- a/dask/dataframe/io/sql.py
+++ b/dask/dataframe/io/sql.py
@@ -438,6 +438,18 @@ def _read_sql_chunk(q, uri, meta, engine_kwargs=None, **kwargs):
         return df.astype(meta.dtypes.to_dict(), copy=False)
 
 
+def _to_sql_chunk(d, uri, engine_kwargs=None, **kwargs):
+    import sqlalchemy as sa
+
+    engine_kwargs = engine_kwargs or {}
+    engine = sa.create_engine(uri, **engine_kwargs)
+
+    q = d.to_sql(con=engine, **kwargs)
+    engine.dispose()
+
+    return q
+
+
 def _gt14() -> bool:
     """
     Check if sqlalchemy.__version__ is at least 1.4.0, when several
@@ -789,18 +801,14 @@ def to_sql(
     >>> result
     [(0, 0, '00'), (1, 1, '11'), (2, 2, '22'), (3, 3, '33')]
     """
-    import sqlalchemy as sa
-
     if not isinstance(uri, str):
         raise ValueError(f"Expected URI to be a string, got {type(uri)}.")
 
-    engine_kwargs = {} if engine_kwargs is None else engine_kwargs
-    engine = sa.create_engine(uri, **engine_kwargs)
-
     # This is the only argument we add on top of what Pandas supports
     kwargs = dict(
         name=name,
-        con=engine,
+        uri=uri,
+        engine_kwargs=engine_kwargs,
         schema=schema,
         if_exists=if_exists,
         index=index,
@@ -810,11 +818,7 @@ def to_sql(
         method=method,
     )
 
-    def make_meta(meta):
-        return meta.to_sql(**kwargs)
-
-    make_meta = delayed(make_meta)
-    meta_task = make_meta(df._meta)
+    meta_task = delayed(_to_sql_chunk)(df._meta, **kwargs)
 
     # Partitions should always append to the empty table created from `meta` above
     worker_kwargs = dict(kwargs, if_exists="append")
@@ -823,7 +827,8 @@ def make_meta(meta):
         # Perform the meta insert, then one task that inserts all blocks concurrently:
         result = [
             _extra_deps(
-                d.to_sql,
+                _to_sql_chunk,
+                d,
                 extras=meta_task,
                 **worker_kwargs,
                 dask_key_name="to_sql-%s" % tokenize(d, **worker_kwargs),
@@ -837,7 +842,8 @@ def make_meta(meta):
         for d in df.to_delayed():
             result.append(
                 _extra_deps(
-                    d.to_sql,
+                    _to_sql_chunk,
+                    d,
                     extras=last,
                     **worker_kwargs,
                     dask_key_name="to_sql-%s" % tokenize(d, **worker_kwargs),


Test patch:
diff --git a/dask/tests/test_distributed.py b/dask/tests/test_distributed.py
--- a/dask/tests/test_distributed.py
+++ b/dask/tests/test_distributed.py
@@ -18,7 +18,7 @@
 from dask.delayed import Delayed
 from dask.distributed import futures_of, wait
 from dask.highlevelgraph import HighLevelGraph, MaterializedLayer
-from dask.utils import get_named_args, tmpdir
+from dask.utils import get_named_args, tmpdir, tmpfile
 
 if "should_check_state" in get_named_args(gen_cluster):
     gen_cluster = partial(gen_cluster, should_check_state=False)
@@ -659,3 +659,27 @@ async def test_pack_MaterializedLayer_handles_futures_in_graph_properly(c, s, a,
     packed = hlg.__dask_distributed_pack__(c, ["z"], {})
     unpacked = HighLevelGraph.__dask_distributed_unpack__(packed)
     assert unpacked["deps"] == {"x": {fut.key}, "y": {fut.key}, "z": {"y"}}
+
+
+@gen_cluster(client=True)
+async def test_to_sql_engine_kwargs(c, s, a, b):
+    # https://github.com/dask/dask/issues/8738
+    pd = pytest.importorskip("pandas")
+    dd = pytest.importorskip("dask.dataframe")
+    pytest.importorskip("sqlalchemy")
+
+    df = pd.DataFrame({"a": range(10), "b": range(10)})
+    df.index.name = "index"
+    ddf = dd.from_pandas(df, npartitions=1)
+    with tmpfile() as f:
+        uri = f"sqlite:///{f}"
+        result = ddf.to_sql(
+            "test", uri, index=True, engine_kwargs={"echo": False}, compute=False
+        )
+        await c.compute(result)
+
+        dd.utils.assert_eq(
+            ddf,
+            dd.read_sql_table("test", uri, "index"),
+            check_divisions=False,
+        )