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/core/util.py b/great_expectations/core/util.py
--- a/great_expectations/core/util.py
+++ b/great_expectations/core/util.py
@@ -3,6 +3,7 @@
 import decimal
 import logging
 import os
+import re
 import sys
 import uuid
 from collections import OrderedDict
@@ -403,6 +404,32 @@ def datetime_to_int(dt: datetime.date) -> int:
     return int(dt.strftime("%Y%m%d%H%M%S"))
 
 
+class AzureUrl:
+    """
+    Parses an Azure Blob Storage URL into its separate components
+    """
+
+    def __init__(self, url: str):
+        search = re.search(
+            r"^(?:https?://)?(.+?).blob.core.windows.net/([^/]+)/(.+)$", url
+        )
+        self._account_name = search.group(1)
+        self._container = search.group(2)
+        self._blob = search.group(3)
+
+    @property
+    def account_name(self):
+        return self._account_name
+
+    @property
+    def container(self):
+        return self._container
+
+    @property
+    def blob(self):
+        return self._blob
+
+
 # S3Url class courtesy: https://stackoverflow.com/questions/42641315/s3-urls-get-bucket-name-and-path
 class S3Url:
     """
diff --git a/great_expectations/datasource/data_connector/runtime_data_connector.py b/great_expectations/datasource/data_connector/runtime_data_connector.py
--- a/great_expectations/datasource/data_connector/runtime_data_connector.py
+++ b/great_expectations/datasource/data_connector/runtime_data_connector.py
@@ -1,6 +1,5 @@
 import logging
-from typing import Any, List, Optional, Tuple, Union, cast
-from urllib.parse import urlparse
+from typing import Any, List, Optional, Tuple, Union
 
 import great_expectations.exceptions as ge_exceptions
 from great_expectations.core.batch import (
@@ -10,6 +9,7 @@
     RuntimeBatchRequest,
 )
 from great_expectations.core.batch_spec import (
+    AzureBatchSpec,
     BatchMarkers,
     BatchSpec,
     PathBatchSpec,
@@ -224,18 +224,19 @@ def build_batch_spec(
         batch_spec: BatchSpec = super().build_batch_spec(
             batch_definition=batch_definition
         )
-        if runtime_parameters.get("batch_data") is not None:
+        if "batch_data" in runtime_parameters:
             batch_spec["batch_data"] = runtime_parameters.get("batch_data")
             return RuntimeDataBatchSpec(batch_spec)
-        elif runtime_parameters.get("query"):
+        elif "query" in runtime_parameters:
             batch_spec["query"] = runtime_parameters.get("query")
             return RuntimeQueryBatchSpec(batch_spec)
-        elif runtime_parameters.get("path"):
-            path = runtime_parameters.get("path")
+        elif "path" in runtime_parameters:
+            path: str = runtime_parameters["path"]
             batch_spec["path"] = path
-            parsed_url = urlparse(path)
-            if "s3" in parsed_url.scheme:
+            if "s3" in path:
                 return S3BatchSpec(batch_spec)
+            elif "blob.core.windows.net" in path:
+                return AzureBatchSpec(batch_spec)
             else:
                 return PathBatchSpec(batch_spec)
 
diff --git a/great_expectations/execution_engine/pandas_execution_engine.py b/great_expectations/execution_engine/pandas_execution_engine.py
--- a/great_expectations/execution_engine/pandas_execution_engine.py
+++ b/great_expectations/execution_engine/pandas_execution_engine.py
@@ -13,12 +13,13 @@
 import great_expectations.exceptions as ge_exceptions
 from great_expectations.core.batch import BatchMarkers
 from great_expectations.core.batch_spec import (
+    AzureBatchSpec,
     BatchSpec,
     PathBatchSpec,
     RuntimeDataBatchSpec,
     S3BatchSpec,
 )
-from great_expectations.core.util import S3Url, sniff_s3_compression
+from great_expectations.core.util import AzureUrl, S3Url, sniff_s3_compression
 from great_expectations.execution_engine import ExecutionEngine
 from great_expectations.execution_engine.execution_engine import MetricDomainTypes
 from great_expectations.execution_engine.pandas_batch_data import PandasBatchData
@@ -28,6 +29,15 @@
 except ImportError:
     boto3 = None
 
+try:
+    from azure.storage.blob import (
+        BlobClient,
+        BlobServiceClient,
+        StorageStreamDownloader,
+    )
+except:
+    azure = None
+
 
 logger = logging.getLogger(__name__)
 
@@ -78,19 +88,26 @@ def __init__(self, *args, **kwargs):
             "discard_subset_failing_expectations", False
         )
         boto3_options: dict = kwargs.pop("boto3_options", {})
+        azure_options: dict = kwargs.pop("azure_options", {})
 
-        # Try initializing boto3 client. If unsuccessful, we'll catch it when/if a S3BatchSpec is passed in.
+        # Try initializing cloud provider client. If unsuccessful, we'll catch it when/if a BatchSpec is passed in.
         try:
             self._s3 = boto3.client("s3", **boto3_options)
         except (TypeError, AttributeError):
             self._s3 = None
 
+        try:
+            self._azure = BlobServiceClient(**azure_options)
+        except (TypeError, AttributeError):
+            self._azure = None
+
         super().__init__(*args, **kwargs)
 
         self._config.update(
             {
                 "discard_subset_failing_expectations": self.discard_subset_failing_expectations,
                 "boto3_options": boto3_options,
+                "azure_options": azure_options,
             }
         )
 
@@ -139,6 +156,7 @@ def get_batch_data_and_markers(
                     "RuntimeDataBatchSpec must provide a Pandas DataFrame or PandasBatchData object."
                 )
             batch_spec.batch_data = "PandasDataFrame"
+
         elif isinstance(batch_spec, S3BatchSpec):
             if self._s3 is None:
                 raise ge_exceptions.ExecutionEngineError(
@@ -163,15 +181,39 @@ def get_batch_data_and_markers(
             buf = BytesIO(s3_object["Body"].read())
             buf.seek(0)
             df = reader_fn(buf, **reader_options)
+
+        elif isinstance(batch_spec, AzureBatchSpec):
+            if self._azure is None:
+                raise ge_exceptions.ExecutionEngineError(
+                    f"""PandasExecutionEngine has been passed a AzureBatchSpec,
+                        but the ExecutionEngine does not have an Azure client configured. Please check your config."""
+                )
+            azure_engine: BlobServiceClient = self._azure
+            azure_url = AzureUrl(batch_spec.path)
+            reader_method: str = batch_spec.reader_method
+            reader_options: dict = batch_spec.reader_options or {}
+            blob_client: BlobClient = azure_engine.get_blob_client(
+                container=azure_url.container, blob=azure_url.blob
+            )
+            azure_object: StorageStreamDownloader = blob_client.download_blob()
+            logger.debug(
+                f"Fetching Azure blob. Container: {azure_url.container} Blob: {azure_url.blob}"
+            )
+            reader_fn = self._get_reader_fn(reader_method, azure_url.blob)
+            buf = BytesIO(azure_object.readall())
+            buf.seek(0)
+            df = reader_fn(buf, **reader_options)
+
         elif isinstance(batch_spec, PathBatchSpec):
             reader_method: str = batch_spec.reader_method
             reader_options: dict = batch_spec.reader_options
             path: str = batch_spec.path
             reader_fn: Callable = self._get_reader_fn(reader_method, path)
             df = reader_fn(path, **reader_options)
+
         else:
             raise ge_exceptions.BatchSpecError(
-                f"batch_spec must be of type RuntimeDataBatchSpec, PathBatchSpec, or S3BatchSpec, not {batch_spec.__class__.__name__}"
+                f"batch_spec must be of type RuntimeDataBatchSpec, PathBatchSpec, S3BatchSpec, or AzureBatchSpec, not {batch_spec.__class__.__name__}"
             )
 
         df = self._apply_splitting_and_sampling_methods(batch_spec, df)


Test patch:
diff --git a/tests/core/test_util.py b/tests/core/test_util.py
--- a/tests/core/test_util.py
+++ b/tests/core/test_util.py
@@ -2,6 +2,7 @@
 from freezegun import freeze_time
 
 from great_expectations.core.util import (
+    AzureUrl,
     S3Url,
     sniff_s3_compression,
     substitute_all_strftime_format_strings,
@@ -61,3 +62,24 @@ def test_s3_suffix(url, expected):
 )
 def test_sniff_s3_compression(url, expected):
     assert sniff_s3_compression(S3Url(url)) == expected
+
+
+def test_azure_url():
+    url = AzureUrl("my_account.blob.core.windows.net/my_container/my_blob")
+    assert url.account_name == "my_account"
+    assert url.container == "my_container"
+    assert url.blob == "my_blob"
+
+
+def test_azure_url_with_https():
+    url = AzureUrl("https://my_account.blob.core.windows.net/my_container/my_blob")
+    assert url.account_name == "my_account"
+    assert url.container == "my_container"
+    assert url.blob == "my_blob"
+
+
+def test_azure_url_with_nested_blob():
+    url = AzureUrl("my_account.blob.core.windows.net/my_container/a/b/c/d/e/my_blob")
+    assert url.account_name == "my_account"
+    assert url.container == "my_container"
+    assert url.blob == "a/b/c/d/e/my_blob"
diff --git a/tests/data_context/test_data_context_test_yaml_config.py b/tests/data_context/test_data_context_test_yaml_config.py
--- a/tests/data_context/test_data_context_test_yaml_config.py
+++ b/tests/data_context/test_data_context_test_yaml_config.py
@@ -1158,6 +1158,7 @@ def test_golden_path_runtime_data_connector_pandas_datasource_configuration(
         "class_name": "PandasExecutionEngine",
         "discard_subset_failing_expectations": False,
         "boto3_options": {},
+        "azure_options": {},
     }
     assert report_object["data_connectors"]["count"] == 1
 
@@ -1262,6 +1263,7 @@ def test_golden_path_runtime_data_connector_and_inferred_data_connector_pandas_d
         "class_name": "PandasExecutionEngine",
         "discard_subset_failing_expectations": False,
         "boto3_options": {},
+        "azure_options": {},
     }
     assert report_object["data_connectors"]["count"] == 2
     assert report_object["data_connectors"]["default_runtime_data_connector_name"] == {
diff --git a/tests/datasource/test_new_datasource.py b/tests/datasource/test_new_datasource.py
--- a/tests/datasource/test_new_datasource.py
+++ b/tests/datasource/test_new_datasource.py
@@ -203,6 +203,7 @@ def test_basic_pandas_datasource_v013_self_check(basic_pandas_datasource_v013):
             "class_name": "PandasExecutionEngine",
             "discard_subset_failing_expectations": False,
             "boto3_options": {},
+            "azure_options": {},
         },
         "data_connectors": {
             "count": 2,
diff --git a/tests/datasource/test_new_datasource_with_runtime_data_connector.py b/tests/datasource/test_new_datasource_with_runtime_data_connector.py
--- a/tests/datasource/test_new_datasource_with_runtime_data_connector.py
+++ b/tests/datasource/test_new_datasource_with_runtime_data_connector.py
@@ -62,6 +62,7 @@ def test_basic_datasource_runtime_data_connector_self_check(
         },
         "execution_engine": {
             "boto3_options": {},
+            "azure_options": {},
             "caching": True,
             "class_name": "PandasExecutionEngine",
             "discard_subset_failing_expectations": False,
diff --git a/tests/execution_engine/test_pandas_execution_engine.py b/tests/execution_engine/test_pandas_execution_engine.py
--- a/tests/execution_engine/test_pandas_execution_engine.py
+++ b/tests/execution_engine/test_pandas_execution_engine.py
@@ -3,6 +3,7 @@
 import random
 from pathlib import Path
 from typing import List
+from unittest import mock
 
 import boto3
 import pandas as pd
@@ -10,9 +11,16 @@
 from botocore.errorfactory import ClientError
 from moto import mock_s3
 
+try:
+    from azure.storage.blob import BlobServiceClient
+except:
+    azure = None
+
+
 import great_expectations.exceptions as ge_exceptions
 from great_expectations.core.batch import BatchDefinition
 from great_expectations.core.batch_spec import (
+    AzureBatchSpec,
     PathBatchSpec,
     RuntimeDataBatchSpec,
     S3BatchSpec,
@@ -27,7 +35,138 @@
 from tests.expectations.test_util import get_table_columns_metric
 
 
-def test_constructor():
+@pytest.fixture(scope="function")
+def aws_credentials():
+    """Mocked AWS Credentials for moto."""
+    os.environ["AWS_ACCESS_KEY_ID"] = "testing"
+    os.environ["AWS_SECRET_ACCESS_KEY"] = "testing"
+    os.environ["AWS_SECURITY_TOKEN"] = "testing"
+    os.environ["AWS_SESSION_TOKEN"] = "testing"
+
+
+@pytest.fixture
+def s3(aws_credentials):
+    with mock_s3():
+        yield boto3.client("s3", region_name="us-east-1")
+
+
+@pytest.fixture
+def s3_bucket(s3):
+    bucket: str = "test_bucket"
+    s3.create_bucket(Bucket=bucket)
+    return bucket
+
+
+@pytest.fixture
+def test_df_small() -> pd.DataFrame:
+    return pd.DataFrame(data={"col1": [1, 0, 505], "col2": [3, 4, 101]})
+
+
+@pytest.fixture
+def test_df_small_csv_compressed(test_df_small, tmpdir) -> bytes:
+    path = Path(tmpdir) / "file.csv.gz"
+    test_df_small.to_csv(path, index=False, compression="gzip")
+    return path.read_bytes()
+
+
+@pytest.fixture
+def test_df_small_csv(test_df_small, tmpdir) -> bytes:
+    path = Path(tmpdir) / "file.csv"
+    test_df_small.to_csv(path, index=False)
+    return path.read_bytes()
+
+
+@pytest.fixture
+def test_s3_files(s3, s3_bucket, test_df_small_csv):
+    keys: List[str] = [
+        "path/A-100.csv",
+        "path/A-101.csv",
+        "directory/B-1.csv",
+        "directory/B-2.csv",
+        "alpha-1.csv",
+        "alpha-2.csv",
+    ]
+    for key in keys:
+        s3.put_object(Bucket=s3_bucket, Body=test_df_small_csv, Key=key)
+    return s3_bucket, keys
+
+
+@pytest.fixture
+def test_s3_files_parquet(tmpdir, s3, s3_bucket, test_df_small, test_df_small_csv):
+    keys: List[str] = [
+        "path/A-100.csv",
+        "path/A-101.csv",
+        "directory/B-1.parquet",
+        "directory/B-2.parquet",
+        "alpha-1.csv",
+        "alpha-2.csv",
+    ]
+    path = Path(tmpdir) / "file.parquet"
+    test_df_small.to_parquet(path)
+    for key in keys:
+        if key.endswith(".parquet"):
+            with open(path, "rb") as f:
+                s3.put_object(Bucket=s3_bucket, Body=f, Key=key)
+        else:
+            s3.put_object(Bucket=s3_bucket, Body=test_df_small_csv, Key=key)
+    return s3_bucket, keys
+
+
+@pytest.fixture
+def batch_with_split_on_whole_table_s3(test_s3_files) -> S3BatchSpec:
+    bucket, keys = test_s3_files
+    path = keys[0]
+    full_path = f"s3a://{os.path.join(bucket, path)}"
+
+    batch_spec = S3BatchSpec(
+        path=full_path,
+        reader_method="read_csv",
+        splitter_method="_split_on_whole_table",
+    )
+    return batch_spec
+
+
+@pytest.fixture
+def test_s3_files_compressed(s3, s3_bucket, test_df_small_csv_compressed):
+    keys: List[str] = [
+        "path/A-100.csv.gz",
+        "path/A-101.csv.gz",
+        "directory/B-1.csv.gz",
+        "directory/B-2.csv.gz",
+    ]
+
+    for key in keys:
+        s3.put_object(
+            Bucket=s3_bucket,
+            Body=test_df_small_csv_compressed,
+            Key=key,
+        )
+    return s3_bucket, keys
+
+
+@pytest.fixture
+def azure_batch_spec() -> AzureBatchSpec:
+    container = "test_container"
+    keys: List[str] = [
+        "path/A-100.csv",
+        "path/A-101.csv",
+        "directory/B-1.csv",
+        "directory/B-2.csv",
+        "alpha-1.csv",
+        "alpha-2.csv",
+    ]
+    path = keys[0]
+    full_path = os.path.join("mock_account.blob.core.windows.net", container, path)
+
+    batch_spec = AzureBatchSpec(
+        path=full_path,
+        reader_method="read_csv",
+        splitter_method="_split_on_whole_table",
+    )
+    return batch_spec
+
+
+def test_constructor_with_boto3_options():
     # default instantiation
     PandasExecutionEngine()
 
@@ -427,97 +566,6 @@ def test_get_batch_with_split_on_whole_table_filesystem(
     assert test_df.dataframe.shape == (5, 2)
 
 
-@pytest.fixture(scope="function")
-def aws_credentials():
-    """Mocked AWS Credentials for moto."""
-    os.environ["AWS_ACCESS_KEY_ID"] = "testing"
-    os.environ["AWS_SECRET_ACCESS_KEY"] = "testing"
-    os.environ["AWS_SECURITY_TOKEN"] = "testing"
-    os.environ["AWS_SESSION_TOKEN"] = "testing"
-
-
-@pytest.fixture
-def s3(aws_credentials):
-    with mock_s3():
-        yield boto3.client("s3", region_name="us-east-1")
-
-
-@pytest.fixture
-def s3_bucket(s3):
-    bucket: str = "test_bucket"
-    s3.create_bucket(Bucket=bucket)
-    return bucket
-
-
-@pytest.fixture
-def test_df_small() -> pd.DataFrame:
-    return pd.DataFrame(data={"col1": [1, 0, 505], "col2": [3, 4, 101]})
-
-
-@pytest.fixture
-def test_df_small_csv_compressed(test_df_small, tmpdir) -> bytes:
-    path = Path(tmpdir) / "file.csv.gz"
-    test_df_small.to_csv(path, index=False, compression="gzip")
-    return path.read_bytes()
-
-
-@pytest.fixture
-def test_df_small_csv(test_df_small, tmpdir) -> bytes:
-    path = Path(tmpdir) / "file.csv"
-    test_df_small.to_csv(path, index=False)
-    return path.read_bytes()
-
-
-@pytest.fixture
-def test_s3_files(s3, s3_bucket, test_df_small_csv):
-    keys: List[str] = [
-        "path/A-100.csv",
-        "path/A-101.csv",
-        "directory/B-1.csv",
-        "directory/B-2.csv",
-        "alpha-1.csv",
-        "alpha-2.csv",
-    ]
-    for key in keys:
-        s3.put_object(Bucket=s3_bucket, Body=test_df_small_csv, Key=key)
-    return s3_bucket, keys
-
-
-@pytest.fixture
-def test_s3_files_parquet(tmpdir, s3, s3_bucket, test_df_small, test_df_small_csv):
-    keys: List[str] = [
-        "path/A-100.csv",
-        "path/A-101.csv",
-        "directory/B-1.parquet",
-        "directory/B-2.parquet",
-        "alpha-1.csv",
-        "alpha-2.csv",
-    ]
-    path = Path(tmpdir) / "file.parquet"
-    test_df_small.to_parquet(path)
-    for key in keys:
-        if key.endswith(".parquet"):
-            with open(path, "rb") as f:
-                s3.put_object(Bucket=s3_bucket, Body=f, Key=key)
-        else:
-            s3.put_object(Bucket=s3_bucket, Body=test_df_small_csv, Key=key)
-    return s3_bucket, keys
-
-
-@pytest.fixture
-def batch_with_split_on_whole_table_s3(test_s3_files) -> S3BatchSpec:
-    bucket, keys = test_s3_files
-    path = keys[0]
-    full_path = f"s3a://{os.path.join(bucket, path)}"
-
-    batch_spec = S3BatchSpec(
-        path=full_path,
-        reader_method="read_csv",
-        splitter_method="_split_on_whole_table",
-    )
-    return batch_spec
-
-
 def test_get_batch_with_split_on_whole_table_s3(
     batch_with_split_on_whole_table_s3, test_df_small
 ):
@@ -588,24 +636,6 @@ def test_get_batch_with_split_on_whole_table_s3_with_configured_asset_s3_data_co
         )
 
 
-@pytest.fixture
-def test_s3_files_compressed(s3, s3_bucket, test_df_small_csv_compressed):
-    keys: List[str] = [
-        "path/A-100.csv.gz",
-        "path/A-101.csv.gz",
-        "directory/B-1.csv.gz",
-        "directory/B-2.csv.gz",
-    ]
-
-    for key in keys:
-        s3.put_object(
-            Bucket=s3_bucket,
-            Body=test_df_small_csv_compressed,
-            Key=key,
-        )
-    return s3_bucket, keys
-
-
 def test_get_batch_s3_compressed_files(test_s3_files_compressed, test_df_small):
     bucket, keys = test_s3_files_compressed
     path = keys[0]
@@ -864,3 +894,53 @@ def test_get_batch_with_split_on_divided_integer_and_sample_on_list(test_df):
     assert split_df.dataframe.shape == (2, 10)
     assert split_df.dataframe.id.min() == 54
     assert split_df.dataframe.id.max() == 59
+
+
+@mock.patch(
+    "great_expectations.datasource.data_connector.configured_asset_azure_data_connector.BlobServiceClient"
+)
+def test_constructor_with_azure_options(mock_azure_conn):
+    # default instantiation
+    PandasExecutionEngine()
+
+    # instantiation with custom parameters
+    engine = PandasExecutionEngine(discard_subset_failing_expectations=True)
+    assert "discard_subset_failing_expectations" in engine.config
+    assert engine.config.get("discard_subset_failing_expectations") is True
+    custom_azure_options = {"account_url": "my_account_url"}
+    engine = PandasExecutionEngine(azure_options=custom_azure_options)
+    assert "azure_options" in engine.config
+    assert engine.config.get("azure_options")["account_url"] == "my_account_url"
+
+
+@mock.patch(
+    "great_expectations.execution_engine.pandas_execution_engine.BlobServiceClient",
+)
+def test_get_batch_data_with_azure_batch_spec(
+    mock_azure_conn,
+    azure_batch_spec,
+):
+    mock_blob_client = mock_azure_conn().get_blob_client()
+    mock_azure_obj = mock_blob_client.download_blob()
+    mock_azure_obj.readall.return_value = (
+        b"colA,colB,colC\n1,2,3\n4,5,6\n7,8,9"  # (3,3) CSV for testing
+    )
+
+    df = PandasExecutionEngine().get_batch_data(batch_spec=azure_batch_spec)
+
+    mock_azure_conn().get_blob_client.assert_called_with(
+        container="test_container", blob="path/A-100.csv"
+    )
+    mock_azure_obj.readall.assert_called_once()
+
+    assert df.dataframe.shape == (3, 3)
+
+
+def test_get_batch_with_no_azure_configured(azure_batch_spec):
+    # if Azure BlobServiceClient was not configured
+    execution_engine_no_azure = PandasExecutionEngine()
+    execution_engine_no_azure._azure = None
+
+    # Raises error due the connection object not being set
+    with pytest.raises(ge_exceptions.ExecutionEngineError):
+        execution_engine_no_azure.get_batch_data(batch_spec=azure_batch_spec)