Repository: great-expectations/great_expectations

Issue:
Use cleaner solution for non-truncating division in python 2
Prefer `from __future__ import division` to `1.*x/y`


Discussion:
Done inline during improve_partitions push.

Resolving patch:
diff --git a/great_expectations/datasource/fluent/__init__.py b/great_expectations/datasource/fluent/__init__.py
--- a/great_expectations/datasource/fluent/__init__.py
+++ b/great_expectations/datasource/fluent/__init__.py
@@ -41,6 +41,9 @@
 from great_expectations.datasource.fluent.spark_filesystem_datasource import (
     SparkFilesystemDatasource,
 )
+from great_expectations.datasource.fluent.spark_dbfs_datasource import (
+    SparkDBFSDatasource,
+)
 from great_expectations.datasource.fluent.spark_s3_datasource import (
     SparkS3Datasource,
 )
diff --git a/great_expectations/datasource/fluent/spark_dbfs_datasource.py b/great_expectations/datasource/fluent/spark_dbfs_datasource.py
new file mode 100644
--- /dev/null
+++ b/great_expectations/datasource/fluent/spark_dbfs_datasource.py
@@ -0,0 +1,82 @@
+from __future__ import annotations
+
+import logging
+import re
+from typing import TYPE_CHECKING, Optional, Union
+
+from typing_extensions import Literal
+
+from great_expectations.core._docs_decorators import public_api
+from great_expectations.core.util import DBFSPath
+from great_expectations.datasource.fluent import SparkFilesystemDatasource
+from great_expectations.datasource.fluent.data_asset.data_connector import (
+    DBFSDataConnector,
+)
+from great_expectations.datasource.fluent.spark_file_path_datasource import (
+    CSVAsset,
+)
+
+if TYPE_CHECKING:
+    from great_expectations.datasource.fluent.interfaces import (
+        Sorter,
+        SortersDefinition,
+    )
+
+logger = logging.getLogger(__name__)
+
+
+@public_api
+class SparkDBFSDatasource(SparkFilesystemDatasource):
+    """Spark based Datasource for DataBricks File System (DBFS) based data assets."""
+
+    # instance attributes
+    type: Literal["spark_dbfs"] = "spark_dbfs"
+
+    @public_api
+    def add_csv_asset(
+        self,
+        name: str,
+        batching_regex: Optional[Union[re.Pattern, str]] = None,
+        glob_directive: str = "**/*",
+        order_by: Optional[SortersDefinition] = None,
+        **kwargs,
+    ) -> CSVAsset:  # type: ignore[valid-type]
+        """Adds a CSV DataAsset to the present "SparkDBFSDatasource" object.
+
+        Args:
+            name: The name of the CSV asset
+            batching_regex: regex pattern that matches CSV filenames that is used to label the batches
+            glob_directive: glob for selecting files in DBFS directory (defaults to `**/*`) or nested directories (e.g. `*/*/*.csv`)
+            order_by: sorting directive via either list[Sorter] or "+/-  key" syntax: +/- (a/de)scending; + default
+            # TODO: <Alex>ALEX</Alex>
+            kwargs: Extra keyword arguments should correspond to ``pandas.read_csv`` keyword args
+            # TODO: <Alex>ALEX</Alex>
+        """
+        batching_regex_pattern: re.Pattern = self.parse_batching_regex_string(
+            batching_regex=batching_regex
+        )
+        order_by_sorters: list[Sorter] = self.parse_order_by_sorters(order_by=order_by)
+        asset = CSVAsset(
+            name=name,
+            batching_regex=batching_regex_pattern,
+            order_by=order_by_sorters,
+            **kwargs,
+        )
+        asset._data_connector = DBFSDataConnector.build_data_connector(
+            datasource_name=self.name,
+            data_asset_name=name,
+            batching_regex=batching_regex_pattern,
+            base_directory=self.base_directory,
+            glob_directive=glob_directive,
+            data_context_root_directory=self.data_context_root_directory,
+            file_path_template_map_fn=DBFSPath.convert_to_protocol_version,
+        )
+        asset._test_connection_error_message = (
+            DBFSDataConnector.build_test_connection_error_message(
+                data_asset_name=name,
+                batching_regex=batching_regex_pattern,
+                glob_directive=glob_directive,
+                base_directory=self.base_directory,
+            )
+        )
+        return self.add_asset(asset=asset)


