We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN 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

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `great_expectations/datasource/spark_source.py`
Content:
```
1 import logging
2 import time
3 from six import string_types
4 
5 from ..exceptions import BatchKwargsError
6 
7 from .datasource import Datasource, ReaderMethods
8 from great_expectations.datasource.generator.filesystem_path_generator import SubdirReaderGenerator
9 from great_expectations.datasource.generator.databricks_generator import DatabricksTableGenerator
10 from great_expectations.datasource.generator.in_memory_generator import InMemoryGenerator
11 
12 from great_expectations.data_context.types import ClassConfig
13 
14 logger = logging.getLogger(__name__)
15 
16 try:
17     from great_expectations.dataset.sparkdf_dataset import SparkDFDataset
18     from pyspark.sql import SparkSession, DataFrame
19 except ImportError:
20     # TODO: review logging more detail here
21     logger.debug("Unable to load pyspark; install optional spark dependency for support.")
22 
23 
24 class SparkDFDatasource(Datasource):
25     """The SparkDFDatasource produces SparkDFDatasets and supports generators capable of interacting with local
26     filesystem (the default subdir_reader generator) and databricks notebooks.
27     """
28 
29     def __init__(self, name="default", data_context=None, data_asset_type=None, generators=None, **kwargs):
30         if generators is None:
31             # Provide a gentle way to build a datasource with a sane default,
32             # including ability to specify the base_directory
33             base_directory = kwargs.pop("base_directory", "/data")
34             reader_options = kwargs.pop("reader_options", {})
35             generators = {
36                 "default": {
37                     "type": "subdir_reader",
38                     "base_directory": base_directory,
39                     "reader_options": reader_options
40                 }
41         }
42 
43         if data_asset_type is None:
44             data_asset_type = ClassConfig(
45                 class_name="SparkDFDataset"
46             )
47         else:
48             try:
49                 data_asset_type = ClassConfig(**data_asset_type)
50             except TypeError:
51                 # In this case, we allow the passed config, for now, in case they're using a legacy string-only config
52                 pass
53 
54         super(SparkDFDatasource, self).__init__(name, type_="spark",
55                                                 data_context=data_context,
56                                                 data_asset_type=data_asset_type,
57                                                 generators=generators)
58         try:
59             self.spark = SparkSession.builder.getOrCreate()
60         except Exception:
61             logger.error("Unable to load spark context; install optional spark dependency for support.")
62             self.spark = None
63 
64         self._build_generators()
65 
66     def _get_generator_class(self, type_):
67         if type_ == "subdir_reader":
68             return SubdirReaderGenerator
69         elif type_ == "databricks":
70             return DatabricksTableGenerator
71         elif type_ == "memory":
72             return InMemoryGenerator
73         else:
74             raise ValueError("Unrecognized BatchGenerator type %s" % type_)
75 
76     def _get_data_asset(self, batch_kwargs, expectation_suite, caching=True, **kwargs):
77         """class-private implementation of get_data_asset"""
78         if self.spark is None:
79             logger.error("No spark session available")
80             return None
81 
82         batch_kwargs.update(kwargs)
83         reader_options = batch_kwargs.copy()
84 
85         if "data_asset_type" in reader_options:
86             data_asset_type_config = reader_options.pop("data_asset_type")  # Get and remove the config
87             try:
88                 data_asset_type_config = ClassConfig(**data_asset_type_config)
89             except TypeError:
90                 # We tried; we'll pass the config downstream, probably as a string, and handle an error later
91                 pass
92         else:
93             data_asset_type_config = self._data_asset_type
94 
95         data_asset_type = self._get_data_asset_class(data_asset_type_config)
96         if not issubclass(data_asset_type, SparkDFDataset):
97             raise ValueError("SparkDFDatasource cannot instantiate batch with data_asset_type: '%s'. It "
98                              "must be a subclass of SparkDFDataset." % data_asset_type.__name__)
99 
100         if "path" in batch_kwargs:
101             path = reader_options.pop("path")  # We remove this so it is not used as a reader option
102             reader_options.pop("timestamp", "")    # ditto timestamp (but missing ok)
103             reader_method = reader_options.pop("reader_method", None)
104             if reader_method is None:
105                 reader_method = self._guess_reader_method_from_path(path)
106                 if reader_method is None:
107                     raise BatchKwargsError("Unable to determine reader for path: %s" % path, batch_kwargs)
108             else:
109                 try:
110                     reader_method = ReaderMethods[reader_method]
111                 except KeyError:
112                     raise BatchKwargsError("Unknown reader method: %s" % reader_method, batch_kwargs)
113 
114             reader = self.spark.read
115 
116             for option in reader_options.items():
117                 reader = reader.option(*option)
118 
119             if reader_method == ReaderMethods.CSV:
120                 df = reader.csv(path)
121             elif reader_method == ReaderMethods.parquet:
122                 df = reader.parquet(path)
123             elif reader_method == ReaderMethods.delta:
124                 df = reader.format("delta").load(path)
125             else:
126                 raise BatchKwargsError("Unsupported reader: %s" % reader_method.name, batch_kwargs)
127             
128         elif "query" in batch_kwargs:
129             df = self.spark.sql(batch_kwargs.query)
130 
131         elif "df" in batch_kwargs and isinstance(batch_kwargs["df"], (DataFrame, SparkDFDataset)):
132             df = batch_kwargs.pop("df")  # We don't want to store the actual DataFrame in kwargs
133             if isinstance(df, SparkDFDataset):
134                 # Grab just the spark_df reference, since we want to override everything else
135                 df = df.spark_df
136             batch_kwargs["SparkDFRef"] = True
137         else:
138             raise BatchKwargsError("Unrecognized batch_kwargs for spark_source", batch_kwargs)
139 
140         return data_asset_type(df,
141                                expectation_suite=expectation_suite,
142                                data_context=self._data_context,
143                                batch_kwargs=batch_kwargs,
144                                caching=caching)
145 
146     def build_batch_kwargs(self, *args, **kwargs):
147         if len(args) > 0:
148             if isinstance(args[0], (DataFrame, SparkDFDataset)):
149                 kwargs.update({
150                     "df": args[0],
151                     "timestamp": time.time()
152                 })
153             elif isinstance(args[0], string_types):
154                 kwargs.update({
155                     "path": args[0],
156                     "timestamp": time.time()
157                 })
158         else:
159             kwargs.update({
160                 "timestamp": time.time()
161             })
162         return kwargs
163
```

