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/execution_engine/split_and_sample/sqlalchemy_data_splitter.py b/great_expectations/execution_engine/split_and_sample/sqlalchemy_data_splitter.py
--- a/great_expectations/execution_engine/split_and_sample/sqlalchemy_data_splitter.py
+++ b/great_expectations/execution_engine/split_and_sample/sqlalchemy_data_splitter.py
@@ -28,6 +28,7 @@
     sa = None
 
 try:
+    import sqlalchemy.sql.functions.concat as concat
     from sqlalchemy.engine import LegacyRow
     from sqlalchemy.sql import Selectable
     from sqlalchemy.sql.elements import BinaryExpression, BooleanClauseList, Label
@@ -37,6 +38,7 @@
     BinaryExpression = None
     BooleanClauseList = None
     Label = None
+    concat = None
 
 
 class SqlAlchemyDataSplitter(DataSplitter):
@@ -410,22 +412,30 @@ def get_split_query_for_data_for_batch_identifiers_for_split_on_date_parts(
                 ).label("concat_distinct_values")
             ]
         else:
+            """
+            # NOTE: <Alex>6/29/2022</Alex>
+            Certain SQLAlchemy-compliant backends (e.g., Amazon Redshift) allow only binary operators for "CONCAT".
+            """
+            concat_date_parts: concat = sa.func.concat(
+                "",
+                sa.cast(
+                    sa.func.extract(date_parts[0].value, sa.column(column_name)),
+                    sa.String,
+                ),
+            )
+
+            date_part: DatePart
+            for date_part in date_parts[1:]:
+                concat_date_parts = sa.func.concat(
+                    concat_date_parts,
+                    sa.cast(
+                        sa.func.extract(date_part.value, sa.column(column_name)),
+                        sa.String,
+                    ),
+                )
+
             concat_clause: List[Label] = [
-                sa.func.distinct(
-                    sa.func.concat(
-                        *[
-                            (
-                                sa.cast(
-                                    sa.func.extract(
-                                        date_part.value, sa.column(column_name)
-                                    ),
-                                    sa.String,
-                                )
-                            ).label(date_part.value)
-                            for date_part in date_parts
-                        ]
-                    )
-                ).label("concat_distinct_values"),
+                sa.func.distinct(concat_date_parts).label("concat_distinct_values"),
             ]
 
         split_query: Selectable = sa.select(
@@ -479,7 +489,7 @@ def get_data_for_batch_identifiers_for_split_on_date_parts(
 
     def _execute_split_query(
         self,
-        execution_engine: "SqlAlchemyExecutionEngine",
+        execution_engine: "SqlAlchemyExecutionEngine",  # noqa: F821
         split_query: Selectable,  # noqa: F821
     ) -> List[LegacyRow]:
         """Use the provided execution engine to run the split query and fetch all of the results.


Test patch:
diff --git a/tests/execution_engine/split_and_sample/test_sqlalchemy_execution_engine_splitting.py b/tests/execution_engine/split_and_sample/test_sqlalchemy_execution_engine_splitting.py
--- a/tests/execution_engine/split_and_sample/test_sqlalchemy_execution_engine_splitting.py
+++ b/tests/execution_engine/split_and_sample/test_sqlalchemy_execution_engine_splitting.py
@@ -328,25 +328,16 @@ def test_get_split_query_for_data_for_batch_identifiers_for_split_on_date_parts_
 
     assert isinstance(result, sa.sql.Select)
 
-    query_str: str = (
+    expected_query_str: str = "SELECT distinct(concat(concat('', CAST(EXTRACT(year FROM column_name) AS VARCHAR)), CAST(EXTRACT(month FROM column_name) AS VARCHAR))) AS concat_distinct_values, CAST(EXTRACT(year FROM column_name) AS INTEGER) AS year, CAST(EXTRACT(month FROM column_name) AS INTEGER) AS month FROM table_name"
+    actual_query_str: str = (
         str(result.compile(compile_kwargs={"literal_binds": True}))
         .replace("\n", "")
         .replace(" ", "")
         .lower()
     )
     assert (
-        query_str
-        == (
-            "SELECT distinct(concat("
-            "CAST(EXTRACT(year FROM column_name) AS VARCHAR), CAST(EXTRACT(month FROM column_name) AS VARCHAR)"
-            ")) AS concat_distinct_values,"
-            "CAST(EXTRACT(year FROM column_name) AS INTEGER) AS year,"
-            "CAST(EXTRACT(month FROM column_name) AS INTEGER) AS month"
-            "FROM table_name"
-        )
-        .replace("\n", "")
-        .replace(" ", "")
-        .lower()
+        actual_query_str
+        == expected_query_str.replace("\n", "").replace(" ", "").lower()
     )
 
 
diff --git a/tests/integration/db/test_sql_data_splitting.py b/tests/integration/db/test_sql_data_splitting.py
--- a/tests/integration/db/test_sql_data_splitting.py
+++ b/tests/integration/db/test_sql_data_splitting.py
@@ -131,6 +131,7 @@ def _is_dialect_athena(dialect: str) -> bool:
             },
         }
 
+        # noinspection PyTypeChecker
         context.add_datasource(
             name=datasource_name,
             class_name="Datasource",
diff --git a/tests/integration/fixtures/split_and_sample_data/redshift_connection_string.yml b/tests/integration/fixtures/split_and_sample_data/redshift_connection_string.yml
new file mode 100644
--- /dev/null
+++ b/tests/integration/fixtures/split_and_sample_data/redshift_connection_string.yml
@@ -0,0 +1,2 @@
+connection_string: unsupported_see_get_redshift_connection_url
+dialect: redshift
diff --git a/tests/integration/fixtures/split_and_sample_data/splitter_test_cases_and_fixtures.py b/tests/integration/fixtures/split_and_sample_data/splitter_test_cases_and_fixtures.py
--- a/tests/integration/fixtures/split_and_sample_data/splitter_test_cases_and_fixtures.py
+++ b/tests/integration/fixtures/split_and_sample_data/splitter_test_cases_and_fixtures.py
@@ -20,7 +20,8 @@ def test_df(self):
     def test_column_name(self):
         return self._test_column_name
 
-    def years_in_taxi_data(self) -> List[datetime.datetime]:
+    @staticmethod
+    def years_in_taxi_data() -> List[datetime.datetime]:
         return (
             pd.date_range(start="2018-01-01", end="2020-12-31", freq="AS")
             .to_pydatetime()
@@ -30,7 +31,8 @@ def years_in_taxi_data(self) -> List[datetime.datetime]:
     def year_batch_identifier_data(self) -> List[dict]:
         return [{DatePart.YEAR.value: dt.year} for dt in self.years_in_taxi_data()]
 
-    def months_in_taxi_data(self) -> List[datetime.datetime]:
+    @staticmethod
+    def months_in_taxi_data() -> List[datetime.datetime]:
         return (
             pd.date_range(start="2018-01-01", end="2020-12-31", freq="MS")
             .to_pydatetime()
diff --git a/tests/integration/test_script_runner.py b/tests/integration/test_script_runner.py
--- a/tests/integration/test_script_runner.py
+++ b/tests/integration/test_script_runner.py
@@ -365,6 +365,20 @@ class IntegrationTestFixture:
         ),
         extra_backend_dependencies=BackendDependencies.POSTGRESQL,
     ),
+    IntegrationTestFixture(
+        name="split_data_on_datetime_redshift",
+        user_flow_script="tests/integration/db/test_sql_data_splitting.py",
+        data_context_dir="tests/integration/fixtures/no_datasources/great_expectations",
+        data_dir="tests/test_sets/taxi_yellow_tripdata_samples/",
+        util_script="tests/test_utils.py",
+        other_files=(
+            (
+                "tests/integration/fixtures/split_and_sample_data/redshift_connection_string.yml",
+                "connection_string.yml",
+            ),
+        ),
+        extra_backend_dependencies=BackendDependencies.POSTGRESQL,
+    ),
     IntegrationTestFixture(
         name="split_data_on_datetime_mssql",
         user_flow_script="tests/integration/db/test_sql_data_splitting.py",
@@ -393,6 +407,48 @@ class IntegrationTestFixture:
         ),
         extra_backend_dependencies=BackendDependencies.MYSQL,
     ),
+    IntegrationTestFixture(
+        name="split_data_on_datetime_snowflake",
+        user_flow_script="tests/integration/db/test_sql_data_splitting.py",
+        data_context_dir="tests/integration/fixtures/no_datasources/great_expectations",
+        data_dir="tests/test_sets/taxi_yellow_tripdata_samples/",
+        util_script="tests/test_utils.py",
+        other_files=(
+            (
+                "tests/integration/fixtures/split_and_sample_data/snowflake_connection_string.yml",
+                "connection_string.yml",
+            ),
+        ),
+        extra_backend_dependencies=BackendDependencies.SNOWFLAKE,
+    ),
+    IntegrationTestFixture(
+        name="split_data_on_datetime_bigquery",
+        user_flow_script="tests/integration/db/test_sql_data_splitting.py",
+        data_context_dir="tests/integration/fixtures/no_datasources/great_expectations",
+        data_dir="tests/test_sets/taxi_yellow_tripdata_samples/",
+        util_script="tests/test_utils.py",
+        other_files=(
+            (
+                "tests/integration/fixtures/split_and_sample_data/bigquery_connection_string.yml",
+                "connection_string.yml",
+            ),
+        ),
+        extra_backend_dependencies=BackendDependencies.BIGQUERY,
+    ),
+    IntegrationTestFixture(
+        name="split_data_on_datetime_awsathena",
+        user_flow_script="tests/integration/db/test_sql_data_splitting.py",
+        data_context_dir="tests/integration/fixtures/no_datasources/great_expectations",
+        data_dir="tests/test_sets/taxi_yellow_tripdata_samples/",
+        util_script="tests/test_utils.py",
+        other_files=(
+            (
+                "tests/integration/fixtures/split_and_sample_data/awsathena_connection_string.yml",
+                "connection_string.yml",
+            ),
+        ),
+        extra_backend_dependencies=BackendDependencies.AWS,
+    ),
     IntegrationTestFixture(
         name="sample_data_using_limit_postgres",
         user_flow_script="tests/integration/db/test_sql_data_sampling.py",
@@ -510,20 +566,6 @@ class IntegrationTestFixture:
         extra_backend_dependencies=BackendDependencies.SNOWFLAKE,
         util_script="tests/test_utils.py",
     ),
-    IntegrationTestFixture(
-        name="split_data_on_datetime_snowflake",
-        user_flow_script="tests/integration/db/test_sql_data_splitting.py",
-        data_context_dir="tests/integration/fixtures/no_datasources/great_expectations",
-        data_dir="tests/test_sets/taxi_yellow_tripdata_samples/",
-        util_script="tests/test_utils.py",
-        other_files=(
-            (
-                "tests/integration/fixtures/split_and_sample_data/snowflake_connection_string.yml",
-                "connection_string.yml",
-            ),
-        ),
-        extra_backend_dependencies=BackendDependencies.SNOWFLAKE,
-    ),
     IntegrationTestFixture(
         name="sample_data_using_limit_snowflake",
         user_flow_script="tests/integration/db/test_sql_data_sampling.py",
@@ -645,20 +687,6 @@ class IntegrationTestFixture:
         data_context_dir="tests/integration/fixtures/no_datasources/great_expectations",
         extra_backend_dependencies=BackendDependencies.BIGQUERY,
     ),
-    IntegrationTestFixture(
-        name="split_data_on_datetime_bigquery",
-        user_flow_script="tests/integration/db/test_sql_data_splitting.py",
-        data_context_dir="tests/integration/fixtures/no_datasources/great_expectations",
-        data_dir="tests/test_sets/taxi_yellow_tripdata_samples/",
-        util_script="tests/test_utils.py",
-        other_files=(
-            (
-                "tests/integration/fixtures/split_and_sample_data/bigquery_connection_string.yml",
-                "connection_string.yml",
-            ),
-        ),
-        extra_backend_dependencies=BackendDependencies.BIGQUERY,
-    ),
     IntegrationTestFixture(
         name="sample_data_using_limit_bigquery",
         user_flow_script="tests/integration/db/test_sql_data_sampling.py",
@@ -878,20 +906,6 @@ class IntegrationTestFixture:
         extra_backend_dependencies=BackendDependencies.AWS,
         util_script="tests/test_utils.py",
     ),
-    IntegrationTestFixture(
-        name="split_data_on_datetime_awsathena",
-        user_flow_script="tests/integration/db/test_sql_data_splitting.py",
-        data_context_dir="tests/integration/fixtures/no_datasources/great_expectations",
-        data_dir="tests/test_sets/taxi_yellow_tripdata_samples/",
-        util_script="tests/test_utils.py",
-        other_files=(
-            (
-                "tests/integration/fixtures/split_and_sample_data/awsathena_connection_string.yml",
-                "connection_string.yml",
-            ),
-        ),
-        extra_backend_dependencies=BackendDependencies.AWS,
-    ),
     IntegrationTestFixture(
         name="sample_data_using_limit_awsathena",
         user_flow_script="tests/integration/db/test_sql_data_sampling.py",
@@ -990,6 +1004,7 @@ def _execute_integration_test(
                 dest_dir = os.path.dirname(dest_file)
                 if not os.path.exists(dest_dir):
                     os.makedirs(dest_dir)
+
                 shutil.copyfile(src=source_file, dst=dest_file)
 
         # UAT Script
diff --git a/tests/test_utils.py b/tests/test_utils.py
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -449,20 +449,38 @@ def delete_config_from_filesystem(
     )
 
 
+# noinspection PyPep8Naming
 def get_snowflake_connection_url() -> str:
     """Get snowflake connection url from environment variables.
 
     Returns:
-        String of the snowflake connection url.
+        String of the snowflake connection URL.
     """
-    sfAccount = os.environ.get("SNOWFLAKE_ACCOUNT")
     sfUser = os.environ.get("SNOWFLAKE_USER")
     sfPswd = os.environ.get("SNOWFLAKE_PW")
+    sfAccount = os.environ.get("SNOWFLAKE_ACCOUNT")
     sfDatabase = os.environ.get("SNOWFLAKE_DATABASE")
     sfSchema = os.environ.get("SNOWFLAKE_SCHEMA")
     sfWarehouse = os.environ.get("SNOWFLAKE_WAREHOUSE")
+    sfRole = os.environ.get("SNOWFLAKE_ROLE") or "PUBLIC"
+
+    return f"snowflake://{sfUser}:{sfPswd}@{sfAccount}/{sfDatabase}/{sfSchema}?warehouse={sfWarehouse}&role={sfRole}"
+
+
+def get_redshift_connection_url() -> str:
+    """Get Amazon Redshift connection url from environment variables.
+
+    Returns:
+        String of the Amazon Redshift connection URL.
+    """
+    host = os.environ.get("REDSHIFT_HOST")
+    port = os.environ.get("REDSHIFT_PORT")
+    user = os.environ.get("REDSHIFT_USERNAME")
+    pswd = os.environ.get("REDSHIFT_PASSWORD")
+    db = os.environ.get("REDSHIFT_DATABASE")
+    ssl = os.environ.get("REDSHIFT_SSLMODE")
 
-    return f"snowflake://{sfUser}:{sfPswd}@{sfAccount}/{sfDatabase}/{sfSchema}?warehouse={sfWarehouse}"
+    return f"redshift+psycopg2://{user}:{pswd}@{host}:{port}/{db}?sslmode={ssl}"
 
 
 def get_bigquery_table_prefix() -> str:
@@ -877,6 +895,8 @@ def get_connection_string_and_dialect(
     dialect: str = db_config["dialect"]
     if dialect == "snowflake":
         connection_string: str = get_snowflake_connection_url()
+    elif dialect == "redshift":
+        connection_string: str = get_redshift_connection_url()
     elif dialect == "bigquery":
         connection_string: str = get_bigquery_connection_url()
     elif dialect == "awsathena":