Test patch:
diff --git a/tests/datasource/data_connector/test_configured_asset_dbfs_data_connector.py b/tests/datasource/data_connector/test_configured_asset_dbfs_data_connector.py
--- a/tests/datasource/data_connector/test_configured_asset_dbfs_data_connector.py
+++ b/tests/datasource/data_connector/test_configured_asset_dbfs_data_connector.py
@@ -1,6 +1,8 @@
+from __future__ import annotations
+
 import os
 import pathlib
-from typing import List
+from typing import TYPE_CHECKING, List
 
 import boto3
 import botocore
@@ -17,12 +19,16 @@
 from great_expectations.execution_engine import PandasExecutionEngine
 from tests.test_utils import create_files_in_directory
 
+if TYPE_CHECKING:
+    from pyfakefs.fake_filesystem import FakeFilesystem
+
+
 yaml = YAML()
 
 
 @pytest.mark.integration
 @pytest.mark.slow  # 1.05s
-def test__get_full_file_path_for_asset_pandas(fs):
+def test__get_full_file_path_for_asset_pandas(fs: FakeFilesystem):
     """
     What does this test and why?
     File paths in DBFS need to use the `dbfs:/` protocol base instead of `/dbfs/` when
diff --git a/tests/datasource/data_connector/test_inferred_asset_dbfs_data_connector.py b/tests/datasource/data_connector/test_inferred_asset_dbfs_data_connector.py
--- a/tests/datasource/data_connector/test_inferred_asset_dbfs_data_connector.py
+++ b/tests/datasource/data_connector/test_inferred_asset_dbfs_data_connector.py
@@ -1,6 +1,8 @@
+from __future__ import annotations
+
 import os
 import pathlib
-from typing import List
+from typing import TYPE_CHECKING, List
 
 import boto3
 import botocore
@@ -13,9 +15,12 @@
 from great_expectations.execution_engine import PandasExecutionEngine
 from tests.test_utils import create_files_in_directory
 
+if TYPE_CHECKING:
+    from pyfakefs.fake_filesystem import FakeFilesystem
+
 
 @pytest.mark.integration
-def test__get_full_file_path_pandas(fs):
+def test__get_full_file_path_pandas(fs: FakeFilesystem):
     """
     What does this test and why?
     File paths in DBFS need to use the `dbfs:/` protocol base instead of `/dbfs/` when
diff --git a/tests/datasource/fluent/data_asset/data_connector/test_dbfs_data_connector.py b/tests/datasource/fluent/data_asset/data_connector/test_dbfs_data_connector.py
--- a/tests/datasource/fluent/data_asset/data_connector/test_dbfs_data_connector.py
+++ b/tests/datasource/fluent/data_asset/data_connector/test_dbfs_data_connector.py
@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import os
 import pathlib
 import re
@@ -14,6 +16,8 @@
 from tests.test_utils import create_files_in_directory
 
 if TYPE_CHECKING:
+    from pyfakefs.fake_filesystem import FakeFilesystem
+
     from great_expectations.datasource.fluent.data_asset.data_connector import (
         DataConnector,
     )
@@ -21,7 +25,7 @@
 
 @pytest.mark.integration
 @pytest.mark.slow  # 1.05s
