[MAINTENANCE] Add schema to Redshift ConnectionDetails (#11431)

Author: shiplet

great_expectations/datasource/fluent/redshift_datasource.py

@@ -1,14 +1,10 @@
 from __future__ import annotations
 
 from enum import Enum
-from typing import TYPE_CHECKING, Literal, Type, Union
+from typing import TYPE_CHECKING, Literal, Optional, Type, Union
 
 from great_expectations._docs_decorators import public_api
-from great_expectations.compatibility.pydantic import (
-    AnyUrl,
-    BaseModel,
-    validator,
-)
+from great_expectations.compatibility.pydantic import AnyUrl, BaseModel, Field, validator
 from great_expectations.compatibility.typing_extensions import override
 from great_expectations.datasource.fluent.config_str import ConfigStr
 from great_expectations.datasource.fluent.sql_datasource import SQLDatasource
@@ -52,6 +48,12 @@ class RedshiftConnectionDetails(BaseModel):
     port: int
     database: str
     sslmode: RedshiftSSLModes
+    schema_: Optional[str] = Field(
+        default=None, alias="schema", description="`schema` that the Datasource is mapped to."
+    )
+
+    class Config:
+        allow_population_by_field_name = True
 
 
 @public_api
@@ -88,6 +90,8 @@ def _build_connection_string_from_connection_details(
         else:
             raise TypeError("Invalid connection_string type: ", type(connection_string))  # noqa: TRY003
         connection_string = f"redshift+psycopg2://{connection_details.user}:{connection_details.password}@{connection_details.host}:{connection_details.port}/{connection_details.database}?sslmode={connection_details.sslmode.value}"
+        if connection_details.schema_:
+            connection_string += f"&options=-csearch_path%3D{connection_details.schema_}"
         return connection_string
 
     @property

great_expectations/datasource/fluent/schemas/RedshiftDatasource.json

@@ -596,6 +596,11 @@
                 },
                 "sslmode": {
                     "$ref": "#/definitions/RedshiftSSLModes"
+                },
+                "schema": {
+                    "title": "Schema",
+                    "description": "`schema` that the Datasource is mapped to.",
+                    "type": "string"
                 }
             },
             "required": [

tests/datasource/fluent/test_redshift_datasource.py

@@ -1,4 +1,5 @@
 import logging
+from typing import Union
 
 import pytest
 from pytest_mock import MockerFixture
@@ -152,3 +153,72 @@ def test_connection_updating_plain_connection_string():
         f"Expected RedshiftDsn for plain connection string, "
         f"got {type(datasource.connection_string)}"
     )
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize(
+    "connection_input,expected_connection_string",
+    [
+        pytest.param(
+            {
+                "user": "user",
+                "password": "password",
+                "host": "host",
+                "port": 1234,
+                "database": "database",
+                "sslmode": RedshiftSSLModes.ALLOW,
+                "schema": "my_schema",
+            },
+            "redshift+psycopg2://user:password@host:1234/database?sslmode=allow&options=-csearch_path%3Dmy_schema",
+            id="dict type with schema",
+        ),
+        pytest.param(
+            RedshiftConnectionDetails(
+                user="user",
+                password="password",
+                host="host",
+                port=1234,
+                database="database",
+                sslmode=RedshiftSSLModes.ALLOW,
+                schema="my_schema",
+            ),
+            "redshift+psycopg2://user:password@host:1234/database?sslmode=allow&options=-csearch_path%3Dmy_schema",
+            id="RedshiftConnectionDetails with schema_",
+        ),
+        pytest.param(
+            {
+                "user": "user",
+                "password": "password",
+                "host": "host",
+                "port": 1234,
+                "database": "database",
+                "sslmode": RedshiftSSLModes.ALLOW,
+            },
+            "redshift+psycopg2://user:password@host:1234/database?sslmode=allow",
+            id="dict type without schema",
+        ),
+    ],
+)
+def test_schema_property_in_connection_string(
+    connection_input: Union[ConfigStr, RedshiftDsn],
+    expected_connection_string: str,
+    sa,
+    mocker: MockerFixture,
+    ephemeral_context_with_defaults: EphemeralDataContext,
+    scheme,
+):
+    create_engine_spy = mocker.patch.object(sa, "create_engine")
+
+    context = ephemeral_context_with_defaults
+    data_source = context.data_sources.add_redshift(
+        name="redshift_test_schema",
+        connection_string=connection_input,
+    )
+    data_source.get_engine()
+
+    expected_kwargs = RedshiftDsn(
+        expected_connection_string,
+        scheme=scheme,
+    )
+
+    create_engine_spy.assert_called_once_with(expected_kwargs)