From c1a792ad95d64ecf476eb6cf9ae04aa88abd87d0 Mon Sep 17 00:00:00 2001 From: stewartbryson Date: Fri, 24 Mar 2023 11:20:37 -0400 Subject: [PATCH] Support for partial connection properties. --- .../SnowflakeEphemeralTask.groovy | 46 +++++++++++++++++-- .../stewartbryson/SnowflakeExtension.groovy | 8 ++-- .../github/stewartbryson/SnowflakeJvm.groovy | 2 +- .../github/stewartbryson/SnowflakeTask.groovy | 28 +++++++---- src/markdown/README.md | 2 +- 5 files changed, 66 insertions(+), 20 deletions(-) diff --git a/plugin/src/main/groovy/io/github/stewartbryson/SnowflakeEphemeralTask.groovy b/plugin/src/main/groovy/io/github/stewartbryson/SnowflakeEphemeralTask.groovy index b4e5b43..57a1938 100644 --- a/plugin/src/main/groovy/io/github/stewartbryson/SnowflakeEphemeralTask.groovy +++ b/plugin/src/main/groovy/io/github/stewartbryson/SnowflakeEphemeralTask.groovy @@ -3,6 +3,7 @@ package io.github.stewartbryson import groovy.util.logging.Slf4j import org.gradle.api.tasks.CacheableTask import org.gradle.api.tasks.Input +import org.gradle.api.tasks.Internal import org.gradle.api.tasks.Optional import org.gradle.api.tasks.options.Option @@ -37,6 +38,24 @@ abstract class SnowflakeEphemeralTask extends SnowflakeTask { @Option(option = "ephemeral-name", description = "Optional: specify the ephemeral database name instead of relying on an autogenerated value.") String ephemeralName = extension.ephemeralName + /** + * The connection database in case it wasn't in the connection. + */ + @Internal + String connectionDatabase + + /** + * The connection schema in case it wasn't in the connection. + */ + @Internal + String connectionSchema + + /** + * The connection role in case it wasn't in the connection. + */ + @Internal + String connectionRole + /** * Create an ephemeral Snowflake clone and return a session to it. * @@ -44,9 +63,22 @@ abstract class SnowflakeEphemeralTask extends SnowflakeTask { */ def createClone() { if (useEphemeral) { - session.jdbcConnection().createStatement().execute("create database if not exists ${ephemeralName} clone $database") - session.jdbcConnection().createStatement().execute("grant ownership on database ${ephemeralName} to $role") - session.jdbcConnection().createStatement().execute("use schema ${ephemeralName}.${schema}") + // record connection attributes + try { + connectionDatabase = getSingleValue('SELECT CURRENT_DATABASE()') + connectionSchema = getSingleValue('SELECT CURRENT_SCHEMA()') + connectionRole = getSingleValue('SELECT CURRENT_ROLE()') + log.debug "Connection database, schema, role: $connectionDatabase, $connectionSchema, $connectionRole" + } catch (Exception e) { + throw new Exception("Connection context is not available.", e) + } + try { + session.jdbcConnection().createStatement().execute("create database if not exists ${ephemeralName} clone ${connectionDatabase}") + session.jdbcConnection().createStatement().execute("grant ownership on database ${ephemeralName} to ${connectionRole}") + session.jdbcConnection().createStatement().execute("use schema ${ephemeralName}.${connectionSchema}") + } catch (Exception e) { + throw new Exception("Cloning ephemeral clone failed.", e) + } log.warn "Ephemeral clone $ephemeralName created." } } @@ -57,8 +89,12 @@ abstract class SnowflakeEphemeralTask extends SnowflakeTask { def dropClone() { if (useEphemeral && !keepEphemeral) { // drop the ephemeral database - session.jdbcConnection().createStatement().execute("drop database if exists ${ephemeralName}") - session.jdbcConnection().createStatement().execute("use schema ${database}.${schema}") + try { + session.jdbcConnection().createStatement().execute("drop database if exists ${ephemeralName}") + session.jdbcConnection().createStatement().execute("use schema ${connectionDatabase}.${connectionSchema}") + } catch (Exception e) { + throw new Exception("Dropping ephemeral clone failed.", e) + } log.warn "Ephemeral clone $ephemeralName dropped." } } diff --git a/plugin/src/main/groovy/io/github/stewartbryson/SnowflakeExtension.groovy b/plugin/src/main/groovy/io/github/stewartbryson/SnowflakeExtension.groovy index d85589c..1fc2400 100644 --- a/plugin/src/main/groovy/io/github/stewartbryson/SnowflakeExtension.groovy +++ b/plugin/src/main/groovy/io/github/stewartbryson/SnowflakeExtension.groovy @@ -40,17 +40,17 @@ class SnowflakeExtension { */ String database /** - * The Snowflake schema to connect with. Default: 'public'. + * The Snowflake schema to connect with. */ - String schema = 'public' + String schema /** * The Snowflake role to connect with. */ String role /** - * The Snowflake warehouse to connect with. Default: 'compute_wh'. + * The Snowflake warehouse to connect with. */ - String warehouse = "compute_wh" + String warehouse /** * The Snowflake stage to upload to. Default: 'maven'. */ diff --git a/plugin/src/main/groovy/io/github/stewartbryson/SnowflakeJvm.groovy b/plugin/src/main/groovy/io/github/stewartbryson/SnowflakeJvm.groovy index d2b6457..48ef0af 100644 --- a/plugin/src/main/groovy/io/github/stewartbryson/SnowflakeJvm.groovy +++ b/plugin/src/main/groovy/io/github/stewartbryson/SnowflakeJvm.groovy @@ -116,7 +116,7 @@ abstract class SnowflakeJvm extends SnowflakeEphemeralTask { } } else if (extension.publishUrl) { // ensure that the stage and the publishUrl are aligned - String selectStage = getColumnValue("select stage_url from information_schema.stages where stage_name=upper('$stage') and stage_schema=upper('$schema') and stage_type='External Named'") + String selectStage = getSingleValue("select stage_url from information_schema.stages where stage_name=upper('$stage') and stage_schema=upper('$schema') and stage_type='External Named'") assert selectStage == extension.publishUrl } diff --git a/plugin/src/main/groovy/io/github/stewartbryson/SnowflakeTask.groovy b/plugin/src/main/groovy/io/github/stewartbryson/SnowflakeTask.groovy index 14e54e4..35bda0c 100644 --- a/plugin/src/main/groovy/io/github/stewartbryson/SnowflakeTask.groovy +++ b/plugin/src/main/groovy/io/github/stewartbryson/SnowflakeTask.groovy @@ -30,7 +30,7 @@ abstract class SnowflakeTask extends DefaultTask { } /** - * The Snowflake account URL, for instance: https://gradle-snowflake.us-east-1.snowflakecomputing.com:443. Overrides {@link SnowflakeExtension#account}. + * The Snowflake account URL, for instance: https://gradle-snowflake.us-east-1.snowflakecomputing.com. Overrides {@link SnowflakeExtension#account}. */ @Optional @Input @@ -103,14 +103,24 @@ abstract class SnowflakeTask extends DefaultTask { */ Session createSession() { Map props = [ - url : account, - user : user, - password : password, - role : role, - warehouse: warehouse, - db : database, - schema : schema + url : account, + user : user, + password: password ] + // add the optional connection parameters + if (role != 'null') { + props.role = role + } + if (warehouse != 'null') { + props.warehouse = warehouse + } + if (database != 'null') { + props.db = database + } + if (schema != 'null') { + props.schema = schema + } + Map printable = props.clone() printable.password = "*********" log.info "Snowflake config: $printable" @@ -136,7 +146,7 @@ abstract class SnowflakeTask extends DefaultTask { * * @return a scalar column value. */ - def getColumnValue(String sql) { + def getSingleValue(String sql) { Statement statement = session.jdbcConnection().createStatement() ResultSet rs = statement.executeQuery(sql) def columnValue diff --git a/src/markdown/README.md b/src/markdown/README.md index 9f2e48d..27474d4 100644 --- a/src/markdown/README.md +++ b/src/markdown/README.md @@ -407,7 +407,7 @@ BUILD SUCCESSFUL in 1s 1 actionable task: 1 executed ``` The `functionalTest` task contains all the tests that actually make a connection to Snowflake and test a deployment. -> WARNING: Ensure that the credential you provide below are for a safe development database. +> WARNING: Ensure that the credential provided below are for a safe development database. To run `functionalTest`, create the following entries in `~/.gradle/gradle.properties`: ```properties