Set Dask-sql as the default Fugue Dask engine when installed

continuous_integration/environment-3.10-jdk11-dev.yaml

@@ -6,6 +6,7 @@ dependencies:
 - dask-ml>=2022.1.22
 - dask>=2022.3.0
 - fastapi>=0.69.0
+- fugue>=0.7.0
 - intake>=0.6.0
 - jpype1>=1.0.2
 - jsonschema
@@ -31,14 +32,3 @@ dependencies:
 - tpot
 - tzlocal>=2.1
 - uvicorn>=0.11.3
-# fugue dependencies; remove when we conda install fugue
-- adagio
-- antlr4-python3-runtime<4.10
-- ciso8601
-- fs
-- pip
-- qpd
-- triad
-- pip:
-  # FIXME: tests are failing with fugue 0.7.0
-  - fugue[sql]>=0.5.3,<0.7.0

continuous_integration/environment-3.10-jdk8-dev.yaml

@@ -6,6 +6,7 @@ dependencies:
 - dask-ml>=2022.1.22
 - dask>=2022.3.0
 - fastapi>=0.69.0
+- fugue>=0.7.0
 - intake>=0.6.0
 - jpype1>=1.0.2
 - jsonschema
@@ -31,14 +32,3 @@ dependencies:
 - tpot
 - tzlocal>=2.1
 - uvicorn>=0.11.3
-# fugue dependencies; remove when we conda install fugue
-- adagio
-- antlr4-python3-runtime<4.10
-- ciso8601
-- fs
-- pip
-- qpd
-- triad
-- pip:
-  # FIXME: tests are failing with fugue 0.7.0
-  - fugue[sql]>=0.5.3,<0.7.0

continuous_integration/environment-3.8-jdk11-dev.yaml

@@ -6,6 +6,7 @@ dependencies:
 - dask-ml=2022.1.22
 - dask=2022.3.0
 - fastapi=0.69.0
+- fugue=0.7.0
 - intake=0.6.0
 - jpype1=1.0.2
 - jsonschema
@@ -31,13 +32,3 @@ dependencies:
 - tpot
 - tzlocal=2.1
 - uvicorn=0.11.3
-# fugue dependencies; remove when we conda install fugue
-- adagio
-- antlr4-python3-runtime<4.10
-- ciso8601
-- fs
-- pip
-- qpd
-- triad
-- pip:
-  - fugue[sql]==0.5.3

continuous_integration/environment-3.8-jdk8-dev.yaml

@@ -6,6 +6,7 @@ dependencies:
 - dask-ml=2022.1.22
 - dask=2022.3.0
 - fastapi=0.69.0
+- fugue=0.7.0
 - intake=0.6.0
 - jpype1=1.0.2
 - jsonschema
@@ -31,13 +32,3 @@ dependencies:
 - tpot
 - tzlocal=2.1
 - uvicorn=0.11.3
-# fugue dependencies; remove when we conda install fugue
-- adagio
-- antlr4-python3-runtime<4.10
-- ciso8601
-- fs
-- pip
-- qpd
-- triad
-- pip:
-  - fugue[sql]==0.5.3

continuous_integration/environment-3.9-jdk11-dev.yaml

@@ -6,6 +6,7 @@ dependencies:
 - dask-ml>=2022.1.22
 - dask>=2022.3.0
 - fastapi>=0.69.0
+- fugue>=0.7.0
 - intake>=0.6.0
 - jpype1>=1.0.2
 - jsonschema
@@ -31,14 +32,3 @@ dependencies:
 - tpot
 - tzlocal>=2.1
 - uvicorn>=0.11.3
-# fugue dependencies; remove when we conda install fugue
-- adagio
-- antlr4-python3-runtime<4.10
-- ciso8601
-- fs
-- pip
-- qpd
-- triad
-- pip:
-  # FIXME: tests are failing with fugue 0.7.0
-  - fugue[sql]>=0.5.3,<0.7.0

