s3 connector (for data detection & discovery) POC - fides

CHANGELOG.md

@@ -23,6 +23,7 @@ The types of changes are:
 - Added option in FidesJS SDK to only disable notice-served API [#4965](https://github.com/ethyca/fides/pull/4965)
 - External ID support for consent management [#4927](https://github.com/ethyca/fides/pull/4927)
 - Added access and erasure support for the Greenhouse Harvest integration [#4945](https://github.com/ethyca/fides/pull/4945)
+- Add an S3 connection type (currently used for discovery and detection only) [#4930](https://github.com/ethyca/fides/pull/4930)
 - Support for Limited FIDES__CELERY__* Env Vars [#4980](https://github.com/ethyca/fides/pull/4980)
 - Implement sending emails via property-specific messaging templates [#4950](https://github.com/ethyca/fides/pull/4950)
 

clients/admin-ui/src/features/datastore-connections/constants.ts

@@ -46,6 +46,7 @@ export const CONNECTION_TYPE_LOGO_MAP = new Map<ConnectionType, string>([
   [ConnectionType.MYSQL, "mysql.svg"],
   [ConnectionType.POSTGRES, "postgres.svg"],
   [ConnectionType.REDSHIFT, "redshift.svg"],
+  [ConnectionType.S3, "s3.svg"],
   [ConnectionType.SCYLLA, "scylla.svg"],
   [ConnectionType.SNOWFLAKE, "snowflake.svg"],
   [ConnectionType.SOVRN, "sovrn.svg"],

clients/admin-ui/src/types/api/models/ConnectionType.ts

@@ -23,6 +23,7 @@ export enum ConnectionType {
   MANUAL_WEBHOOK = "manual_webhook",
   TIMESCALE = "timescale",
   FIDES = "fides",
+  S3 = "s3",
   SCYLLA = "scylla",
   GENERIC_ERASURE_EMAIL = "generic_erasure_email",
   GENERIC_CONSENT_EMAIL = "generic_consent_email",

src/fides/api/alembic/migrations/versions/2736c942faa2_property_specific_messaging_db_models.py

@@ -1,7 +1,7 @@
 """property specific messaging db models
 
 Revision ID: 2736c942faa2
-Revises: efddde14da21
+Revises: 52a5f1a957bc
 Create Date: 2024-05-28 14:26:09.012859
 
 """

src/fides/api/alembic/migrations/versions/cb344673f633_add_s3_connection_type.py

@@ -0,0 +1,47 @@
+"""add s3 connection type
+
+Revision ID: cb344673f633
+Revises: 3304082a6cee
+Create Date: 2024-05-31 20:46:08.829330
+
+"""
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "cb344673f633"
+down_revision = "3304082a6cee"
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # Add 's3' to ConnectionType enum
+    op.execute("alter type connectiontype rename to connectiontype_old")
+    op.execute(
+        "create type connectiontype as enum('mongodb', 'mysql', 'https', 'snowflake', 'redshift', 'mssql', 'mariadb', 'bigquery', 'saas', 'manual', 'manual_webhook', 'timescale', 'fides', 'sovrn', 'attentive', 'dynamodb', 'postgres', 'generic_consent_email', 'generic_erasure_email', 'scylla', 's3')"
+    )
+    op.execute(
+        (
+            "alter table connectionconfig alter column connection_type type connectiontype using "
+            "connection_type::text::connectiontype"
+        )
+    )
+    op.execute("drop type connectiontype_old")
+
+
+def downgrade():
+    # Remove 's3' from ConnectionType enum
+    op.execute("delete from connectionconfig where connection_type in ('s3')")
+    op.execute("alter type connectiontype rename to connectiontype_old")
+    op.execute(
+        "create type connectiontype as enum('mongodb', 'mysql', 'https', 'snowflake', 'redshift', 'mssql', 'mariadb', 'bigquery', 'saas', 'manual', 'manual_webhook', 'timescale', 'fides', 'sovrn', 'attentive', 'dynamodb', 'postgres', 'generic_consent_email', 'generic_erasure_email', 'scylla')"
+    )
+    op.execute(
+        (
+            "alter table connectionconfig alter column connection_type type connectiontype using "
+            "connection_type::text::connectiontype"
+        )
+    )
+    op.execute("drop type connectiontype_old")

src/fides/api/api/v1/endpoints/storage_endpoints.py

@@ -38,8 +38,8 @@
 from fides.api.schemas.storage.data_upload_location_response import DataUpload
 from fides.api.schemas.storage.storage import (
     FULLY_CONFIGURED_STORAGE_TYPES,
+    AWSAuthMethod,
     BulkPutStorageConfigResponse,
-    S3AuthMethod,
     StorageConfigStatus,
     StorageConfigStatusMessage,
     StorageDestination,
@@ -420,7 +420,7 @@ def get_storage_status(
 def _storage_config_requires_secrets(storage_config: StorageConfig) -> bool:
     return (
         storage_config.details.get(StorageDetails.AUTH_METHOD.value, None)
-        == S3AuthMethod.SECRET_KEYS.value
+        == AWSAuthMethod.SECRET_KEYS.value
     )
 
 

src/fides/api/models/connectionconfig.py

@@ -53,6 +53,7 @@ class ConnectionType(enum.Enum):
     fides = "fides"
     generic_erasure_email = "generic_erasure_email"  # Run after the traversal
     generic_consent_email = "generic_consent_email"  # Run after the traversal
+    s3 = "s3"
     scylla = "scylla"
 
     @property
@@ -76,6 +77,7 @@ def human_readable(self) -> str:
             ConnectionType.mysql.value: "MySQL",
             ConnectionType.postgres.value: "PostgreSQL",
             ConnectionType.redshift.value: "Amazon Redshift",
+            ConnectionType.s3.value: "Amazon S3",
             ConnectionType.saas.value: "SaaS",
             ConnectionType.scylla.value: "Scylla DB",
             ConnectionType.snowflake.value: "Snowflake",

src/fides/api/schemas/connection_configuration/__init__.py

@@ -80,6 +80,12 @@
 from fides.api.schemas.connection_configuration.connection_secrets_redshift import (
     RedshiftSchema as RedshiftSchema,
 )
+from fides.api.schemas.connection_configuration.connection_secrets_s3 import (
+    S3DocsSchema as S3DocsSchema,
+)
+from fides.api.schemas.connection_configuration.connection_secrets_s3 import (
+    S3Schema as S3Schema,
+)
 from fides.api.schemas.connection_configuration.connection_secrets_saas import (
     SaaSSchema as SaaSSchema,
 )
@@ -127,6 +133,7 @@
     ConnectionType.postgres.value: PostgreSQLSchema,
     ConnectionType.redshift.value: RedshiftSchema,
     ConnectionType.saas.value: SaaSSchema,
+    ConnectionType.s3.value: S3Schema,
     ConnectionType.scylla.value: ScyllaSchema,
     ConnectionType.snowflake.value: SnowflakeSchema,
     ConnectionType.sovrn.value: SovrnSchema,
@@ -176,5 +183,6 @@ def get_connection_secrets_schema(
     FidesDocsSchema,
     SovrnDocsSchema,
     DynamoDBDocsSchema,
+    S3DocsSchema,
     ScyllaDocsSchema,
 ]

src/fides/api/schemas/connection_configuration/connection_secrets_dynamodb.py

@@ -25,6 +25,8 @@ class DynamoDBSchema(ConnectionConfigSecretsSchema):
         sensitive=True,
     )
 
+    # TODO: include an aws_assume_role_arn and more closely follow the pattern in `connection_secrets_s3`
+
     _required_components: List[str] = [
         "region_name",
         "aws_access_key_id",

src/fides/api/schemas/connection_configuration/connection_secrets_s3.py

@@ -0,0 +1,54 @@
+from typing import Dict, List, Optional
+
+from pydantic import Field, root_validator
+
+from fides.api.schemas.base_class import NoValidationSchema
+from fides.api.schemas.connection_configuration.connection_secrets import (
+    ConnectionConfigSecretsSchema,
+)
+from fides.api.schemas.storage.storage import AWSAuthMethod
+
+
+class S3Schema(ConnectionConfigSecretsSchema):
+    """Schema to validate the secrets needed to connect to Amazon S3"""
+
+    auth_method: AWSAuthMethod = Field(
+        title="Authentication Method",
+        description="Determines which type of authentication method to use for connecting to Amazon S3",
+    )
+
+    aws_access_key_id: Optional[str] = Field(
+        title="Access Key ID",
+        description="Part of the credentials that provide access to your AWS account. This is required if using secret key authentication.",
+    )
+    aws_secret_access_key: Optional[str] = Field(
+        title="Secret Access Key",
+        description="Part of the credentials that provide access to your AWS account. This is required if using secret key authentication.",
+        sensitive=True,
+    )
+
+    aws_assume_role_arn: Optional[str] = Field(
+        title="Assume Role ARN",
+        description="If provided, the ARN of the role that should be assumed to connect to s3.",
+    )
+
+    _required_components: List[str] = ["auth_method"]
+
+    @root_validator(pre=True)
+    @classmethod
+    def keys_provided_if_needed(cls, values: Dict) -> Dict:
+        """
+        Validates that both access and secret access keys are provided if using a `secret_keys` auth method.
+        """
+        if values.get("auth_method") == AWSAuthMethod.SECRET_KEYS.value and not (
+            values.get("aws_access_key_id") and values.get("aws_secret_access_key")
+        ):
+            raise ValueError(
+                f"An Access Key ID and a Secret Access Key must be provided if using the `{AWSAuthMethod.SECRET_KEYS.value}` Authentication Method"
+            )
+
+        return values
+
+
+class S3DocsSchema(S3Schema, NoValidationSchema):
+    """S3 Secrets Schema for API Docs"""

src/fides/api/schemas/storage/storage.py

@@ -48,15 +48,15 @@ class Config:
         extra = Extra.forbid
 
 
-class S3AuthMethod(Enum):
+class AWSAuthMethod(str, Enum):
     AUTOMATIC = "automatic"
     SECRET_KEYS = "secret_keys"
 
 
 class StorageDetailsS3(FileBasedStorageDetails):
     """The details required to represent an AWS S3 storage bucket."""
 
-    auth_method: S3AuthMethod
+    auth_method: AWSAuthMethod
     bucket: str
     max_retries: Optional[int] = 0
 

src/fides/api/service/connectors/__init__.py

@@ -30,6 +30,7 @@
 from fides.api.service.connectors.mongodb_connector import (
     MongoDBConnector as MongoDBConnector,
 )
+from fides.api.service.connectors.s3_connector import S3Connector
 from fides.api.service.connectors.saas_connector import SaaSConnector as SaaSConnector
 from fides.api.service.connectors.scylla_connector import ScyllaConnector
 from fides.api.service.connectors.sql_connector import (
@@ -75,6 +76,7 @@
     ConnectionType.snowflake.value: SnowflakeConnector,
     ConnectionType.sovrn.value: SovrnConnector,
     ConnectionType.timescale.value: TimescaleConnector,
+    ConnectionType.s3.value: S3Connector,
 }
 
 

src/fides/api/service/connectors/s3_connector.py

@@ -0,0 +1,73 @@
+from typing import Any, Dict, List, Optional
+
+from loguru import logger
+
+from fides.api.common_exceptions import ConnectionException
+from fides.api.graph.execution import ExecutionNode
+from fides.api.models.connectionconfig import ConnectionTestStatus
+from fides.api.models.policy import Policy
+from fides.api.models.privacy_request import PrivacyRequest, RequestTask
+from fides.api.schemas.connection_configuration.connection_secrets_s3 import S3Schema
+from fides.api.service.connectors.base_connector import BaseConnector
+from fides.api.service.connectors.query_config import QueryConfig
+from fides.api.util.aws_util import get_aws_session
+from fides.api.util.collection_util import Row
+
+
+class S3Connector(BaseConnector):
+    """
+    AWS S3 Connector - this is currently used just to test connections to S3.
+
+    NOTE: No DSR processing is yet supported for S3.
+    """
+
+    def create_client(self) -> Any:  # type: ignore
+        """Returns a client for s3"""
+        config = S3Schema(**self.configuration.secrets or {})
+        return get_aws_session(
+            auth_method=config.auth_method.value,
+            storage_secrets=config.dict(),  # type: ignore
+            assume_role_arn=config.aws_assume_role_arn,
+        )
+
+    def query_config(self, node: ExecutionNode) -> QueryConfig[Any]:
+        """DSR execution not yet supported for S3"""
+        raise NotImplementedError()
+
+    def test_connection(self) -> Optional[ConnectionTestStatus]:
+        """
+        Connects to AWS S3 and gets caller identity to validate credentials.
+        """
+        logger.info("Starting test connection to {}", self.configuration.key)
+        try:
+            session = self.client()
+            sts_client = session.client("sts")
+            sts_client.get_caller_identity()
+        except Exception as error:
+            raise ConnectionException(str(error))
+
+        return ConnectionTestStatus.succeeded
+
+    def retrieve_data(
+        self,
+        node: ExecutionNode,
+        policy: Policy,
+        privacy_request: PrivacyRequest,
+        request_task: RequestTask,
+        input_data: Dict[str, List[Any]],
+    ) -> List[Row]:
+        """DSR execution not yet supported for S3"""
+
+    def mask_data(
+        self,
+        node: ExecutionNode,
+        policy: Policy,
+        privacy_request: PrivacyRequest,
+        request_task: RequestTask,
+        rows: List[Row],
+    ) -> int:
+        """DSR execution not yet supported for S3"""
+
+    def close(self) -> None:
+        """Close any held resources"""
+        # no held resources for S3 connector

src/fides/api/service/storage/storage_authenticator_service.py

@@ -4,11 +4,11 @@
 
 from fides.api.schemas.storage.storage import (
     SUPPORTED_STORAGE_SECRETS,
-    S3AuthMethod,
+    AWSAuthMethod,
     StorageSecrets,
     StorageType,
 )
-from fides.api.util.storage_authenticator import get_s3_session
+from fides.api.util.aws_util import get_aws_session
 
 
 def secrets_are_valid(
@@ -31,7 +31,7 @@ def secrets_are_valid(
 def _s3_authenticator(secrets: Dict[StorageSecrets, Any]) -> bool:
     """Authenticates secrets for s3, returns true if secrets are valid"""
     try:
-        get_s3_session(S3AuthMethod.SECRET_KEYS.value, secrets.dict())  # type: ignore
+        get_aws_session(AWSAuthMethod.SECRET_KEYS.value, secrets.dict())  # type: ignore
         return True
     except ClientError:
         return False

src/fides/api/task/task_resources.py

@@ -24,6 +24,7 @@
     TimescaleConnector,
 )
 from fides.api.service.connectors.base_email_connector import BaseEmailConnector
+from fides.api.service.connectors.s3_connector import S3Connector
 from fides.api.util.cache import get_cache
 from fides.api.util.collection_util import Row, extract_key_for_address
 
@@ -74,6 +75,8 @@
             return DynamoDBConnector(connection_config)
         if connection_config.connection_type == ConnectionType.fides:
             return FidesConnector(connection_config)
+        if connection_config.connection_type == ConnectionType.s3:
+            return S3Connector(connection_config)
         raise NotImplementedError(
             f"No connector available for {connection_config.connection_type}"
         )