Path: `great_expectations/datasource/generator/databricks_generator.py`
Content:
```
1 import time
2 import logging
3 
4 from great_expectations.datasource.generator.batch_generator import BatchGenerator
5 
6 logger = logging.getLogger(__name__)
7 
8 try:
9     from pyspark.sql import SparkSession
10 except ImportError:
11     logger.debug("Unable to load spark context; install optional spark dependency for support.")
12 
13 
14 class DatabricksTableGenerator(BatchGenerator):
15     """Meant to be used in a Databricks notebook
16     """
17 
18     def __init__(self, name, type_, datasource, database):
19         super(DatabricksTableGenerator, self).__init__(name, type_, datasource)
20         # this should grab the already instantiated SparkSession available on Databricks notebooks
21         self.spark = datasource.spark
22         self.database = database
23 
24     def get_available_data_asset_names(self):
25         tables = self.spark.sql('show tables in {}'.format(self.database))
26         return set([row.tableName for row in tables.collect()])
27 
28     def _get_iterator(self, data_asset_name, **kwargs):
29         query = 'select * from {}.{}'.format(self.database, data_asset_name)
30         if kwargs.get('partition'):
31             if not kwargs.get('date_field'):
32                 raise Exception('Must specify date_field when using partition.')
33             query += ' where {} = "{}"'.format(kwargs.get('date_field'), kwargs.get('partition'))
34         return iter(
35             {
36                 "query": query,
37                 "timestamp": time.time()
38             }
39         )
40
```

