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/cli/datasource.py b/great_expectations/cli/datasource.py
--- a/great_expectations/cli/datasource.py
+++ b/great_expectations/cli/datasource.py
@@ -47,6 +47,7 @@ class SupportedDatabaseBackends(enum.Enum):
     SNOWFLAKE = "Snowflake"
     BIGQUERY = "BigQuery"
     TRINO = "Trino"
+    ATHENA = "Athena"
     OTHER = "other - Do you have a working SQLAlchemy connection string?"
     # TODO MSSQL
 
@@ -746,6 +747,43 @@ def credentials_snippet(self) -> str:
 table_name = "{self.table_name}"'''
 
 
+class AthenaCredentialYamlHelper(SQLCredentialYamlHelper):
+    def __init__(self, datasource_name: Optional[str]) -> None:
+        super().__init__(
+            datasource_name=datasource_name,
+            usage_stats_payload={
+                "type": "sqlalchemy",
+                "db": SupportedDatabaseBackends.ATHENA.value,
+                "api_version": "v3",
+            },
+        )
+
+    def verify_libraries_installed(self) -> bool:
+        # noinspection SpellCheckingInspection
+        pyathena_success: bool = verify_library_dependent_modules(
+            python_import_name="pyathena",
+            pip_library_name="pyathena[SQLAlchemy]",
+            module_names_to_reload=CLI_ONLY_SQLALCHEMY_ORDERED_DEPENDENCY_MODULE_NAMES,
+        )
+        return pyathena_success
+
+    def credentials_snippet(self) -> str:
+        return '''\
+# The SQLAlchemy url/connection string for the Athena connection
+# (reference: https://docs.greatexpectations.io/docs/guides/connecting_to_your_data/database/athena or https://github.com/laughingman7743/PyAthena/#sqlalchemy)"""
+
+schema_name = "YOUR_SCHEMA"  # or database name. It is optional
+table_name = "YOUR_TABLE_NAME"
+region = "YOUR_REGION"
+s3_path = "s3://YOUR_S3_BUCKET/path/to/"  # ignore partitioning
+
+connection_string = f"awsathena+rest://@athena.{region}.amazonaws.com/{schema_name}?s3_staging_dir={s3_path}"
+            '''
+
+    def _yaml_innards(self) -> str:
+        return "\n  connection_string: {connection_string}"
+
+
 class ConnectionStringCredentialYamlHelper(SQLCredentialYamlHelper):
     def __init__(self, datasource_name: Optional[str]) -> None:
         super().__init__(
@@ -785,6 +823,7 @@ def _yaml_innards(self) -> str:
     BigqueryCredentialYamlHelper,
     ConnectionStringCredentialYamlHelper,
     TrinoCredentialYamlHelper,
+    AthenaCredentialYamlHelper,
 ]
 
 
@@ -798,6 +837,7 @@ def _get_sql_yaml_helper_class(
         SupportedDatabaseBackends.SNOWFLAKE: SnowflakeCredentialYamlHelper,
         SupportedDatabaseBackends.BIGQUERY: BigqueryCredentialYamlHelper,
         SupportedDatabaseBackends.TRINO: TrinoCredentialYamlHelper,
+        SupportedDatabaseBackends.ATHENA: AthenaCredentialYamlHelper,
         SupportedDatabaseBackends.OTHER: ConnectionStringCredentialYamlHelper,
     }
     helper_class = helper_class_by_backend[selected_database]


Test patch:
diff --git a/tests/cli/test_datasource_athena.py b/tests/cli/test_datasource_athena.py
new file mode 100644
--- /dev/null
+++ b/tests/cli/test_datasource_athena.py
@@ -0,0 +1,118 @@
+import os
+from unittest import mock
+
+import nbformat
+import pytest
+from click.testing import CliRunner
+from nbconvert.preprocessors import ExecutePreprocessor
+
+from great_expectations.cli import cli
+from great_expectations.util import get_context
+from tests.cli.utils import assert_no_logging_messages_or_tracebacks
+
+
+@mock.patch(
+    "great_expectations.core.usage_statistics.usage_statistics.UsageStatisticsHandler.emit"
+)
+@mock.patch("subprocess.call", return_value=True, side_effect=None)
+@pytest.mark.slow  # 1.85s
+def test_cli_athena_datasource_new_connection_string(
+    mock_subprocess,
+    mock_emit,
+    empty_data_context_stats_enabled,
+    caplog,
+    monkeypatch,
+):
+    root_dir = empty_data_context_stats_enabled.root_directory
+    context = empty_data_context_stats_enabled
+    assert context.list_datasources() == []
+
+    runner = CliRunner(mix_stderr=False)
+    monkeypatch.chdir(os.path.dirname(context.root_directory))
+    result = runner.invoke(
+        cli,
+        "--v3-api datasource new",
+        input="2\n7\n",
+        catch_exceptions=False,
+    )
+    stdout = result.stdout
+
+    assert "What data would you like Great Expectations to connect to?" in stdout
+
+    assert "Athena" in stdout
+
+    assert result.exit_code == 0
+
+    uncommitted_dir = os.path.join(root_dir, context.GX_UNCOMMITTED_DIR)
+    expected_notebook = os.path.join(uncommitted_dir, "datasource_new.ipynb")
+
+    assert os.path.isfile(expected_notebook)
+    mock_subprocess.assert_called_once_with(["jupyter", "notebook", expected_notebook])
+
+    expected_call_args_list = [
+        mock.call(
+            {"event_payload": {}, "event": "data_context.__init__", "success": True}
+        ),
+        mock.call(
+            {
+                "event": "cli.datasource.new.begin",
+                "event_payload": {"api_version": "v3"},
+                "success": True,
+            }
+        ),
+        mock.call(
+            {
+                "event": "cli.new_ds_choice",
+                "event_payload": {
+                    "type": "sqlalchemy",
+                    "db": "Athena",
+                    "api_version": "v3",
+                },
+                "success": True,
+            }
+        ),
+        mock.call(
+            {
+                "event": "cli.datasource.new.end",
+                "event_payload": {"api_version": "v3"},
+                "success": True,
+            }
+        ),
+    ]
+
+    assert mock_emit.call_args_list == expected_call_args_list
+    assert mock_emit.call_count == len(expected_call_args_list)
+
+    # Run notebook
+    with open(expected_notebook) as f:
+        nb = nbformat.read(f, as_version=4)
+
+    # Mock the user adding a connection string into the notebook by overwriting the right cell
+    credentials_cell = nb["cells"][5]["source"]
+
+    credentials = (
+        "connection_string",
+        "schema_name",
+        "table_name",
+        "region",
+        "s3_path",
+    )
+    for credential in credentials:
+        assert credential in credentials_cell
+
+    region_test = "us-east-1"
+    schema_test = "test_data"
+    table_name_test = "table_test"
+    connection_string = f"awsathena+rest://@athena.{region_test}.amazonaws.com/{schema_test}?s3_staging_dir=s3://YOUR_S3_BUCKET/path/to/"
+
+    nb["cells"][5]["source"] = credentials_cell.replace("YOUR_REGION", region_test)
+
+    credentials_cell = nb["cells"][5]["source"]
+    nb["cells"][5]["source"] = credentials_cell.replace("YOUR_SCHEMA", schema_test)
+
+    credentials_cell = nb["cells"][5]["source"]
+    nb["cells"][5]["source"] = credentials_cell.replace(
+        "YOUR_TABLE_NAME", table_name_test
+    )
+
+    assert_no_logging_messages_or_tracebacks(caplog, result)
diff --git a/tests/cli/test_datasource_new_helpers.py b/tests/cli/test_datasource_new_helpers.py
--- a/tests/cli/test_datasource_new_helpers.py
+++ b/tests/cli/test_datasource_new_helpers.py
@@ -4,6 +4,7 @@
 
 from great_expectations import DataContext
 from great_expectations.cli.datasource import (
+    AthenaCredentialYamlHelper,
     BigqueryCredentialYamlHelper,
     ConnectionStringCredentialYamlHelper,
     MySQLCredentialYamlHelper,
@@ -791,3 +792,67 @@ def test_check_if_datasource_name_exists(
         )
         is False
     )
+
+
+@mock.patch(
+    "great_expectations.core.usage_statistics.usage_statistics.UsageStatisticsHandler.emit"
+)
+def test_AthenaCredentialYamlHelper(mock_emit, empty_data_context_stats_enabled):
+    helper = AthenaCredentialYamlHelper("my_datasource")
+    assert (
+        helper.credentials_snippet()
+        == '''\
+# The SQLAlchemy url/connection string for the Athena connection
+# (reference: https://docs.greatexpectations.io/docs/guides/connecting_to_your_data/database/athena or https://github.com/laughingman7743/PyAthena/#sqlalchemy)"""
+
+schema_name = "YOUR_SCHEMA"  # or database name. It is optional
+table_name = "YOUR_TABLE_NAME"
+region = "YOUR_REGION"
+s3_path = "s3://YOUR_S3_BUCKET/path/to/"  # ignore partitioning
+
+connection_string = f"awsathena+rest://@athena.{region}.amazonaws.com/{schema_name}?s3_staging_dir={s3_path}"
+            '''
+    )
+
+    assert (
+        helper.yaml_snippet()
+        == '''f"""
+name: {datasource_name}
+class_name: Datasource
+execution_engine:
+  class_name: SqlAlchemyExecutionEngine
+  connection_string: {connection_string}
+data_connectors:
+  default_runtime_data_connector_name:
+    class_name: RuntimeDataConnector
+    batch_identifiers:
+      - default_identifier_name
+  default_inferred_data_connector_name:
+    class_name: InferredAssetSqlDataConnector
+    include_schema_name: True
+    introspection_directives:
+      schema_name: {schema_name}
+  default_configured_data_connector_name:
+    class_name: ConfiguredAssetSqlDataConnector
+    assets:
+      {table_name}:
+        class_name: Asset
+        schema_name: {schema_name}
+"""'''
+    )
+
+    helper.send_backend_choice_usage_message(empty_data_context_stats_enabled)
+    assert mock_emit.call_count == 1
+    assert mock_emit.call_args_list == [
+        mock.call(
+            {
+                "event": "cli.new_ds_choice",
+                "event_payload": {
+                    "type": "sqlalchemy",
+                    "db": "Athena",
+                    "api_version": "v3",
+                },
+                "success": True,
+            }
+        ),
+    ]
diff --git a/tests/cli/test_datasource_sqlite.py b/tests/cli/test_datasource_sqlite.py
--- a/tests/cli/test_datasource_sqlite.py
+++ b/tests/cli/test_datasource_sqlite.py
@@ -237,7 +237,7 @@ def test_cli_datasource_new_connection_string(
     result = runner.invoke(
         cli,
         "--v3-api datasource new",
-        input="2\n7\n",
+        input="2\n8\n",
         catch_exceptions=False,
     )
     stdout = result.stdout