Skip to content

Commit af93464

Browse files
AcevedoRloicmathieu
AcevedoR
and
loicmathieu
authored
feat(snowflake): add SnowflakeCLI task (#551)
* feat(snowflake): add SnowflakeCLI task * doc(snowflake): add link to Snowflake doc for command parameter * Update plugin-jdbc-snowflake/src/main/java/io/kestra/plugin/jdbc/snowflake/SnowflakeCLI.java Co-authored-by: Loïc Mathieu <loikeseke@gmail.com> * feat(snowflake): remove useless beforeCommands parameter --------- Co-authored-by: Loïc Mathieu <loikeseke@gmail.com>
1 parent 0d212cd commit af93464

File tree

3 files changed

+230
-0
lines changed

3 files changed

+230
-0
lines changed

‎plugin-jdbc-snowflake/build.gradle‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,5 +16,8 @@ dependencies {
1616
implementation("net.snowflake:snowflake-jdbc:3.21.0")
1717
implementation project(':plugin-jdbc')
1818

19+
compileOnly group: "io.kestra", name: "script", version: kestraVersion
20+
21+
testImplementation group: "io.kestra", name: "script", version: kestraVersion
1922
testImplementation project(':plugin-jdbc').sourceSets.test.output
2023
}

‎plugin-jdbc-snowflake/src/main/java/io/kestra/plugin/jdbc/snowflake/SnowflakeCLI.java‎

Lines changed: 176 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,176 @@
1+
package io.kestra.plugin.jdbc.snowflake;
2+
3+
import io.kestra.core.models.annotations.Example;
4+
import io.kestra.core.models.annotations.Plugin;
5+
import io.kestra.core.models.annotations.PluginProperty;
6+
import io.kestra.core.models.property.Property;
7+
import io.kestra.core.models.tasks.*;
8+
import io.kestra.core.models.tasks.runners.TaskRunner;
9+
import io.kestra.core.runners.RunContext;
10+
import io.kestra.plugin.scripts.exec.scripts.models.ScriptOutput;
11+
import io.kestra.plugin.scripts.exec.scripts.runners.CommandsWrapper;
12+
import io.kestra.plugin.scripts.runner.docker.Docker;
13+
import io.swagger.v3.oas.annotations.media.Schema;
14+
import jakarta.validation.Valid;
15+
import jakarta.validation.constraints.NotNull;
16+
import lombok.*;
17+
import lombok.experimental.SuperBuilder;
18+
19+
import java.util.*;
20+
import java.util.stream.Stream;
21+
22+
@SuperBuilder
23+
@ToString
24+
@EqualsAndHashCode
25+
@Getter
26+
@NoArgsConstructor
27+
@Schema(
28+
title = "Run Snowflake commands."
29+
)
30+
@Plugin(
31+
examples = {
32+
@Example(
33+
title = "Show basic infos and connection status",
34+
full = true,
35+
code = {
36+
"""
37+
id: snowflake
38+
namespace: company.team
39+
tasks:
40+
- id: log_info_and_connection_status
41+
type: io.kestra.plugin.jdbc.snowflake.SnowflakeCLI
42+
account: snowflake_account
43+
username: snowflake_username
44+
password: snowflake_password
45+
commands:
46+
- snow --info
47+
- snow connection test
48+
"""
49+
}
50+
),
51+
@Example(
52+
title = "List Snowflake staged files",
53+
full = true,
54+
code = {
55+
"""
56+
id: snowflake
57+
namespace: company.team
58+
tasks:
59+
- id: list_stage_files
60+
type: io.kestra.plugin.jdbc.snowflake.SnowflakeCLI
61+
account: snowflake_account
62+
username: snowflake_username
63+
password: snowflake_password
64+
commands:
65+
- snow stage list-files @MY_WAREHOUSE.MY_SCHEMA.%MY_TABLE_STAGE_NAME
66+
"""
67+
}
68+
),
69+
@Example(
70+
title = "Run Snowflake SQL select via CLI",
71+
full = true,
72+
code = {
73+
"""
74+
id: snowflake
75+
namespace: company.team
76+
tasks:
77+
- id: query
78+
type: io.kestra.plugin.jdbc.snowflake.SnowflakeCLI
79+
account: snowflake_account
80+
username: snowflake_username
81+
password: snowflake_password
82+
commands:
83+
- snow sql --query="SELECT 1"
84+
"""
85+
}
86+
)
87+
}
88+
)
89+
public class SnowflakeCLI extends Task implements RunnableTask<ScriptOutput>, NamespaceFilesInterface, InputFilesInterface, OutputFilesInterface {
90+
91+
public static final String DEFAULT_IMAGE = "ghcr.io/kestra-io/snowflake-cli";
92+
93+
@Schema(
94+
title = "The commands to run. Please refer to SnowflakeCLI documentation https://docs.snowflake.com/en/developer-guide/snowflake-cli/command-reference/overview"
95+
)
96+
@NotNull
97+
protected Property<List<String>> commands;
98+
99+
@Schema(
100+
title = "The account to use for authentication."
101+
)
102+
@NotNull
103+
protected Property<String> account;
104+
105+
@Schema(
106+
title = "The username to use for authentication."
107+
)
108+
@NotNull
109+
protected Property<String> username;
110+
111+
@Schema(
112+
title = "The password to use for authentication."
113+
)
114+
@NotNull
115+
protected Property<String> password;
116+
117+
@Schema(
118+
title = "Additional environment variables for the current process."
119+
)
120+
@PluginProperty(dynamic = true)
121+
protected Map<String, String> env;
122+
123+
@Schema(
124+
title = "The task runner to use."
125+
)
126+
@Valid
127+
@PluginProperty
128+
@Builder.Default
129+
private TaskRunner<?> taskRunner = Docker.instance();
130+
131+
@Schema(
132+
title = "The snowflake container image."
133+
)
134+
@PluginProperty(dynamic = true)
135+
@Builder.Default
136+
private String containerImage = DEFAULT_IMAGE;
137+
138+
private NamespaceFiles namespaceFiles;
139+
140+
private Object inputFiles;
141+
142+
private Property<List<String>> outputFiles;
143+
144+
@Override
145+
public ScriptOutput run(RunContext runContext) throws Exception {
146+
var renderedOutputFiles = runContext.render(this.outputFiles).asList(String.class);
147+
var renderedCommands = runContext.render(this.commands).asList(String.class);
148+
var envVarsWithDefaultAuthentication = Optional.ofNullable(env).orElse(new HashMap<>());
149+
envVarsWithDefaultAuthentication.putAll(Map.of(
150+
"SNOWFLAKE_ACCOUNT", runContext.render(this.account).as(String.class).orElseThrow(),
151+
"SNOWFLAKE_USER", runContext.render(this.username).as(String.class).orElseThrow(),
152+
"SNOWFLAKE_PASSWORD", runContext.render(this.password).as(String.class).orElseThrow()
153+
));
154+
155+
return new CommandsWrapper(runContext)
156+
.withWarningOnStdErr(true)
157+
.withTaskRunner(this.taskRunner)
158+
.withContainerImage(this.containerImage)
159+
.withEnv(envVarsWithDefaultAuthentication)
160+
.withNamespaceFiles(namespaceFiles)
161+
.withInputFiles(inputFiles)
162+
.withOutputFiles(renderedOutputFiles.isEmpty() ? null : renderedOutputFiles)
163+
.withInterpreter(Property.of(List.of("/bin/sh", "-c")))
164+
.withCommands(
165+
Property.of(
166+
Stream.concat(
167+
Stream.of(
168+
"snow connection add --connection-name=default-connection --default --no-interactive"
169+
),
170+
renderedCommands.stream()
171+
).toList()
172+
)
173+
)
174+
.run();
175+
}
176+
}

‎plugin-jdbc-snowflake/src/test/java/io/kestra/plugin/jdbc/snowflake/SnowflakeCLITest.java‎

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
1+
package io.kestra.plugin.jdbc.snowflake;
2+
3+
import io.kestra.core.junit.annotations.KestraTest;
4+
import io.kestra.core.models.property.Property;
5+
import io.kestra.core.runners.RunContext;
6+
import io.kestra.core.runners.RunContextFactory;
7+
import io.kestra.core.utils.IdUtils;
8+
import io.kestra.core.utils.TestsUtils;
9+
import io.kestra.plugin.scripts.exec.scripts.models.ScriptOutput;
10+
import jakarta.inject.Inject;
11+
import org.junit.jupiter.api.Disabled;
12+
import org.junit.jupiter.api.Test;
13+
14+
import java.util.List;
15+
import java.util.Map;
16+
17+
import static org.hamcrest.MatcherAssert.assertThat;
18+
import static org.hamcrest.Matchers.is;
19+
20+
@KestraTest
21+
@Disabled("Create a Snowflake account for unit testing")
22+
public class SnowflakeCLITest {
23+
24+
@Inject
25+
private RunContextFactory runContextFactory;
26+
27+
protected String account = "";
28+
protected String username = "";
29+
protected String password = "";
30+
31+
@Test
32+
void run() throws Exception {
33+
34+
var snowflakeCLI = SnowflakeCLI.builder()
35+
.id(IdUtils.create())
36+
.type(SnowflakeCLI.class.getName())
37+
.account(Property.of(account))
38+
.username(Property.of(username))
39+
.password(Property.of(password))
40+
.commands(
41+
Property.of(List.of("snow connection test")))
42+
.build();
43+
44+
RunContext runContext = TestsUtils.mockRunContext(runContextFactory, snowflakeCLI, Map.of());
45+
46+
ScriptOutput output = snowflakeCLI.run(runContext);
47+
48+
assertThat(output.getExitCode(), is(0));
49+
}
50+
51+
}

0 commit comments

Comments
 (0)