[FEATURE] SnowflakeDatasource update

great_expectations/datasource/fluent/__init__.py

@@ -8,6 +8,7 @@
     Sorter,
     BatchMetadata,
     GxDatasourceWarning,
+    GxContextWarning,
     TestConnectionError,
 )
 from great_expectations.datasource.fluent.invalid_datasource import (

great_expectations/datasource/fluent/config_str.py

@@ -2,13 +2,22 @@
 
 import logging
 import warnings
-from typing import TYPE_CHECKING, Mapping
-
-from great_expectations.compatibility.pydantic import SecretStr
+from typing import (
+    TYPE_CHECKING,
+    ClassVar,
+    Literal,
+    Mapping,
+    Optional,
+    TypedDict,
+)
+
+from great_expectations.compatibility.pydantic import AnyUrl, SecretStr, parse_obj_as
 from great_expectations.compatibility.typing_extensions import override
 from great_expectations.core.config_substitutor import TEMPLATE_STR_REGEX
 
 if TYPE_CHECKING:
+    from typing_extensions import Self, TypeAlias
+
     from great_expectations.core.config_provider import _ConfigurationProvider
     from great_expectations.datasource.fluent import Datasource
 
@@ -54,8 +63,15 @@ def __repr__(self) -> str:
         return f"{self.__class__.__name__}({self._display()!r})"
 
     @classmethod
-    def _validate_template_str_format(cls, v):
-        if TEMPLATE_STR_REGEX.search(v):
+    def str_contains_config_template(cls, v: str) -> bool:
+        """
+        Returns True if the input string contains a config template string.
+        """
+        return TEMPLATE_STR_REGEX.search(v) is not None
+
+    @classmethod
+    def _validate_template_str_format(cls, v: str) -> str | None:
+        if cls.str_contains_config_template(v):
             return v
         raise ValueError(
             cls.__name__
@@ -72,6 +88,125 @@ def __get_validators__(cls):
         yield cls.validate
 
 
+UriParts: TypeAlias = Literal[  # https://docs.pydantic.dev/1.10/usage/types/#url-properties
+    "scheme", "host", "user", "password", "port", "path", "query", "fragment", "tld"
+]
+
+
+class UriPartsDict(TypedDict, total=False):
+    scheme: str
+    user: str | None
+    password: str | None
+    ipv4: str | None
+    ipv6: str | None
+    domain: str | None
+    port: str | None
+    path: str | None
+    query: str | None
+    fragment: str | None
+
+
+class ConfigUri(AnyUrl, ConfigStr):  # type: ignore[misc] # Mixin "validate" signature mismatch
+    """
+    Special type that enables great_expectation config variable substitution for the
+    `user` and `password` section of a URI.
+
+    Example:
+    ```
+    "snowflake://${MY_USER}:${MY_PASSWORD}@account/database/schema/table"
+    ```
+
+    Note: this type is meant to used as part of pydantic model.
+    To use this outside of a model see the pydantic docs below.
+    https://docs.pydantic.dev/usage/models/#parsing-data-into-a-specified-type
+    """
+
+    ALLOWED_SUBSTITUTIONS: ClassVar[set[UriParts]] = {"user", "password"}
+
+    min_length: int = 1
+    max_length: int = 2**16
+
+    def __init__(  # noqa: PLR0913 # for compatibility with AnyUrl
+        self,
+        template_str: str,
+        *,
+        scheme: str,
+        user: Optional[str] = None,
+        password: Optional[str] = None,
+        host: Optional[str] = None,
+        tld: Optional[str] = None,
+        host_type: str = "domain",
+        port: Optional[str] = None,
+        path: Optional[str] = None,
+        query: Optional[str] = None,
+        fragment: Optional[str] = None,
+    ) -> None:
+        if template_str:  # may have already been set in __new__
+            self.template_str: str = template_str
+        self._secret_value = template_str  # for compatibility with SecretStr
+        super().__init__(
+            template_str,
+            scheme=scheme,
+            user=user,
+            password=password,
+            host=host,
+            tld=tld,
+            host_type=host_type,
+            port=port,
+            path=path,
+            query=query,
+            fragment=fragment,
+        )
+
+    def __new__(cls: type[Self], template_str: Optional[str], **kwargs) -> Self:
+        """custom __new__ for compatibility with pydantic.parse_obj_as()"""
+        built_url = cls.build(**kwargs) if template_str is None else template_str
+        instance = str.__new__(cls, built_url)
+        instance.template_str = str(instance)
+        return instance
+
+    @classmethod
+    @override
+    def validate_parts(cls, parts: UriPartsDict, validate_port: bool = True) -> UriPartsDict:
+        """
+        Ensure that only the `user` and `password` parts have config template strings.
+        Also validate that all parts of the URI are valid.
+        """
+        allowed_substitutions = sorted(cls.ALLOWED_SUBSTITUTIONS)
+
+        for name, part in parts.items():
+            if not part:
+                continue
+            if (
+                cls.str_contains_config_template(part)  # type: ignore[arg-type] # is str
+                and name not in cls.ALLOWED_SUBSTITUTIONS
+            ):
+                raise ValueError(  # noqa: TRY003
+                    f"Only {', '.join(allowed_substitutions)} may use config substitution; '{name}'"
+                    " substitution not allowed"
+                )
+        return AnyUrl.validate_parts(parts, validate_port)
+
+    @override
+    def get_config_value(self, config_provider: _ConfigurationProvider) -> AnyUrl:
+        """
+        Resolve the config template string to its string value according to the passed
+        _ConfigurationProvider.
+        Parse the resolved URI string into an `AnyUrl` object.
+        """
+        LOGGER.info(f"Substituting '{self}'")
+        raw_value = config_provider.substitute_config(self.template_str)
+        return parse_obj_as(AnyUrl, raw_value)
+
+    @classmethod
+    def __get_validators__(cls):
+        # one or more validators may be yielded which will be called in the
+        # order to validate the input, each validator will receive as an input
+        # the value returned from the previous validator
+        yield ConfigStr._validate_template_str_format
+        yield cls.validate  # equivalent to AnyUrl.validate
+
+
 def _check_config_substitutions_needed(
     datasource: Datasource,
     options: Mapping,

great_expectations/datasource/fluent/interfaces.py

@@ -181,6 +181,14 @@ class GxDatasourceWarning(UserWarning):
     """
 
 
+class GxContextWarning(GxDatasourceWarning):
+    """
+    Warning related to a Datasource that with a missing context.
+    Usually because the Datasource was created directly rather than using a
+    `context.sources` factory method.
+    """
+
+
 class GxSerializationWarning(GxDatasourceWarning):
     pass
 

great_expectations/datasource/fluent/schemas/SnowflakeDatasource.json

@@ -51,8 +51,9 @@
                 },
                 {
                     "type": "string",
-                    "writeOnly": true,
-                    "format": "password"
+                    "minLength": 1,
+                    "maxLength": 65536,
+                    "format": "uri"
                 },
                 {
                     "type": "string",
@@ -538,7 +539,7 @@
         },
         "ConnectionDetails": {
             "title": "ConnectionDetails",
-            "description": "Information needed to connect to a Snowflake database.\nAlternative to a connection string.",
+            "description": "Information needed to connect to a Snowflake database.\nAlternative to a connection string.\n\nhttps://docs.snowflake.com/en/developer-guide/python-connector/sqlalchemy#additional-connection-parameters",
             "type": "object",
             "properties": {
                 "account": {
@@ -564,10 +565,12 @@
                 },
                 "database": {
                     "title": "Database",
+                    "description": "`database` that the Datasource is mapped to.",
                     "type": "string"
                 },
                 "schema": {
                     "title": "Schema",
+                    "description": "`schema` that the Datasource is mapped to.",
                     "type": "string"
                 },
                 "warehouse": {
@@ -587,7 +590,9 @@
             "required": [
                 "account",
                 "user",
-                "password"
+                "password",
+                "database",
+                "schema"
             ],
             "additionalProperties": false
         }