Create a Service User with Custom Service User Objects

CREATE OR REPLACE PROCEDURE SETUP_ALTR_SERVICE_ACCOUNT(
  SERVICE_USER STRING DEFAULT 'ALTR_SERVICE_USER',
  DATABASE_NAMES ARRAY DEFAULT [],
  SERVICE_ROLE STRING DEFAULT NULL,
  SERVICE_WAREHOUSE STRING DEFAULT NULL
) RETURNS STRING
LANGUAGE JAVASCRIPT
VOLATILE
EXECUTE AS CALLER
AS $$
  // ********************************************
  // Stored Procedure Caller Report
  // ********************************************
  
  function AltrServiceAccountException(message) {
    this.message = message;
  }

  function StoredProcedureReport() {
    this.currentRole = "";
    this.isSuccess = true;
    this.successMessages = [];
    this.failMessages = [];
    this.unknownDBList = [];
    this.skippedDBList = [];

    this.fail = function(message) {
      this.isSuccess = false;
      this.failMessages.push(message);
    }

    this.success = function(message) {
      this.successMessages.push(message);
    }

    this.unknownDB = function(message) {
      this.unknownDBList.push(message);
    }

    this.skipped = function(message) {
      this.skippedDBList.push(message);
    }

    this.callerReport = function() {
      let report = 'ROLE EXECUTED AS: ' + this.currentRole + "\n\n";

      function appendMessages(prefix, messages) {
        const totalPrefix = '\n[' + prefix + ']: ';
        for (const message of messages) {
          report += totalPrefix + message;
        }
      }

      if (this.isSuccess) {
        report += 'SUCCEEDED!';
      } else {
        report += 'FAILED!';
        appendMessages('FAILURE', this.failMessages);
      }

      appendMessages('SUCCESS', this.successMessages);
      appendMessages('SKIPPED', this.skippedDBList);
      appendMessages('UNKNOWN OBJECT FOUND', this.unknownDBList);

      return report;
    }
  }

  var RUN_AS_ALTR = false;
  const storedProcedureReport = new StoredProcedureReport();

  // ********************************************
  // Utilities
  // ********************************************

  function isNullOrEmpty(str) {
    return typeof('string') !== typeof(str) || str.length < 1;
  }

  function tryGetColumnValueAsString(resultSet, columnName) {
    try {
      return resultSet.getColumnValueAsString(columnName);
    } catch (err) {
      return null;
    }
  }

  // https://docs.snowflake.com/en/sql-reference/stored-procedures-api.html
  function execQuery(queryString) {
    return snowflake.execute({sqlText:queryString});
  }

  function execStatement(queryString, handleError) {
    try {
      const resultSet = execQuery(queryString);
      storedProcedureReport.success(queryString);
      return true;
    } catch (error) {
      let reportFailure = true;
      if ('function' === typeof(handleError)) {
        try {
          reportFailure = handleError(queryString, error);
        } catch (error) {
          reportFailure = true;
        }
      }
      if (reportFailure) {
        storedProcedureReport.fail(queryString + ": " + error);
      }
    }
    return false;
  }

  // https://docs.snowflake.com/en/sql-reference/identifiers-syntax.html
  function delimitIdentifier(identifier) {
    return '"' + identifier.replace(/\"/g, '""') + '"';
  }

  function getDelimitedComment(objectType) {
    return "'This " + objectType + " is used by ALTR to help simplify governance and control over data in Snowflake. Please do not modify without speaking with ALTR Support.'";
  }

  function isUknownDB(queryString, error) {
    const errorString = '' + error;
    let idx = -1;
    idx = errorString.indexOf('Database');
    if (idx !== -1) {
      if (errorString.includes(' does not exist or not authorized.')) {
        storedProcedureReport.unknownDB(errorString.substring(idx));
        return false;
      }
    }
    return true;
  }

  function isDatabaseInList(databaseName) {
    return DATABASE_NAMES.find(function(element) { return element === databaseName });
  }

  // ********************************************
  // Main functionality
  // ********************************************

  function permissionRole(delimitedRoleName) {
    // ********************************************
    // Grant account-level privileges on the role
    // ********************************************

    let targetRoleExists = true;
    execStatement('GRANT CREATE DATABASE ON ACCOUNT TO ROLE ' + delimitedRoleName, function(queryString, error) {
      const errorString = '' + error;
      if (errorString.includes('Role ') && errorString.includes(' does not exist or not authorized.')) {
        // Use this initial grant to check validity of the input role.
        targetRoleExists = false;

        let extraInfo = '';
        const upperCaseTargetRoleName = ALTR_SERVICE_ROLE_NAME.toUpperCase();
        if (ALTR_SERVICE_ROLE_NAME !== upperCaseTargetRoleName) {
          extraInfo = '. Perhaps you meant role ' + upperCaseTargetRoleName;
        }

        storedProcedureReport.fail('You called this stored procedure with a role that does not exist: ' + ALTR_SERVICE_ROLE_NAME + extraInfo + '.');
        if (RUN_AS_ALTR) {
          throw new AltrServiceAccountException(storedProcedureReport.callerReport());
        }
        return false;
      } else {
        // Some other error when executing the grant. Report the
        // error, but keep executing.
        return true;
      }
    });

    if (!targetRoleExists) return;

    execStatement('GRANT APPLY MASKING POLICY ON ACCOUNT TO ROLE ' + delimitedRoleName, function(queryString, error) {
      // This stored procedure is written to execute on all editions of Snowflake,
      // including those that do not support masking policies. So, if we happen
      // to be on the latter, swallow this error silently and move on.
      return !('' + error).includes('Unsupported feature');
    });

    execStatement('GRANT CREATE INTEGRATION ON ACCOUNT TO ROLE ' + delimitedRoleName);
    execStatement('GRANT APPLY TAG ON ACCOUNT TO ROLE ' + delimitedRoleName);
    execStatement('GRANT APPLY ROW ACCESS POLICY ON ACCOUNT TO ROLE ' + delimitedRoleName);

    // ********************************************
    // Grant privileges on warehouses to role
    // ********************************************

    const warehouseNames = [];

    let resultSet = execQuery('SHOW WAREHOUSES');
    while (resultSet.next()) {
      warehouseNames.push(resultSet.getColumnValueAsString('name'));
    }

    for (const warehouseName of warehouseNames) {
      execStatement('GRANT MONITOR ON WAREHOUSE ' + delimitIdentifier(warehouseName) + ' TO ROLE ' + delimitedRoleName);
    }

    // ********************************************
    // Grant database-level privileges to role
    // ********************************************

    const databaseNames = [];

    resultSet = execQuery('SHOW DATABASES'); // Do not use TERSE here; we need some of the columns it trims
    while (resultSet.next()) {
      
      const name = resultSet.getColumnValueAsString('name');
      if (DATABASE_NAMES.length !== 0) {
          // We need to filter the result set down to the DATABASE_NAMES they have specified
          const database = DATABASE_NAMES.find(function(element) { return element === name });
          if (!database) {
              // the database provided in the array is not here (at least yet), on to the next.
              continue;
          }
      }
    
      // Exclude TRANSIENT databases. There are no TEMPORARY databases as of this writing,
      // but we might as well exclude them if they every implement them.
      const options = tryGetColumnValueAsString(resultSet, 'options');
      if (!isNullOrEmpty(options)) {
        const optionsUpper = options.toUpperCase();
        if (optionsUpper.includes('TRANSIENT') || optionsUpper.includes('TEMPORARY')) {
          storedProcedureReport.skipped('Database ' + name + ' because it has options: ' + options);
          continue;
        }
      }

      // Exclude any non-standard databases (so, shared and native app).
      const kind = tryGetColumnValueAsString(resultSet, 'kind');
      if (!isNullOrEmpty(kind) && !kind.toUpperCase().includes('STANDARD')) {
        storedProcedureReport.skipped('Database ' + name + ' because it is type: ' + kind);
        continue;
      }

      // Older versions of Snowflake did not provide us with the "kind"
      // column in SHOW DATABASES, so we had to check the "origin" of the
      // database. If there is an origin, it means the database "comes from
      // somewhere", and is, therefore, a SHARED database.
      const origin = tryGetColumnValueAsString(resultSet, 'origin');
      if (!isNullOrEmpty(origin)) {
        storedProcedureReport.skipped('Database ' + name + ' because it has origin: ' + origin);
        continue;
      }

      // Just a regular database, but exclude databases with empty or null names.
      if (!isNullOrEmpty(name)) {
        databaseNames.push(name);
      } else {
        storedProcedureReport.skipped('Database ' + name + ' because it has null or empty name');
      }
    }

    for (const databaseName of databaseNames) {
      const delimitedDatabaseName = delimitIdentifier(databaseName);
      let querySuccess = execStatement('GRANT USAGE ON DATABASE ' + delimitedDatabaseName + ' TO ROLE ' + delimitedRoleName, isUknownDB);
      if (querySuccess === false)
        continue;
      execStatement('GRANT CREATE SCHEMA ON DATABASE ' + delimitedDatabaseName + ' TO ROLE ' + delimitedRoleName, isUknownDB);
      execStatement('GRANT USAGE ON FUTURE SCHEMAS IN DATABASE ' + delimitedDatabaseName + ' TO ROLE ' + delimitedRoleName, isUknownDB);
      execStatement('GRANT USAGE ON ALL SCHEMAS IN DATABASE ' + delimitedDatabaseName + ' TO ROLE ' + delimitedRoleName, isUknownDB);
      execStatement('GRANT SELECT ON FUTURE TABLES IN DATABASE ' + delimitedDatabaseName + ' TO ROLE ' + delimitedRoleName, isUknownDB);
      execStatement('GRANT SELECT ON ALL TABLES IN DATABASE ' + delimitedDatabaseName + ' TO ROLE ' + delimitedRoleName, isUknownDB);
      execStatement('GRANT SELECT ON FUTURE VIEWS IN DATABASE ' + delimitedDatabaseName + ' TO ROLE ' + delimitedRoleName, isUknownDB);
      execStatement('GRANT SELECT ON ALL VIEWS IN DATABASE ' + delimitedDatabaseName + ' TO ROLE ' + delimitedRoleName, isUknownDB);
      execStatement('GRANT SELECT ON FUTURE MATERIALIZED VIEWS IN DATABASE ' + delimitedDatabaseName + ' TO ROLE ' + delimitedRoleName, isUknownDB);
      execStatement('GRANT SELECT ON ALL MATERIALIZED VIEWS IN DATABASE ' + delimitedDatabaseName + ' TO ROLE ' + delimitedRoleName, isUknownDB);
      execStatement('GRANT CREATE TAG ON FUTURE SCHEMAS IN DATABASE ' + delimitedDatabaseName + ' TO ROLE ' + delimitedRoleName, isUknownDB);
      execStatement('GRANT CREATE TAG ON ALL SCHEMAS IN DATABASE ' + delimitedDatabaseName + ' TO ROLE ' + delimitedRoleName, isUknownDB);
      execStatement('GRANT CREATE ROW ACCESS POLICY ON FUTURE SCHEMAS IN DATABASE ' + delimitedDatabaseName + ' TO ROLE ' + delimitedRoleName, isUknownDB);
      execStatement('GRANT CREATE ROW ACCESS POLICY ON ALL SCHEMAS IN DATABASE ' + delimitedDatabaseName + ' TO ROLE ' + delimitedRoleName, isUknownDB);
    }

    execStatement('GRANT ALL ON FUTURE SCHEMAS IN DATABASE ALTR_DSAAS_DB TO ROLE ' + delimitedRoleName, isUknownDB);
    execStatement('GRANT ALL ON ALL SCHEMAS IN DATABASE ALTR_DSAAS_DB TO ROLE ' + delimitedRoleName, isUknownDB);
    execStatement('GRANT ALL ON FUTURE FUNCTIONS IN DATABASE ALTR_DSAAS_DB TO ROLE ' + delimitedRoleName, isUknownDB);
    execStatement('GRANT ALL ON ALL FUNCTIONS IN DATABASE ALTR_DSAAS_DB TO ROLE ' + delimitedRoleName, isUknownDB);
    execStatement('GRANT ALL ON DATABASE ALTR_DSAAS_DB TO ROLE ' + delimitedRoleName, isUknownDB);
    execStatement('GRANT MANAGE GRANTS ON ACCOUNT TO ROLE ' + delimitedRoleName, isUknownDB);

    execStatement('GRANT IMPORTED PRIVILEGES ON DATABASE SNOWFLAKE TO ROLE ' + delimitedRoleName, isUknownDB);
  }

  function permissionUser(delimitedRoleName, delimitedUserName, delimitedWarehouseName) {
    // Create the role for ALTR and assign the role to SYSADMIN as recommended by
    // Snowflake to help with troubleshooting and debugging if necessary.
    execStatement('CREATE ROLE IF NOT EXISTS ' + delimitedRoleName);
    execStatement('ALTER ROLE ' + delimitedRoleName + ' SET COMMENT = ' + getDelimitedComment('role'));
    execStatement('GRANT ROLE ' + delimitedRoleName + ' TO ROLE SYSADMIN');

    // Create the warehouse for ALTR and grant usage of it to the ALTR role.
    // Though XSMALL is the default for WAREHOUSE_SIZE and TRUE is the default
    // for AUTO_RESUME, we include them anyway because if the defaults of these
    // properties ever change, there will be serious functionality and cost concerns.
    execStatement('CREATE WAREHOUSE IF NOT EXISTS ' + delimitedWarehouseName
      + ' WITH WAREHOUSE_SIZE = XSMALL SCALING_POLICY = ECONOMY AUTO_RESUME = TRUE INITIALLY_SUSPENDED = TRUE');
    execStatement('ALTER WAREHOUSE ' + delimitedWarehouseName + ' SET AUTO_SUSPEND = 30 COMMENT = ' + getDelimitedComment('warehouse'));
    execStatement('GRANT USAGE ON WAREHOUSE ' + delimitedWarehouseName + ' TO ROLE ' + delimitedRoleName);

    // Transfer ownership of the warehouse to customer SYSADMIN. They can then
    // assign resource monitors and other customer-specific configuration.
    // Though this procedure will necessarily run as ACCOUNTADMIN, Snowflake
    // recommeneds that database objects be owned by SYSADMIN.
    execStatement('GRANT OWNERSHIP ON WAREHOUSE ' + delimitedWarehouseName + ' TO ROLE SYSADMIN COPY CURRENT GRANTS');

    // Permission the ALTR role. Be sure to do this after you create the
    // ALTR warehouse so the ALTR role gets monitor on that warehouse.
    permissionRole(delimitedRoleName);

    // Finally, assign the role to the user (the latter must already be created)
    // and set appropriate defaults.
    execStatement('GRANT ROLE ' + delimitedRoleName + ' TO USER ' + delimitedUserName);
    execStatement('ALTER USER ' + delimitedUserName + ' SET DEFAULT_ROLE = ' + delimitedRoleName);
    execStatement('ALTER USER ' + delimitedUserName + ' SET DEFAULT_WAREHOUSE = ' + delimitedWarehouseName);
    execStatement('ALTER USER ' + delimitedUserName + ' SET COMMENT = ' + getDelimitedComment('user'));
  }

  // ********************************************
  // Main
  // ********************************************

  let paramsResultSet = execQuery("SHOW PARAMETERS LIKE '%QUERY_TAG%' IN SESSION");
  while (paramsResultSet.next()) {
    let queryTagStr = '';
    queryTagStr = paramsResultSet.getColumnValueAsString('value');
    if (!isNullOrEmpty(queryTagStr) && queryTagStr === "ALTR") {
      RUN_AS_ALTR = true;
      break;
    }
  }

  // get the current role that the user is before we switch to accountadmin
  let currentRoles = [];
  let resultSet = execQuery("SELECT CURRENT_ROLE() AS CURRENT_ROLE;");
  while (resultSet.next()) {
     currentRoles.push(tryGetColumnValueAsString(resultSet, "CURRENT_ROLE"));
  }
  storedProcedureReport.currentRole = currentRoles[0];

  // This script must execute as ACCOUNTADMIN so that
  // we can grant the target role permissions on objects
  // that the target role may not even be able to see.
  try {
    execQuery("USE ROLE ACCOUNTADMIN");
  } catch (error) {
    storedProcedureReport.fail('You must call this stored procedure with a user and role that can assume role ACCOUNTADMIN: ' + error);
    if (RUN_AS_ALTR) {
      throw new AltrServiceAccountException(storedProcedureReport.callerReport());
    }
    return storedProcedureReport.callerReport();
  }

  // Create the database asap so that we can grant the role permissions on it.
  execStatement('CREATE DATABASE IF NOT EXISTS ALTR_DSAAS_DB'); // ownership of this database will be accountadmin

  if (SERVICE_USER === "PC_ALTR_USER") {
    // if this user is Partner Connect and SERVICE_ROLE is not defined, assume PC_ALTR_ROLE, else use what they gave.
    let pcRole = !SERVICE_ROLE ? '"PC_ALTR_ROLE"' : SERVICE_ROLE;
    let pcWarehouse = !SERVICE_WAREHOUSE ? '"PC_ALTR_WH"' : SERVICE_WAREHOUSE;
    execStatement('ALTER WAREHOUSE ' + pcWarehouse +' SET AUTO_SUSPEND = 30');
    permissionRole(pcRole);
  } else {
    let serviceUser = !SERVICE_USER ? '"ALTR_SERVICE_USER"' : SERVICE_USER;
    let serviceRole = !SERVICE_ROLE ? '"ALTR_SERVICE_ROLE"' : SERVICE_ROLE;
    let serviceWarehouse = !SERVICE_WAREHOUSE ? '"ALTR_SERVICE_WH"' : SERVICE_WAREHOUSE;
    permissionUser(serviceRole, serviceUser, serviceWarehouse);
  }

  report = storedProcedureReport.callerReport();
  if (storedProcedureReport.isSuccess) {
    execStatement("ALTER PROCEDURE IF EXISTS {{.Database}}.{{.Schema}}.SETUP_ALTR_SERVICE_ACCOUNT(STRING, ARRAY, STRING, STRING) SET COMMENT = '" + new Date().toUTCString() + "'");
  } else if (RUN_AS_ALTR) {
    throw new AltrServiceAccountException(report)
  }
  return report;
$$;