continuous_integration/environment-3.9-jdk8-dev.yaml

@@ -6,6 +6,7 @@ dependencies:
 - dask-ml>=2022.1.22
 - dask>=2022.3.0
 - fastapi>=0.69.0
+- fugue>=0.7.0
 - intake>=0.6.0
 - jpype1>=1.0.2
 - jsonschema
@@ -31,14 +32,3 @@ dependencies:
 - tpot
 - tzlocal>=2.1
 - uvicorn>=0.11.3
-# fugue dependencies; remove when we conda install fugue
-- adagio
-- antlr4-python3-runtime<4.10
-- ciso8601
-- fs
-- pip
-- qpd
-- triad
-- pip:
-  # FIXME: tests are failing with fugue 0.7.0
-  - fugue[sql]>=0.5.3,<0.7.0

dask_sql/integrations/fugue.py

@@ -1,8 +1,10 @@
 try:
     import fugue
     import fugue_dask
+    from dask.distributed import Client
     from fugue import WorkflowDataFrame, register_execution_engine
     from fugue_sql import FugueSQLWorkflow
+    from triad import run_at_def
     from triad.utils.convert import get_caller_global_local_vars
 except ImportError:  # pragma: no cover
     raise ImportError(
@@ -15,9 +17,25 @@
 
 from dask_sql.context import Context
 
-register_execution_engine(
-    "dask", lambda conf: DaskSQLExecutionEngine(conf), on_dup="overwrite"
-)
+
+@run_at_def
+def _register_engines() -> None:
+    """Register (overwrite) the default Dask execution engine of Fugue. This
+    function is invoked as an entrypoint, users don't need to call it explicitly.
+    """
+    register_execution_engine(
+        "dask",
+        lambda conf, **kwargs: DaskSQLExecutionEngine(conf=conf),
+        on_dup="overwrite",
+    )
+
+    register_execution_engine(
+        Client,
+        lambda engine, conf, **kwargs: DaskSQLExecutionEngine(
+            dask_client=engine, conf=conf
+        ),
+        on_dup="overwrite",
+    )
 
 
 class DaskSQLEngine(fugue.execution.execution_engine.SQLEngine):

docs/environment.yml

@@ -10,8 +10,7 @@ dependencies:
   - maven>=3.6.0
   - dask>=2022.3.0
   - pandas>=1.1.2
-  # FIXME: tests are failing with fugue 0.7.0
-  - fugue>=0.5.3,<0.7.0
+  - fugue>=0.7.0
   - jpype1>=1.0.2
   - fastapi>=0.69.0
   - uvicorn>=0.11.3

docs/requirements-docs.txt

@@ -3,8 +3,7 @@ sphinx-tabs
 dask-sphinx-theme>=3.0.0
 dask>=2022.3.0
 pandas>=1.1.2
-# FIXME: tests are failing with fugue 0.7.0
-fugue>=0.5.3,<0.7.0
+fugue>=0.7.0
 jpype1>=1.0.2
 fastapi>=0.69.0
 uvicorn>=0.11.3

docs/source/fugue.rst

@@ -8,7 +8,7 @@ In order to offer a "best of both worlds" solution, dask-sql includes several op
 dask-sql as a FugueSQL engine
 -----------------------------
 
-FugueSQL users unfamiliar with dask-sql can take advantage of its functionality with minimal code changes by passing :class:`dask_sql.integrations.fugue.DaskSQLExecutionEngine` into the ``FugueSQLWorkflow`` being used to execute commands.
+FugueSQL users unfamiliar with dask-sql can take advantage of its functionality by installing it in an environment alongside Fugue; this will automatically register :class:`dask_sql.integrations.fugue.DaskSQLExecutionEngine` as the default Dask execution engine for FugueSQL queries.
 For more information and sample usage, see `Fugue — dask-sql as a FugueSQL engine <https://fugue-tutorials.readthedocs.io/tutorials/integrations/dasksql.html>`_.
 
 Using FugueSQL on an existing ``Context``