def test_create_schema(self, mock_session_helper, mock_execute): mock_execute.return_value = MockScopeObj() mock_session_helper.db_session_scope.return_value.__enter__ = scope_execute_mock schema_name = "my_string" db_name = "my_database" iam_role = "my_iam_role" with mock_session_helper.db_session_scope() as mock_scope: publish_redshift.create_schema( schema_name, db_name, iam_role, mock_session_helper) mock_scope.execute.assert_called_once_with(f"CREATE EXTERNAL SCHEMA IF NOT EXISTS {schema_name} \ FROM DATA CATALOG \ database '{db_name}' \ iam_role '{iam_role}';")
def custom_publish(bucket: str, key: str, partitions: List[str], dataframe: pd.DataFrame, custom_redshift_columns: dict, redshift_params: dict = None) -> List[str]: """ Dataframe to S3 Parquet Publisher with a CUSTOM redshift column definition. Custom publish allows custom defined redshift column definitions to be used and enables support for Redshift's decimal data type. This function handles the portion of work that will see a dataframe converted to parquet and then published to the given S3 location. It supports partitions and will use the custom redshift columns defined in the custom_redshift_columns dictionary when creating the table schema for the parquet file. View the Custom Publishes section of s3parq's readme file for more guidance on formatting the custom_redshift_columns dictionary. It also has the option to automatically publish up to Redshift Spectrum for the newly published parquet files. Args: bucket (str): S3 Bucket name key (str): S3 key to lead to the desired dataset partitions (List[str]): List of columns that should be partitioned on dataframe (pd.DataFrame): Dataframe to be published custom_redshift_columns (dict): This dictionary contains custom column data type definitions for redshift. The params should be formatted as follows: - column name (str) - data type (str) redshift_params (dict, Optional): This dictionary should be provided in the following format in order for data to be published to Spectrum. Leave out entirely to avoid publishing to Spectrum. The params should be formatted as follows: - schema_name (str): Name of the Spectrum schema to publish to - table_name (str): Name of the table to write the dataset as - iam_role (str): Role to take while writing data to Spectrum - region (str): AWS region for Spectrum - cluster_id (str): Spectrum cluster id - host (str): Redshift Spectrum host name - port (str): Redshift Spectrum port to use - db_name (str): Redshift Spectrum database name to use - ec2_user (str): If on ec2, the user that should be used Returns: A str list of all the newly published object keys """ logger.debug("Running custom publish...") session_helper = None if redshift_params: if "index" in dataframe.columns: raise ValueError( "'index' is a reserved keyword in Redshift. Please remove or rename your DataFrame's 'index' column." ) logger.debug( "Found redshift parameters. Checking validity of params...") redshift_params = validate_redshift_params(redshift_params) logger.debug("Redshift parameters valid. Opening Session helper.") session_helper = SessionHelper( region=redshift_params['region'], cluster_id=redshift_params['cluster_id'], host=redshift_params['host'], port=redshift_params['port'], db_name=redshift_params['db_name'], ec2_user=redshift_params['ec2_user']) session_helper.configure_session_helper() publish_redshift.create_schema(redshift_params['schema_name'], redshift_params['db_name'], redshift_params['iam_role'], session_helper) logger.debug( f"Schema {redshift_params['schema_name']} created. Creating table {redshift_params['table_name']}..." ) publish_redshift.create_custom_table(redshift_params['table_name'], redshift_params['schema_name'], partitions, s3_url(bucket, key), custom_redshift_columns, session_helper) logger.debug(f"Custom table {redshift_params['table_name']} created.") logger.debug("Checking publish params...") check_empty_dataframe(dataframe) check_dataframe_for_timedelta(dataframe) check_partitions(partitions, dataframe) logger.debug("Publish params valid.") logger.debug("Begin writing to S3..") files = [] for frame_params in _sized_dataframes(dataframe): logger.info( f"Publishing dataframe chunk : {frame_params['lower']} to {frame_params['upper']}" ) frame = pd.DataFrame( dataframe[frame_params['lower']:frame_params['upper']]) _gen_parquet_to_s3(bucket=bucket, key=key, dataframe=frame, partitions=partitions, custom_redshift_columns=custom_redshift_columns) published_files = _assign_partition_meta( bucket=bucket, key=key, dataframe=frame, partitions=partitions, session_helper=session_helper, redshift_params=redshift_params, custom_redshift_columns=custom_redshift_columns) files = files + published_files logger.info("Done writing to S3.") return files
def publish(bucket: str, key: str, partitions: List['str'], dataframe: pd.DataFrame, redshift_params=None) -> None: """Redshift Params: ARGS: schema_name: str table_name: str iam_role: str region: str cluster_id: str host: str port: str db_name: str """ session_helper = None if redshift_params: if "index" in dataframe.columns: raise ValueError( "'index' is a reserved keyword in Redshift. Please remove or rename your DataFrame's 'index' column." ) logger.debug( "Found redshift parameters. Checking validity of params...") check_redshift_params(redshift_params) logger.debug("Redshift parameters valid. Opening Session helper.") session_helper = SessionHelper( region=redshift_params['region'], cluster_id=redshift_params['cluster_id'], host=redshift_params['host'], port=redshift_params['port'], db_name=redshift_params['db_name']) session_helper.configure_session_helper() publish_redshift.create_schema(redshift_params['schema_name'], redshift_params['db_name'], redshift_params['iam_role'], session_helper) logger.debug( f"Schema {redshift_params['schema_name']} created. Creating table {redshift_params['table_name']}..." ) df_types = _get_dataframe_datatypes(dataframe, partitions) partition_types = _get_dataframe_datatypes(dataframe, partitions, True) publish_redshift.create_table(redshift_params['table_name'], redshift_params['schema_name'], df_types, partition_types, s3_url(bucket, key), session_helper) logger.debug(f"Table {redshift_params['table_name']} created.") logger.info("Checking params...") check_empty_dataframe(dataframe) check_dataframe_for_timedelta(dataframe) check_partitions(partitions, dataframe) logger.info("Params valid.") logger.debug("Begin writing to S3..") files = [] for frame in _sized_dataframes(dataframe): _gen_parquet_to_s3(bucket=bucket, key=key, dataframe=frame, partitions=partitions) published_files = _assign_partition_meta( bucket=bucket, key=key, dataframe=frame, partitions=partitions, session_helper=session_helper, redshift_params=redshift_params) files = files + published_files logger.debug("Done writing to S3.") return files