Path: `great_expectations/datasource/generator/__init__.py`
Content:
```
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

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 @@
     """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 @@
             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 @@
                 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

{"golden_diff": "diff --git a/great_expectations/datasource/generator/__init__.py b/great_expectations/datasource/generator/__init__.py\n--- a/great_expectations/datasource/generator/__init__.py\n+++ b/great_expectations/datasource/generator/__init__.py\n@@ -0,0 +1 @@\n+from .databricks_generator import DatabricksTableGenerator\n\\ No newline at end of file\ndiff --git a/great_expectations/datasource/generator/databricks_generator.py b/great_expectations/datasource/generator/databricks_generator.py\n--- a/great_expectations/datasource/generator/databricks_generator.py\n+++ b/great_expectations/datasource/generator/databricks_generator.py\n@@ -15,13 +15,22 @@\n     \"\"\"Meant to be used in a Databricks notebook\n     \"\"\"\n \n-    def __init__(self, name, type_, datasource, database):\n-        super(DatabricksTableGenerator, self).__init__(name, type_, datasource)\n-        # this should grab the already instantiated SparkSession available on Databricks notebooks\n-        self.spark = datasource.spark\n+    def __init__(self, name=\"default\",\n+                 datasource=None,\n+                 database=\"default\"):\n+        super(DatabricksTableGenerator, self).__init__(name, type_=\"databricks\", datasource=datasource)\n         self.database = database\n+        try:\n+            self.spark = SparkSession.builder.getOrCreate()\n+        except Exception:\n+            logger.error(\"Unable to load spark context; install optional spark dependency for support.\")\n+            self.spark = None\n \n     def get_available_data_asset_names(self):\n+        if self.spark is None:\n+            logger.warning(\"No sparkSession available to query for tables.\")\n+            return set()\n+\n         tables = self.spark.sql('show tables in {}'.format(self.database))\n         return set([row.tableName for row in tables.collect()])\n \n@@ -31,9 +40,9 @@\n             if not kwargs.get('date_field'):\n                 raise Exception('Must specify date_field when using partition.')\n             query += ' where {} = \"{}\"'.format(kwargs.get('date_field'), kwargs.get('partition'))\n-        return iter(\n+        return iter([\n             {\n                 \"query\": query,\n                 \"timestamp\": time.time()\n             }\n-        )\n+        ])\ndiff --git a/great_expectations/datasource/spark_source.py b/great_expectations/datasource/spark_source.py\n--- a/great_expectations/datasource/spark_source.py\n+++ b/great_expectations/datasource/spark_source.py\n@@ -126,7 +126,7 @@\n                 raise BatchKwargsError(\"Unsupported reader: %s\" % reader_method.name, batch_kwargs)\n             \n         elif \"query\" in batch_kwargs:\n-            df = self.spark.sql(batch_kwargs.query)\n+            df = self.spark.sql(batch_kwargs[\"query\"])\n \n         elif \"df\" in batch_kwargs and isinstance(batch_kwargs[\"df\"], (DataFrame, SparkDFDataset)):\n             df = batch_kwargs.pop(\"df\")  # We don't want to store the actual DataFrame in kwargs\n", "issue": "SparkDFDatasource can't use DatabricksTableGenerator\nI'm trying to use `SparkDFDatasource` with a `DatabricksTableGenerator`. To me , it seems multiple problems are happening:\r\n\r\n* `Datasource.add_generator` doesn't pass `type_` although `DatabricksTableGenerator requires it:\r\nhttps://github.com/great-expectations/great_expectations/blob/master/great_expectations/datasource/datasource.py#L202\r\n\r\n* `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\"]`:\r\nhttps://github.com/great-expectations/great_expectations/blob/master/great_expectations/datasource/spark_source.py#L97\r\n\r\n* In general, I'm not sure why `DatabricksTableGenerator._get_iterator()` is returning `dict_keyiterator` object. Shouldn't it give back the batch_kwargs dict?\r\nhttps://github.com/great-expectations/great_expectations/blob/master/great_expectations/datasource/generator/databricks_generator.py#L28\n", "before_files": [{"content": "import logging\nimport time\nfrom six import string_types\n\nfrom ..exceptions import BatchKwargsError\n\nfrom .datasource import Datasource, ReaderMethods\nfrom great_expectations.datasource.generator.filesystem_path_generator import SubdirReaderGenerator\nfrom great_expectations.datasource.generator.databricks_generator import DatabricksTableGenerator\nfrom great_expectations.datasource.generator.in_memory_generator import InMemoryGenerator\n\nfrom great_expectations.data_context.types import ClassConfig\n\nlogger = logging.getLogger(__name__)\n\ntry:\n    from great_expectations.dataset.sparkdf_dataset import SparkDFDataset\n    from pyspark.sql import SparkSession, DataFrame\nexcept ImportError:\n    # TODO: review logging more detail here\n    logger.debug(\"Unable to load pyspark; install optional spark dependency for support.\")\n\n\nclass SparkDFDatasource(Datasource):\n    \"\"\"The SparkDFDatasource produces SparkDFDatasets and supports generators capable of interacting with local\n    filesystem (the default subdir_reader generator) and databricks notebooks.\n    \"\"\"\n\n    def __init__(self, name=\"default\", data_context=None, data_asset_type=None, generators=None, **kwargs):\n        if generators is None:\n            # Provide a gentle way to build a datasource with a sane default,\n            # including ability to specify the base_directory\n            base_directory = kwargs.pop(\"base_directory\", \"/data\")\n            reader_options = kwargs.pop(\"reader_options\", {})\n            generators = {\n                \"default\": {\n                    \"type\": \"subdir_reader\",\n                    \"base_directory\": base_directory,\n                    \"reader_options\": reader_options\n                }\n        }\n\n        if data_asset_type is None:\n            data_asset_type = ClassConfig(\n                class_name=\"SparkDFDataset\"\n            )\n        else:\n            try:\n                data_asset_type = ClassConfig(**data_asset_type)\n            except TypeError:\n                # In this case, we allow the passed config, for now, in case they're using a legacy string-only config\n                pass\n\n        super(SparkDFDatasource, self).__init__(name, type_=\"spark\",\n                                                data_context=data_context,\n                                                data_asset_type=data_asset_type,\n                                                generators=generators)\n        try:\n            self.spark = SparkSession.builder.getOrCreate()\n        except Exception:\n            logger.error(\"Unable to load spark context; install optional spark dependency for support.\")\n            self.spark = None\n\n        self._build_generators()\n\n    def _get_generator_class(self, type_):\n        if type_ == \"subdir_reader\":\n            return SubdirReaderGenerator\n        elif type_ == \"databricks\":\n            return DatabricksTableGenerator\n        elif type_ == \"memory\":\n            return InMemoryGenerator\n        else:\n            raise ValueError(\"Unrecognized BatchGenerator type %s\" % type_)\n\n    def _get_data_asset(self, batch_kwargs, expectation_suite, caching=True, **kwargs):\n        \"\"\"class-private implementation of get_data_asset\"\"\"\n        if self.spark is None:\n            logger.error(\"No spark session available\")\n            return None\n\n        batch_kwargs.update(kwargs)\n        reader_options = batch_kwargs.copy()\n\n        if \"data_asset_type\" in reader_options:\n            data_asset_type_config = reader_options.pop(\"data_asset_type\")  # Get and remove the config\n            try:\n                data_asset_type_config = ClassConfig(**data_asset_type_config)\n            except TypeError:\n                # We tried; we'll pass the config downstream, probably as a string, and handle an error later\n                pass\n        else:\n            data_asset_type_config = self._data_asset_type\n\n        data_asset_type = self._get_data_asset_class(data_asset_type_config)\n        if not issubclass(data_asset_type, SparkDFDataset):\n            raise ValueError(\"SparkDFDatasource cannot instantiate batch with data_asset_type: '%s'. It \"\n                             \"must be a subclass of SparkDFDataset.\" % data_asset_type.__name__)\n\n        if \"path\" in batch_kwargs:\n            path = reader_options.pop(\"path\")  # We remove this so it is not used as a reader option\n            reader_options.pop(\"timestamp\", \"\")    # ditto timestamp (but missing ok)\n            reader_method = reader_options.pop(\"reader_method\", None)\n            if reader_method is None:\n                reader_method = self._guess_reader_method_from_path(path)\n                if reader_method is None:\n                    raise BatchKwargsError(\"Unable to determine reader for path: %s\" % path, batch_kwargs)\n            else:\n                try:\n                    reader_method = ReaderMethods[reader_method]\n                except KeyError:\n                    raise BatchKwargsError(\"Unknown reader method: %s\" % reader_method, batch_kwargs)\n\n            reader = self.spark.read\n\n            for option in reader_options.items():\n                reader = reader.option(*option)\n\n            if reader_method == ReaderMethods.CSV:\n                df = reader.csv(path)\n            elif reader_method == ReaderMethods.parquet:\n                df = reader.parquet(path)\n            elif reader_method == ReaderMethods.delta:\n                df = reader.format(\"delta\").load(path)\n            else:\n                raise BatchKwargsError(\"Unsupported reader: %s\" % reader_method.name, batch_kwargs)\n            \n        elif \"query\" in batch_kwargs:\n            df = self.spark.sql(batch_kwargs.query)\n\n        elif \"df\" in batch_kwargs and isinstance(batch_kwargs[\"df\"], (DataFrame, SparkDFDataset)):\n            df = batch_kwargs.pop(\"df\")  # We don't want to store the actual DataFrame in kwargs\n            if isinstance(df, SparkDFDataset):\n                # Grab just the spark_df reference, since we want to override everything else\n                df = df.spark_df\n            batch_kwargs[\"SparkDFRef\"] = True\n        else:\n            raise BatchKwargsError(\"Unrecognized batch_kwargs for spark_source\", batch_kwargs)\n\n        return data_asset_type(df,\n                               expectation_suite=expectation_suite,\n                               data_context=self._data_context,\n                               batch_kwargs=batch_kwargs,\n                               caching=caching)\n\n    def build_batch_kwargs(self, *args, **kwargs):\n        if len(args) > 0:\n            if isinstance(args[0], (DataFrame, SparkDFDataset)):\n                kwargs.update({\n                    \"df\": args[0],\n                    \"timestamp\": time.time()\n                })\n            elif isinstance(args[0], string_types):\n                kwargs.update({\n                    \"path\": args[0],\n                    \"timestamp\": time.time()\n                })\n        else:\n            kwargs.update({\n                \"timestamp\": time.time()\n            })\n        return kwargs\n", "path": "great_expectations/datasource/spark_source.py"}, {"content": "import time\nimport logging\n\nfrom great_expectations.datasource.generator.batch_generator import BatchGenerator\n\nlogger = logging.getLogger(__name__)\n\ntry:\n    from pyspark.sql import SparkSession\nexcept ImportError:\n    logger.debug(\"Unable to load spark context; install optional spark dependency for support.\")\n\n\nclass DatabricksTableGenerator(BatchGenerator):\n    \"\"\"Meant to be used in a Databricks notebook\n    \"\"\"\n\n    def __init__(self, name, type_, datasource, database):\n        super(DatabricksTableGenerator, self).__init__(name, type_, datasource)\n        # this should grab the already instantiated SparkSession available on Databricks notebooks\n        self.spark = datasource.spark\n        self.database = database\n\n    def get_available_data_asset_names(self):\n        tables = self.spark.sql('show tables in {}'.format(self.database))\n        return set([row.tableName for row in tables.collect()])\n\n    def _get_iterator(self, data_asset_name, **kwargs):\n        query = 'select * from {}.{}'.format(self.database, data_asset_name)\n        if kwargs.get('partition'):\n            if not kwargs.get('date_field'):\n                raise Exception('Must specify date_field when using partition.')\n            query += ' where {} = \"{}\"'.format(kwargs.get('date_field'), kwargs.get('partition'))\n        return iter(\n            {\n                \"query\": query,\n                \"timestamp\": time.time()\n            }\n        )\n", "path": "great_expectations/datasource/generator/databricks_generator.py"}, {"content": "", "path": "great_expectations/datasource/generator/__init__.py"}], "after_files": [{"content": "import logging\nimport time\nfrom six import string_types\n\nfrom ..exceptions import BatchKwargsError\n\nfrom .datasource import Datasource, ReaderMethods\nfrom great_expectations.datasource.generator.filesystem_path_generator import SubdirReaderGenerator\nfrom great_expectations.datasource.generator.databricks_generator import DatabricksTableGenerator\nfrom great_expectations.datasource.generator.in_memory_generator import InMemoryGenerator\n\nfrom great_expectations.data_context.types import ClassConfig\n\nlogger = logging.getLogger(__name__)\n\ntry:\n    from great_expectations.dataset.sparkdf_dataset import SparkDFDataset\n    from pyspark.sql import SparkSession, DataFrame\nexcept ImportError:\n    # TODO: review logging more detail here\n    logger.debug(\"Unable to load pyspark; install optional spark dependency for support.\")\n\n\nclass SparkDFDatasource(Datasource):\n    \"\"\"The SparkDFDatasource produces SparkDFDatasets and supports generators capable of interacting with local\n    filesystem (the default subdir_reader generator) and databricks notebooks.\n    \"\"\"\n\n    def __init__(self, name=\"default\", data_context=None, data_asset_type=None, generators=None, **kwargs):\n        if generators is None:\n            # Provide a gentle way to build a datasource with a sane default,\n            # including ability to specify the base_directory\n            base_directory = kwargs.pop(\"base_directory\", \"/data\")\n            reader_options = kwargs.pop(\"reader_options\", {})\n            generators = {\n                \"default\": {\n                    \"type\": \"subdir_reader\",\n                    \"base_directory\": base_directory,\n                    \"reader_options\": reader_options\n                }\n        }\n\n        if data_asset_type is None:\n            data_asset_type = ClassConfig(\n                class_name=\"SparkDFDataset\"\n            )\n        else:\n            try:\n                data_asset_type = ClassConfig(**data_asset_type)\n            except TypeError:\n                # In this case, we allow the passed config, for now, in case they're using a legacy string-only config\n                pass\n\n        super(SparkDFDatasource, self).__init__(name, type_=\"spark\",\n                                                data_context=data_context,\n                                                data_asset_type=data_asset_type,\n                                                generators=generators)\n        try:\n            self.spark = SparkSession.builder.getOrCreate()\n        except Exception:\n            logger.error(\"Unable to load spark context; install optional spark dependency for support.\")\n            self.spark = None\n\n        self._build_generators()\n\n    def _get_generator_class(self, type_):\n        if type_ == \"subdir_reader\":\n            return SubdirReaderGenerator\n        elif type_ == \"databricks\":\n            return DatabricksTableGenerator\n        elif type_ == \"memory\":\n            return InMemoryGenerator\n        else:\n            raise ValueError(\"Unrecognized BatchGenerator type %s\" % type_)\n\n    def _get_data_asset(self, batch_kwargs, expectation_suite, caching=True, **kwargs):\n        \"\"\"class-private implementation of get_data_asset\"\"\"\n        if self.spark is None:\n            logger.error(\"No spark session available\")\n            return None\n\n        batch_kwargs.update(kwargs)\n        reader_options = batch_kwargs.copy()\n\n        if \"data_asset_type\" in reader_options:\n            data_asset_type_config = reader_options.pop(\"data_asset_type\")  # Get and remove the config\n            try:\n                data_asset_type_config = ClassConfig(**data_asset_type_config)\n            except TypeError:\n                # We tried; we'll pass the config downstream, probably as a string, and handle an error later\n                pass\n        else:\n            data_asset_type_config = self._data_asset_type\n\n        data_asset_type = self._get_data_asset_class(data_asset_type_config)\n        if not issubclass(data_asset_type, SparkDFDataset):\n            raise ValueError(\"SparkDFDatasource cannot instantiate batch with data_asset_type: '%s'. It \"\n                             \"must be a subclass of SparkDFDataset.\" % data_asset_type.__name__)\n\n        if \"path\" in batch_kwargs:\n            path = reader_options.pop(\"path\")  # We remove this so it is not used as a reader option\n            reader_options.pop(\"timestamp\", \"\")    # ditto timestamp (but missing ok)\n            reader_method = reader_options.pop(\"reader_method\", None)\n            if reader_method is None:\n                reader_method = self._guess_reader_method_from_path(path)\n                if reader_method is None:\n                    raise BatchKwargsError(\"Unable to determine reader for path: %s\" % path, batch_kwargs)\n            else:\n                try:\n                    reader_method = ReaderMethods[reader_method]\n                except KeyError:\n                    raise BatchKwargsError(\"Unknown reader method: %s\" % reader_method, batch_kwargs)\n\n            reader = self.spark.read\n\n            for option in reader_options.items():\n                reader = reader.option(*option)\n\n            if reader_method == ReaderMethods.CSV:\n                df = reader.csv(path)\n            elif reader_method == ReaderMethods.parquet:\n                df = reader.parquet(path)\n            elif reader_method == ReaderMethods.delta:\n                df = reader.format(\"delta\").load(path)\n            else:\n                raise BatchKwargsError(\"Unsupported reader: %s\" % reader_method.name, batch_kwargs)\n            \n        elif \"query\" in batch_kwargs:\n            df = self.spark.sql(batch_kwargs[\"query\"])\n\n        elif \"df\" in batch_kwargs and isinstance(batch_kwargs[\"df\"], (DataFrame, SparkDFDataset)):\n            df = batch_kwargs.pop(\"df\")  # We don't want to store the actual DataFrame in kwargs\n            if isinstance(df, SparkDFDataset):\n                # Grab just the spark_df reference, since we want to override everything else\n                df = df.spark_df\n            batch_kwargs[\"SparkDFRef\"] = True\n        else:\n            raise BatchKwargsError(\"Unrecognized batch_kwargs for spark_source\", batch_kwargs)\n\n        return data_asset_type(df,\n                               expectation_suite=expectation_suite,\n                               data_context=self._data_context,\n                               batch_kwargs=batch_kwargs,\n                               caching=caching)\n\n    def build_batch_kwargs(self, *args, **kwargs):\n        if len(args) > 0:\n            if isinstance(args[0], (DataFrame, SparkDFDataset)):\n                kwargs.update({\n                    \"df\": args[0],\n                    \"timestamp\": time.time()\n                })\n            elif isinstance(args[0], string_types):\n                kwargs.update({\n                    \"path\": args[0],\n                    \"timestamp\": time.time()\n                })\n        else:\n            kwargs.update({\n                \"timestamp\": time.time()\n            })\n        return kwargs\n", "path": "great_expectations/datasource/spark_source.py"}, {"content": "import time\nimport logging\n\nfrom great_expectations.datasource.generator.batch_generator import BatchGenerator\n\nlogger = logging.getLogger(__name__)\n\ntry:\n    from pyspark.sql import SparkSession\nexcept ImportError:\n    logger.debug(\"Unable to load spark context; install optional spark dependency for support.\")\n\n\nclass DatabricksTableGenerator(BatchGenerator):\n    \"\"\"Meant to be used in a Databricks notebook\n    \"\"\"\n\n    def __init__(self, name=\"default\",\n                 datasource=None,\n                 database=\"default\"):\n        super(DatabricksTableGenerator, self).__init__(name, type_=\"databricks\", datasource=datasource)\n        self.database = database\n        try:\n            self.spark = SparkSession.builder.getOrCreate()\n        except Exception:\n            logger.error(\"Unable to load spark context; install optional spark dependency for support.\")\n            self.spark = None\n\n    def get_available_data_asset_names(self):\n        if self.spark is None:\n            logger.warning(\"No sparkSession available to query for tables.\")\n            return set()\n\n        tables = self.spark.sql('show tables in {}'.format(self.database))\n        return set([row.tableName for row in tables.collect()])\n\n    def _get_iterator(self, data_asset_name, **kwargs):\n        query = 'select * from {}.{}'.format(self.database, data_asset_name)\n        if kwargs.get('partition'):\n            if not kwargs.get('date_field'):\n                raise Exception('Must specify date_field when using partition.')\n            query += ' where {} = \"{}\"'.format(kwargs.get('date_field'), kwargs.get('partition'))\n        return iter([\n            {\n                \"query\": query,\n                \"timestamp\": time.time()\n            }\n        ])\n", "path": "great_expectations/datasource/generator/databricks_generator.py"}, {"content": "from .databricks_generator import DatabricksTableGenerator", "path": "great_expectations/datasource/generator/__init__.py"}]}