Repository: great-expectations/great_expectations

Issue:
SparkDFDatasource can't use DatabricksTableGenerator
I'm trying to use `SparkDFDatasource` with a `DatabricksTableGenerator`. To me , it seems multiple problems are happening:

* `Datasource.add_generator` doesn't pass `type_` although `DatabricksTableGenerator requires it:
https://github.com/great-expectations/great_expectations/blob/master/great_expectations/datasource/datasource.py#L202

* `Datasource.get_batch` takes `batch_kwargs` as a `dict`, passing it to `SparkDFDatasource._get_data_asset`, which then tries `batch_kwargs.query`, rather than `batch_kwargs["query"]`:
https://github.com/great-expectations/great_expectations/blob/master/great_expectations/datasource/spark_source.py#L97

* In general, I'm not sure why `DatabricksTableGenerator._get_iterator()` is returning `dict_keyiterator` object. Shouldn't it give back the batch_kwargs dict?
https://github.com/great-expectations/great_expectations/blob/master/great_expectations/datasource/generator/databricks_generator.py#L28


Discussion:
Thanks for flagging -- I'm taking a look at this now.
You're exactly right on all three counts. We're planning to add integration testing for databricks but haven't yet. I'm going to make these changes and add some basic tests around them for the new release on Monday.

Resolving patch:
diff --git a/great_expectations/datasource/generator/__init__.py b/great_expectations/datasource/generator/__init__.py
--- a/great_expectations/datasource/generator/__init__.py
+++ b/great_expectations/datasource/generator/__init__.py
@@ -0,0 +1 @@
+from .databricks_generator import DatabricksTableGenerator
\ No newline at end of file
diff --git a/great_expectations/datasource/generator/databricks_generator.py b/great_expectations/datasource/generator/databricks_generator.py
--- a/great_expectations/datasource/generator/databricks_generator.py
+++ b/great_expectations/datasource/generator/databricks_generator.py
@@ -15,13 +15,22 @@ class DatabricksTableGenerator(BatchGenerator):
     """Meant to be used in a Databricks notebook
     """
 
-    def __init__(self, name, type_, datasource, database):
-        super(DatabricksTableGenerator, self).__init__(name, type_, datasource)
-        # this should grab the already instantiated SparkSession available on Databricks notebooks
-        self.spark = datasource.spark
+    def __init__(self, name="default",
+                 datasource=None,
+                 database="default"):
+        super(DatabricksTableGenerator, self).__init__(name, type_="databricks", datasource=datasource)
         self.database = database
+        try:
+            self.spark = SparkSession.builder.getOrCreate()
+        except Exception:
+            logger.error("Unable to load spark context; install optional spark dependency for support.")
+            self.spark = None
 
     def get_available_data_asset_names(self):
+        if self.spark is None:
+            logger.warning("No sparkSession available to query for tables.")
+            return set()
+
         tables = self.spark.sql('show tables in {}'.format(self.database))
         return set([row.tableName for row in tables.collect()])
 
@@ -31,9 +40,9 @@ def _get_iterator(self, data_asset_name, **kwargs):
             if not kwargs.get('date_field'):
                 raise Exception('Must specify date_field when using partition.')
             query += ' where {} = "{}"'.format(kwargs.get('date_field'), kwargs.get('partition'))
-        return iter(
+        return iter([
             {
                 "query": query,
                 "timestamp": time.time()
             }
-        )
+        ])
diff --git a/great_expectations/datasource/spark_source.py b/great_expectations/datasource/spark_source.py
--- a/great_expectations/datasource/spark_source.py
+++ b/great_expectations/datasource/spark_source.py
@@ -126,7 +126,7 @@ def _get_data_asset(self, batch_kwargs, expectation_suite, caching=True, **kwarg
                 raise BatchKwargsError("Unsupported reader: %s" % reader_method.name, batch_kwargs)
             
         elif "query" in batch_kwargs:
-            df = self.spark.sql(batch_kwargs.query)
+            df = self.spark.sql(batch_kwargs["query"])
 
         elif "df" in batch_kwargs and isinstance(batch_kwargs["df"], (DataFrame, SparkDFDataset)):
             df = batch_kwargs.pop("df")  # We don't want to store the actual DataFrame in kwargs


Test patch:
diff --git a/tests/datasource/test_batch_generators.py b/tests/datasource/test_batch_generators.py
--- a/tests/datasource/test_batch_generators.py
+++ b/tests/datasource/test_batch_generators.py
@@ -4,7 +4,7 @@
 
 from great_expectations.exceptions import DataContextError
 from great_expectations.datasource import SubdirReaderGenerator, GlobReaderGenerator
-
+from great_expectations.datasource.generator import DatabricksTableGenerator
 
 def test_file_kwargs_generator(data_context, filesystem_csv):
     base_dir = filesystem_csv
@@ -140,3 +140,16 @@ def test_file_kwargs_generator_extensions(tmp_path_factory):
 
     g1_assets = g1.get_available_data_asset_names()
     assert g1_assets == {"f2", "f4", "f6", "f7", "f8", "f9", "f0"}
+
+
+def test_databricks_generator():
+    generator = DatabricksTableGenerator()
+    available_assets = generator.get_available_data_asset_names()
+
+    # We have no tables available
+    assert available_assets == set()
+
+    databricks_kwargs_iterator = generator.get_iterator("foo")
+    kwargs = [batch_kwargs for batch_kwargs in databricks_kwargs_iterator]
+    assert "timestamp" in kwargs[0]
+    assert "select * from" in kwargs[0]["query"].lower()