Skip to content

Commit d24aac3

Browse files
ferenc-csakyferenc-csaky
committed
feat: Support referencing individual tables from SQRL scripts in IMPORT and EXPORT 
1 parent f90a253 commit d24aac3

22 files changed

Lines changed: 1487 additions & 57 deletions

‎sqrl-planner/src/main/java/com/datasqrl/engine/stream/flink/FlinkCalciteParser.java‎

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,62 @@
1+
/*
2+
* Copyright © 2021 DataSQRL (contact@datasqrl.com)
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
package com.datasqrl.engine.stream.flink;
17+
18+
import static com.google.common.base.Preconditions.checkArgument;
19+
20+
import jakarta.annotation.Nullable;
21+
import java.util.Optional;
22+
import java.util.function.Supplier;
23+
import lombok.AccessLevel;
24+
import lombok.NoArgsConstructor;
25+
import org.apache.calcite.sql.SqlNode;
26+
import org.apache.flink.configuration.Configuration;
27+
import org.apache.flink.table.api.internal.TableEnvironmentImpl;
28+
import org.apache.flink.table.planner.delegation.ParserImpl;
29+
import org.apache.flink.table.planner.parse.CalciteParser;
30+
31+
@NoArgsConstructor(access = AccessLevel.PRIVATE)
32+
public final class FlinkCalciteParser {
33+
34+
public static SqlNode parseSql(String sqlStmt) {
35+
return parseSql(sqlStmt, null);
36+
}
37+
38+
public static SqlNode parseSql(String sqlStmt, @Nullable TableEnvironmentImpl tEnv) {
39+
var parser = getCalciteParser(Optional.ofNullable(tEnv));
40+
41+
var sqlNodeList = parser.parseSqlList(sqlStmt);
42+
var parsed = sqlNodeList.getList();
43+
checkArgument(
44+
parsed.size() == 1,
45+
"Expected exactly 1 SQL statement but found %s. SQL: [%s]",
46+
parsed.size(),
47+
sqlStmt.length() > 500 ? sqlStmt.substring(0, 500) + "..." : sqlStmt);
48+
49+
return parsed.get(0);
50+
}
51+
52+
private static CalciteParser getCalciteParser(Optional<TableEnvironmentImpl> tEnv) {
53+
try {
54+
var tEnvImpl = tEnv.orElse(TableEnvironmentImpl.create(new Configuration()));
55+
var calciteSupplierField = ParserImpl.class.getDeclaredField("calciteParserSupplier");
56+
calciteSupplierField.setAccessible(true);
57+
return ((Supplier<CalciteParser>) calciteSupplierField.get(tEnvImpl.getParser())).get();
58+
} catch (NoSuchFieldException | IllegalAccessException e) {
59+
throw new RuntimeException(e);
60+
}
61+
}
62+
}

‎sqrl-planner/src/main/java/com/datasqrl/error/ErrorCode.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ public enum ErrorCode implements ErrorLabel {
2727
INVALID_TABLE_FUNCTION_ARGUMENTS,
2828
INVALID_SQRL_DEFINITION,
2929
INVALID_SQRL_ADD_COLUMN,
30+
INVALID_INDIVIDUAL_SCRIPT_TABLE,
3031
INVALID_NEXT_BATCH,
3132
BASETABLE_ONLY_ERROR,
3233
MISSING_DEST_TABLE,

sqrl-planner/src/main/java/com/datasqrl/loaders/SqrlDirectoryModule.java renamed to sqrl-planner/src/main/java/com/datasqrl/loaders/DirectorySqrlModule.java

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -21,13 +21,13 @@
2121
import java.util.List;
2222
import java.util.Optional;
2323

24-
public class SqrlDirectoryModule implements SqrlModule {
25-
List<NamespaceObject> nsObjects;
24+
public class DirectorySqrlModule implements SqrlModule {
2625

27-
public SqrlDirectoryModule(List<NamespaceObject> nsObjects) {
28-
if (nsObjects
29-
instanceof
30-
ArrayList) { // check for mutable lists to sort (for consistent tests and behavior)
26+
private final List<NamespaceObject> nsObjects;
27+
28+
public DirectorySqrlModule(List<NamespaceObject> nsObjects) {
29+
// check for mutable lists to sort (for consistent tests and behavior)
30+
if (nsObjects instanceof ArrayList) {
3131
nsObjects.sort(Comparator.comparing(NamespaceObject::name));
3232
}
3333
this.nsObjects = nsObjects;

‎sqrl-planner/src/main/java/com/datasqrl/loaders/FlinkTableNamespaceObject.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,5 +28,5 @@ public Name name() {
2828
return table.name();
2929
}
3030

31-
public record FlinkTable(Name name, String flinkSql, Path flinkSqlFile) {}
31+
public record FlinkTable(Name name, String sql, Path scriptPath) {}
3232
}

‎sqrl-planner/src/main/java/com/datasqrl/loaders/ModuleLoaderImpl.java‎

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
import com.datasqrl.loaders.resolver.ResourceResolver;
2727
import com.datasqrl.loaders.schema.SchemaLoader;
2828
import com.datasqrl.loaders.schema.SchemaLoaderImpl;
29+
import com.datasqrl.planner.parser.SqrlStatementParser;
2930
import com.datasqrl.serializer.Deserializer;
3031
import com.datasqrl.util.FunctionUtil;
3132
import com.fasterxml.jackson.databind.node.ObjectNode;
@@ -57,12 +58,14 @@ public class ModuleLoaderImpl implements ModuleLoader {
5758
CacheBuilder.newBuilder().maximumSize(1000).build();
5859

5960
private final ResourceResolver resourceResolver;
61+
private final SqrlStatementParser sqrlStatementParser;
6062
private final WorkspacePaths workspacePaths;
6163
private final ClasspathFunctionLoader classpathFunctionLoader;
6264
private final ErrorCollector errors;
6365

6466
private ModuleLoaderImpl withResourceResolver(ResourceResolver resourceResolver) {
65-
return new ModuleLoaderImpl(resourceResolver, workspacePaths, classpathFunctionLoader, errors);
67+
return new ModuleLoaderImpl(
68+
resourceResolver, sqrlStatementParser, workspacePaths, classpathFunctionLoader, errors);
6669
}
6770

6871
@Override
@@ -92,7 +95,7 @@ private Optional<SqrlModule> loadFunctionsFromClasspath(NamePath namePath) {
9295
if (lib.isEmpty()) {
9396
return Optional.empty();
9497
}
95-
return Optional.of(new SqrlDirectoryModule(lib));
98+
return Optional.of(new DirectorySqrlModule(lib));
9699
}
97100

98101
private Optional<SqrlModule> loadFromFileSystem(NamePath directory) {
@@ -115,7 +118,7 @@ private Optional<SqrlModule> loadFromFileSystem(NamePath directory) {
115118
.flatMap(url -> load(url, directory).stream())
116119
.collect(Collectors.toList());
117120

118-
return items.isEmpty() ? Optional.empty() : Optional.of(new SqrlDirectoryModule(items));
121+
return items.isEmpty() ? Optional.empty() : Optional.of(new DirectorySqrlModule(items));
119122
}
120123

121124
private Optional<Path> getFile(NamePath directory, Name name) {
@@ -149,7 +152,9 @@ private ScriptSqrlModule loadScript(NamePath namePath, Path path) {
149152
Files.readString(path),
150153
path,
151154
namePath,
152-
withResourceResolver(resourceResolver.getNested(path.getParent())));
155+
sqrlStatementParser,
156+
withResourceResolver(resourceResolver.getNested(path.getParent())),
157+
errors);
153158
}
154159

155160
@SneakyThrows

‎sqrl-planner/src/main/java/com/datasqrl/loaders/ModuleLoaders.java‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@
2727
import com.datasqrl.error.ErrorLocation.FileLocation;
2828
import com.datasqrl.loaders.resolver.FileResourceResolver;
2929
import com.datasqrl.loaders.resolver.ResourceResolver;
30+
import com.datasqrl.planner.parser.SqrlStatementParser;
3031
import jakarta.inject.Inject;
3132
import java.util.Set;
3233
import java.util.stream.Collectors;
@@ -50,15 +51,18 @@ public final class ModuleLoaders {
5051
public ModuleLoaders(
5152
PackageJson packageJson,
5253
ResourceResolver resourceResolver,
54+
SqrlStatementParser sqrlStatementParser,
5355
WorkspacePaths workspacePaths,
5456
ClasspathFunctionLoader classpathFunctionLoader,
5557
ErrorCollector errors) {
5658

5759
mainLoader =
58-
new ModuleLoaderImpl(resourceResolver, workspacePaths, classpathFunctionLoader, errors);
60+
new ModuleLoaderImpl(
61+
resourceResolver, sqrlStatementParser, workspacePaths, classpathFunctionLoader, errors);
5962
rootLoader =
6063
new ModuleLoaderImpl(
6164
new FileResourceResolver(workspacePaths.buildDir()),
65+
sqrlStatementParser,
6266
workspacePaths,
6367
classpathFunctionLoader,
6468
errors);

‎sqrl-planner/src/main/java/com/datasqrl/loaders/ScriptSqrlModule.java‎

Lines changed: 71 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -17,30 +17,43 @@
1717

1818
import com.datasqrl.canonicalizer.Name;
1919
import com.datasqrl.canonicalizer.NamePath;
20+
import com.datasqrl.engine.stream.flink.FlinkCalciteParser;
21+
import com.datasqrl.error.ErrorCode;
22+
import com.datasqrl.error.ErrorCollector;
23+
import com.datasqrl.loaders.FlinkTableNamespaceObject.FlinkTable;
2024
import com.datasqrl.plan.MainScript;
25+
import com.datasqrl.planner.parser.SqrlCreateTableStatement;
26+
import com.datasqrl.planner.parser.SqrlStatementParser;
2127
import java.nio.file.Path;
28+
import java.util.HashMap;
2229
import java.util.List;
30+
import java.util.Map;
2331
import java.util.Optional;
32+
import java.util.function.Supplier;
2433
import lombok.AllArgsConstructor;
2534
import lombok.Getter;
35+
import lombok.RequiredArgsConstructor;
36+
import org.apache.flink.sql.parser.ddl.table.SqlCreateTable;
2637

27-
@AllArgsConstructor
38+
@RequiredArgsConstructor
2839
public class ScriptSqrlModule implements SqrlModule {
2940

3041
private final String scriptContent;
3142
private final Path scriptPath;
3243
private final NamePath namePath;
44+
private final SqrlStatementParser sqrlStatementParser;
3345
private final ModuleLoader moduleLoader;
46+
private final ErrorCollector errors;
47+
48+
private Map<Name, Supplier<NamespaceObject>> tables = null;
3449

3550
@Override
3651
public Optional<NamespaceObject> getNamespaceObject(Name name) {
37-
throw new UnsupportedOperationException(
38-
"Cannot import an individual table from SQRL script. "
39-
+ "Use `%s;` to import entire script in a separate database or `%s.*;` to import script inline."
40-
.formatted(namePath, namePath));
52+
initTables();
53+
54+
return Optional.ofNullable(tables.get(name)).map(Supplier::get);
4155
}
4256

43-
// Planning all tables. Return a single planning ns object and add all tables
4457
@Override
4558
public List<NamespaceObject> getNamespaceObjects() {
4659
return List.of(new ScriptNamespaceObject(true));
@@ -50,6 +63,58 @@ public NamespaceObject asNamespaceObject() {
5063
return new ScriptNamespaceObject(false);
5164
}
5265

66+
/**
67+
* Initializes all {@code CREATE TABLE} statements defined in the scrip for individual table
68+
* imports and exports.
69+
*
70+
* <p>The script is parsed eagerly so table names can be resolved through the module namespace,
71+
* but validation and construction only run when that specific table is referenced.
72+
*/
73+
private void initTables() {
74+
if (tables != null) {
75+
return;
76+
}
77+
78+
tables = new HashMap<>();
79+
80+
var scriptErrors = errors.withScript(scriptPath, scriptContent);
81+
var parsedCreateTables = sqrlStatementParser.parseCreateTables(scriptContent, scriptErrors);
82+
83+
for (var parsedCreateTable : parsedCreateTables) {
84+
var createTableStmt = (SqrlCreateTableStatement) parsedCreateTable.statement().get();
85+
86+
var sql = createTableStmt.toSql();
87+
var sqlNode = FlinkCalciteParser.parseSql(sql);
88+
89+
if (!(sqlNode instanceof SqlCreateTable createTable)) {
90+
throw new IllegalStateException("Expected SqlCreateTable but got " + sqlNode.getClass());
91+
}
92+
93+
var tableName = createTable.getName().getSimple();
94+
95+
// Wrap to a supplier so validation is only triggerred when the table is referenced
96+
Supplier<NamespaceObject> nsObjectSupplier =
97+
() -> {
98+
var finalLocation =
99+
createTableStmt.mapSqlLocation(parsedCreateTable.statement().getFileLocation());
100+
101+
scriptErrors
102+
.atFile(finalLocation)
103+
.checkFatal(
104+
!createTable.getProperties().isEmpty(),
105+
ErrorCode.INVALID_INDIVIDUAL_SCRIPT_TABLE,
106+
"Referenced table '%s' is not an external table",
107+
tableName);
108+
109+
return new FlinkTableNamespaceObject(
110+
new FlinkTable(Name.system(tableName), sql, scriptPath),
111+
moduleLoader.getSchemaLoader());
112+
};
113+
114+
tables.put(Name.system(tableName), nsObjectSupplier);
115+
}
116+
}
117+
53118
@AllArgsConstructor
54119
public class ScriptNamespaceObject implements NamespaceObject {
55120

‎sqrl-planner/src/main/java/com/datasqrl/planner/SqlScriptPlanner.java‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1170,12 +1170,12 @@ public static ExternalFlinkTable fromNamespaceObject(
11701170
var flinkTable = nsObject.table();
11711171

11721172
// Parse SQL
1173-
var tableSql = flinkTable.flinkSql();
1174-
var tableError = errorCollector.withScript(flinkTable.flinkSqlFile(), tableSql);
1173+
var tableSql = flinkTable.sql();
1174+
var tableError = errorCollector.withScript(flinkTable.scriptPath(), tableSql);
11751175
tableSql = SqlScriptStatementSplitter.formatEndOfSqlFile(tableSql);
11761176

11771177
return new ExternalFlinkTable(
1178-
tableSql, nsObject.schemaLoader(), flinkTable.flinkSqlFile(), tableError);
1178+
tableSql, nsObject.schemaLoader(), flinkTable.scriptPath(), tableError);
11791179
}
11801180