-def test__get_full_file_path_pandas(fs):
+def test__get_full_file_path_pandas(fs: FakeFilesystem):
     """
     What does this test and why?
     File paths in DBFS need to use the `dbfs:/` protocol base instead of `/dbfs/` when
diff --git a/tests/datasource/fluent/test_pandas_dbfs_datasource.py b/tests/datasource/fluent/test_pandas_dbfs_datasource.py
--- a/tests/datasource/fluent/test_pandas_dbfs_datasource.py
+++ b/tests/datasource/fluent/test_pandas_dbfs_datasource.py
@@ -4,7 +4,7 @@
 import os
 import pathlib
 import re
-from typing import cast
+from typing import TYPE_CHECKING, cast
 
 import boto3
 import botocore
@@ -25,6 +25,10 @@
 )
 from tests.test_utils import create_files_in_directory
 
+if TYPE_CHECKING:
+    from pyfakefs.fake_filesystem import FakeFilesystem
+
+
 logger = logging.getLogger(__file__)
 
 
@@ -37,7 +41,7 @@
 
 
 @pytest.fixture
-def pandas_dbfs_datasource(fs) -> PandasDBFSDatasource:
+def pandas_dbfs_datasource(fs: FakeFilesystem) -> PandasDBFSDatasource:
     # Copy boto modules into fake filesystem (see https://github.com/spulec/moto/issues/1682#issuecomment-645016188)
     for module in [boto3, botocore]:
         module_dir = pathlib.Path(module.__file__).parent
diff --git a/tests/datasource/fluent/test_spark_dbfs_datasource.py b/tests/datasource/fluent/test_spark_dbfs_datasource.py
new file mode 100644
--- /dev/null
+++ b/tests/datasource/fluent/test_spark_dbfs_datasource.py
@@ -0,0 +1,195 @@
+from __future__ import annotations
+
+import logging
+import os
+import pathlib
+import re
+from typing import TYPE_CHECKING, cast
+
+import boto3
+import botocore
+import pytest
+
+from great_expectations.core.util import DBFSPath
+from great_expectations.datasource.fluent import SparkDBFSDatasource
+from great_expectations.datasource.fluent.data_asset.data_connector import (
+    DBFSDataConnector,
+)
+from great_expectations.datasource.fluent.file_path_data_asset import (
+    _FilePathDataAsset,
+)
+from great_expectations.datasource.fluent.interfaces import TestConnectionError
+from great_expectations.datasource.fluent.spark_file_path_datasource import (
+    CSVAsset,
+)
+from tests.test_utils import create_files_in_directory
+
+if TYPE_CHECKING:
+    from pyfakefs.fake_filesystem import FakeFilesystem
+
+
+logger = logging.getLogger(__file__)
+
+
+@pytest.fixture
+def spark_dbfs_datasource(fs: FakeFilesystem, test_backends) -> SparkDBFSDatasource:
+    if "SparkDFDataset" not in test_backends:
+        pytest.skip("No spark backend selected.")
+
+    # Copy boto modules into fake filesystem (see https://github.com/spulec/moto/issues/1682#issuecomment-645016188)
+    for module in [boto3, botocore]:
+        module_dir = pathlib.Path(module.__file__).parent
+        fs.add_real_directory(module_dir, lazy_read=False)
+
+    # Copy google credentials into fake filesystem if they exist on your filesystem
+    google_cred_file = os.getenv("GOOGLE_APPLICATION_CREDENTIALS")
+    if google_cred_file:
+        fs.add_real_file(google_cred_file)
+
+    base_directory: str = "/dbfs/great_expectations"
+    fs.create_dir(base_directory)
+
+    create_files_in_directory(
+        directory=base_directory,
+        file_name_list=[
+            "alex_20200809_1000.csv",
+            "eugene_20200809_1500.csv",
+            "james_20200811_1009.csv",
+            "abe_20200809_1040.csv",
+            "will_20200809_1002.csv",
+            "james_20200713_1567.csv",
+            "eugene_20201129_1900.csv",
+            "will_20200810_1001.csv",
+            "james_20200810_1003.csv",
+            "alex_20200819_1300.csv",
+        ],
+    )
+
+    return SparkDBFSDatasource(  # type: ignore[call-arg]
+        name="spark_dbfs_datasource",
+        base_directory=pathlib.Path(base_directory),
+    )
+
+
+@pytest.fixture
+def csv_asset(spark_dbfs_datasource: SparkDBFSDatasource) -> _FilePathDataAsset:
+    asset = spark_dbfs_datasource.add_csv_asset(
+        name="csv_asset",
+        batching_regex=r"(?P<name>.+)_(?P<timestamp>.+)_(?P<price>\d{4})\.csv",
+    )
+    return asset
+
+
+@pytest.fixture
+def bad_regex_config(csv_asset: CSVAsset) -> tuple[re.Pattern, str]:
+    regex = re.compile(
+        r"(?P<name>.+)_(?P<ssn>\d{9})_(?P<timestamp>.+)_(?P<price>\d{4})\.csv"
+    )
+    data_connector: DBFSDataConnector = cast(
+        DBFSDataConnector, csv_asset._data_connector
+    )
+    test_connection_error_message = f"""No file at base_directory path "{data_connector._base_directory.resolve()}" matched regular expressions pattern "{data_connector._batching_regex.pattern}" and/or glob_directive "**/*" for DataAsset "csv_asset"."""
+    return regex, test_connection_error_message
+
+
+@pytest.mark.integration
+def test_construct_spark_dbfs_datasource(spark_dbfs_datasource: SparkDBFSDatasource):
+    assert spark_dbfs_datasource.name == "spark_dbfs_datasource"
+
+
+@pytest.mark.integration
+def test_add_csv_asset_to_datasource(spark_dbfs_datasource: SparkDBFSDatasource):
+    asset = spark_dbfs_datasource.add_csv_asset(
+        name="csv_asset",
+        batching_regex=r"(.+)_(.+)_(\d{4})\.csv",
+    )
+    assert asset.name == "csv_asset"
+    assert asset.batching_regex.match("random string") is None
+    assert asset.batching_regex.match("alex_20200819_13D0.csv") is None
+    m1 = asset.batching_regex.match("alex_20200819_1300.csv")
+    assert m1 is not None
+
+
+@pytest.mark.integration
+def test_construct_csv_asset_directly():
+    # noinspection PyTypeChecker
+    asset = CSVAsset(
+        name="csv_asset",
+        batching_regex=r"(.+)_(.+)_(\d{4})\.csv",
+    )
+    assert asset.name == "csv_asset"
+    assert asset.batching_regex.match("random string") is None
+    assert asset.batching_regex.match("alex_20200819_13D0.csv") is None
+    m1 = asset.batching_regex.match("alex_20200819_1300.csv")
+    assert m1 is not None
+
+
+@pytest.mark.integration
+@pytest.mark.xfail(
+    reason="Accessing objects on pyfakefs.fake_filesystem.FakeFilesystem using Spark is not working (this test is conducted using Jupyter notebook manually)."
+)
+def test_get_batch_list_from_fully_specified_batch_request(
+    spark_dbfs_datasource: SparkDBFSDatasource,
+):
+    asset = spark_dbfs_datasource.add_csv_asset(
+        name="csv_asset",
+        batching_regex=r"(?P<name>.+)_(?P<timestamp>.+)_(?P<price>\d{4})\.csv",
+    )
+
+    request = asset.build_batch_request(
+        {"name": "alex", "timestamp": "20200819", "price": "1300"}
+    )
+    batches = asset.get_batch_list_from_batch_request(request)
+    assert len(batches) == 1
+    batch = batches[0]
+    assert batch.batch_request.datasource_name == spark_dbfs_datasource.name
+    assert batch.batch_request.data_asset_name == asset.name
+    assert batch.batch_request.options == (
+        "name",
+        "timestamp",
+        "price",
+        "path",
+    )
+    assert batch.metadata == {
+        "path": "alex_20200819_1300.csv",
+        "name": "alex",
+        "timestamp": "20200819",
+        "price": "1300",
+    }
+    assert (
+        batch.id
+        == "spark_dbfs_datasource-csv_asset-name_alex-timestamp_20200819-price_1300"
+    )
+
+    request = asset.build_batch_request({"name": "alex"})
+    batches = asset.get_batch_list_from_batch_request(request)
+    assert len(batches) == 2
+
+
+@pytest.mark.integration
+def test_test_connection_failures(
+    spark_dbfs_datasource: SparkDBFSDatasource,
+    bad_regex_config: tuple[re.Pattern, str],
+):
+    regex, test_connection_error_message = bad_regex_config
+    csv_asset = CSVAsset(
+        name="csv_asset",
+        batching_regex=regex,
+    )
+    csv_asset._datasource = spark_dbfs_datasource
+    spark_dbfs_datasource.assets = {"csv_asset": csv_asset}
+    csv_asset._data_connector = DBFSDataConnector(
+        datasource_name=spark_dbfs_datasource.name,
+        data_asset_name=csv_asset.name,
+        batching_regex=re.compile(regex),
+        base_directory=spark_dbfs_datasource.base_directory,
+        data_context_root_directory=spark_dbfs_datasource.data_context_root_directory,
+        glob_directive="*.csv",
+        file_path_template_map_fn=DBFSPath.convert_to_file_semantics_version,
+    )
+    csv_asset._test_connection_error_message = test_connection_error_message
+
+    with pytest.raises(TestConnectionError) as e:
+        spark_dbfs_datasource.test_connection()
+
+    assert str(e.value) == str(test_connection_error_message)