Skip to content

Commit

Permalink
Make Snowflake QUERY_ACCELERATION_HISTORY optional (#353)
Browse files Browse the repository at this point in the history
The extraction of Snowflake `QUERY_ACCELERATION_HISTORY` table is sometimes not
possible, therefore I'm changing the TaskCategory from `REQUIRED` to
`OPTIONAL`, so that the extraction doesn't fail in case this table is
not available.

Additionally, cleanup the `TeradataJdbcSelectTask` class, because
otherwise the `TeradataJdbcSelectTask` class would contain a duplicated
code between it and the `JdbcSelectTask` class.
  • Loading branch information
dawidxc authored Feb 8, 2024
1 parent 3a6be98 commit a3ed6f9
Show file tree
Hide file tree
Showing 3 changed files with 31 additions and 11 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@
import com.google.edwmigration.dumper.application.dumper.task.FormatTask;
import com.google.edwmigration.dumper.application.dumper.task.JdbcSelectTask;
import com.google.edwmigration.dumper.application.dumper.task.Task;
import com.google.edwmigration.dumper.application.dumper.task.TaskCategory;
import com.google.edwmigration.dumper.plugin.ext.jdk.annotation.Description;
import com.google.edwmigration.dumper.plugin.lib.dumper.spi.SnowflakeLogsDumpFormat;
import com.google.errorprone.annotations.ForOverride;
Expand Down Expand Up @@ -135,11 +136,22 @@ private static class TaskDescription {
private final String unformattedQuery;
private final Class<? extends Enum<?>> headerClass;

private final TaskCategory taskCategory;

private TaskDescription(
String zipPrefix, String unformattedQuery, Class<? extends Enum<?>> headerClass) {
this.unformattedQuery = unformattedQuery;
this(zipPrefix, unformattedQuery, headerClass, TaskCategory.REQUIRED);
}

private TaskDescription(
String zipPrefix,
String unformattedQuery,
Class<? extends Enum<?>> headerClass,
TaskCategory taskCategory) {
this.zipPrefix = zipPrefix;
this.unformattedQuery = unformattedQuery;
this.headerClass = headerClass;
this.taskCategory = taskCategory;
}
}

Expand Down Expand Up @@ -381,7 +393,7 @@ private static void addJdbcTask(
task.zipPrefix
+ DateTimeFormatter.ISO_OFFSET_DATE_TIME.format(interval.getStartUTC())
+ ".csv";
out.add(new JdbcSelectTask(file, query).withHeaderClass(task.headerClass));
out.add(new JdbcSelectTask(file, query, task.taskCategory).withHeaderClass(task.headerClass));
}

private String getOverrideableQuery(
Expand Down Expand Up @@ -495,7 +507,8 @@ private List<TaskDescription> createTimeSeriesTasks(ConnectorArguments arguments
parseColumnsFromHeader(QueryAccelerationHistoryFormat.Header.class),
"QUERY_ACCELERATION_HISTORY"),
"END_TIME"),
QueryAccelerationHistoryFormat.Header.class),
QueryAccelerationHistoryFormat.Header.class,
TaskCategory.OPTIONAL),
new TaskDescription(
ReplicationGroupUsageHistoryFormat.ZIP_ENTRY_PREFIX,
getOverrideableQuery(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -121,17 +121,10 @@ public String getDescription() {
protected static class TeradataJdbcSelectTask extends JdbcSelectTask {

private final String sqlCount;
private final TaskCategory category;

public TeradataJdbcSelectTask(String targetPath, TaskCategory category, String sqlTemplate) {
super(targetPath, String.format(sqlTemplate, "*"));
super(targetPath, String.format(sqlTemplate, "*"), Preconditions.checkNotNull(category));
this.sqlCount = sqlTemplate.contains("%s") ? String.format(sqlTemplate, "count(*)") : null;
this.category = Preconditions.checkNotNull(category);
}

@Override
public TaskCategory getCategory() {
return category;
}

// This works to execute the count SQL on the same connection as the data SQL,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -33,9 +33,23 @@ public class JdbcSelectTask extends AbstractJdbcTask<Summary> {

@Nonnull private final String sql;

@Nonnull private final TaskCategory taskCategory;

public JdbcSelectTask(@Nonnull String targetPath, @Nonnull String sql) {
this(targetPath, sql, TaskCategory.REQUIRED);
}

public JdbcSelectTask(
@Nonnull String targetPath, @Nonnull String sql, TaskCategory taskCategory) {
super(targetPath);
this.sql = sql;
this.taskCategory = taskCategory;
}

@Override
@Nonnull
public TaskCategory getCategory() {
return taskCategory;
}

@Nonnull
Expand Down

0 comments on commit a3ed6f9

Please sign in to comment.