11811181
public void writeExportWithSchema(FlinkTableBuilder table) {

‎sqrl-planner/src/main/java/com/datasqrl/planner/Sqrl2FlinkSQLTranslator.java‎

Lines changed: 2 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
import com.datasqrl.config.PackageJson.CompilerConfig;
2323
import com.datasqrl.config.SqrlConstants;
2424
import com.datasqrl.config.WorkspacePaths;
25+
import com.datasqrl.engine.stream.flink.FlinkCalciteParser;
2526
import com.datasqrl.engine.stream.flink.FlinkSqlNodes;
2627
import com.datasqrl.engine.stream.flink.FlinkStreamEngine;
2728
import com.datasqrl.engine.stream.flink.sql.RelToFlinkSql;
@@ -136,12 +137,10 @@
136137
import org.apache.flink.table.planner.calcite.FlinkPlannerImpl;
137138
import org.apache.flink.table.planner.calcite.FlinkRelBuilder;
138139
import org.apache.flink.table.planner.calcite.FlinkTypeFactory;
139-
import org.apache.flink.table.planner.delegation.ParserImpl;
140140
import org.apache.flink.table.planner.delegation.PlannerBase;
141141
import org.apache.flink.table.planner.expressions.RexNodeExpression;
142142
import org.apache.flink.table.planner.operations.SqlNodeConvertContext;
143143
import org.apache.flink.table.planner.operations.SqlNodeToOperationConversion;
144-
import org.apache.flink.table.planner.parse.CalciteParser;
145144
import org.apache.flink.table.planner.utils.RowLevelModificationContextUtils;
146145
import org.apache.flink.table.types.CollectionDataType;
147146
import org.apache.flink.table.types.DataType;
@@ -249,24 +248,7 @@ public SqrlRexUtil getRexUtil() {
249248
}
250249

251250
public SqlNode parseSQL(String sqlStatement) {
252-
CalciteParser parser;
253-
try {
254-
// TODO: This is a hack - is there a better way to get the calcite parser?
255-
var calciteSupplierField = ParserImpl.class.getDeclaredField("calciteParserSupplier");
256-
calciteSupplierField.setAccessible(true);
257-
parser = ((Supplier<CalciteParser>) calciteSupplierField.get(tEnv.getParser())).get();
258-
} catch (NoSuchFieldException | IllegalAccessException e) {
259-
throw new RuntimeException(e);
260-
}
261-
var sqlNodeList = parser.parseSqlList(sqlStatement);
262-
var parsed = sqlNodeList.getList();
263-
checkArgument(
264-
parsed.size() == 1,
265-
"Expected exactly 1 SQL statement but found %s. SQL: [%s]",
266-
parsed.size(),
267-
sqlStatement.length() > 500 ? sqlStatement.substring(0, 500) + "..." : sqlStatement);
268-
269-
return parsed.get(0);
251+
return FlinkCalciteParser.parseSql(sqlStatement, tEnv);
270252
}
271253

272254
/**

‎sqrl-planner/src/main/java/com/datasqrl/planner/parser/SqrlStatementParser.java‎

Lines changed: 19 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@
2727
import java.util.ArrayList;
2828
import java.util.List;
2929
import java.util.Map;
30+
import java.util.function.Predicate;
3031
import java.util.regex.Matcher;
3132
import java.util.regex.Pattern;
3233
import org.apache.commons.lang3.tuple.Pair;
@@ -122,21 +123,33 @@ public class SqrlStatementParser {
122123
* @return
123124
*/
124125
public List<ParsedStatement> parseScript(String script, ErrorCollector scriptErrors) {
125-
List<ParsedStatement> sqlStatements = new ArrayList<>();
126+
return parseScript(script, scriptErrors, stmt -> true);
127+
}
128+
129+
public List<ParsedStatement> parseCreateTables(String script, ErrorCollector scriptErrors) {
130+
return parseScript(script, scriptErrors, stmt -> stmt instanceof SqrlCreateTableStatement);
131+
}
132+
133+
private List<ParsedStatement> parseScript(
134+
String script, ErrorCollector scriptErrors, Predicate<SQLStatement> pred) {
135+
var sqlStatements = new ArrayList<ParsedStatement>();
126136
var localErrors = scriptErrors;
127137
try {
128138
var statements = SqlScriptStatementSplitter.splitStatements(script);
129139
for (ParsedObject<String> statement : statements) {
130140
localErrors = scriptErrors.atFile(statement.getFileLocation());
131-
sqlStatements.add(
132-
new ParsedStatement(
133-
new ParsedObject<>(
134-
parseStatement(statement.getObject()), statement.getFileLocation()),
135-
statement.get()));
141+
142+
var parsedStmt = parseStatement(statement.getObject());
143+
if (pred.test(parsedStmt)) {
144+
sqlStatements.add(
145+
new ParsedStatement(
146+
new ParsedObject<>(parsedStmt, statement.getFileLocation()), statement.get()));
147+
}
136148
}
137149
} catch (StatementParserException e) {
138150
throw localErrors.handle(e);
139151
}
152+
140153
return sqlStatements;
141154
}
142155

0 commit comments

Comments
 (0)