Support internal binding and extraction for SQL Server UPDATE table variable targets

RELEASE-NOTES.md

@@ -77,6 +77,7 @@
 1. Encrypt: Support SqlServer update statement for Specifying a table alias as the target object when use encrypt feature - [#38733](https://github.com/apache/shardingsphere/pull/38733)
 1. Encrypt: Support SqlServer update statement for Specifying a view as the target object when use encrypt feature - [#38896](https://github.com/apache/shardingsphere/pull/38896)
 1. Encrypt: Support SqlServer for Using the UPDATE statement with information from another table  when use encrypt feature - [#38926](https://github.com/apache/shardingsphere/pull/38926)
+1. SQL Binder: Support internal binding and extraction for SQL Server UPDATE table variable targets - [#39093](https://github.com/apache/shardingsphere/pull/39093)
 1. Sharding: Fix HASH_MOD routing mismatch for same negative numeric values across numeric Java types with compatibility switch `normalize-numeric-int-range` - [#38327](https://github.com/apache/shardingsphere/pull/38327)
 1. Proxy Native: Support building Proxy Native via GraalVM CE for JDK 25 - [#38682](https://github.com/apache/shardingsphere/pull/38682)
 1. SQL Parser: Support SQLServer table variable declaration parse - [#38904](https://github.com/apache/shardingsphere/pull/38904)

database/connector/core/src/main/java/org/apache/shardingsphere/database/connector/core/metadata/database/metadata/DialectDatabaseMetaData.java

@@ -217,4 +217,13 @@ default DialectProtocolVersionOption getProtocolVersionOption() {
     default DialectFunctionOption getFunctionOption() {
         return new DefaultFunctionOption();
     }
+
+    /**
+     * Get variable table name prefix.
+     *
+     * @return variable table name prefix
+     */
+    default Optional<String> getVariableTableNamePrefix() {
+        return Optional.empty();
+    }
 }

database/connector/dialect/sqlserver/src/main/java/org/apache/shardingsphere/database/connector/sql92/sqlserver/metadata/database/SQLServerDatabaseMetaData.java

@@ -88,6 +88,11 @@ public Optional<DialectAlterTableOption> getAlterTableOption() {
         return Optional.of(new DialectAlterTableOption(false, false, false, new DialectAddColumnOption("ADD", "")));
     }
 
+    @Override
+    public Optional<String> getVariableTableNamePrefix() {
+        return Optional.of("@");
+    }
+
     @Override
     public String getDatabaseType() {
         return "SQLServer";

database/connector/dialect/sqlserver/src/test/java/org/apache/shardingsphere/database/connector/sql92/sqlserver/metadata/database/SQLServerDatabaseMetaDataTest.java

@@ -114,4 +114,9 @@ void assertGetAlterTableOption() {
     void assertGetFunctionOption() {
         assertThat(dialectDatabaseMetaData.getFunctionOption(), isA(SQLServerFunctionOption.class));
     }
+
+    @Test
+    void assertGetVariableTableNamePrefix() {
+        assertThat(dialectDatabaseMetaData.getVariableTableNamePrefix(), is(Optional.of("@")));
+    }
 }

infra/binder/core/pom.xml

@@ -62,5 +62,11 @@
             <version>${project.version}</version>
             <scope>test</scope>
         </dependency>
+        <dependency>
+            <groupId>org.apache.shardingsphere</groupId>
+            <artifactId>shardingsphere-database-connector-sqlserver</artifactId>
+            <version>${project.version}</version>
+            <scope>test</scope>
+        </dependency>
     </dependencies>
 </project>

infra/binder/core/src/main/java/org/apache/shardingsphere/infra/binder/engine/segment/dml/expression/type/ColumnSegmentBinder.java

@@ -323,20 +323,24 @@ private static boolean isSkipColumnBind(final Collection<TableSegmentBinderConte
                 return true;
             }
             if (each instanceof SimpleTableSegmentBinderContext) {
-                return ((SimpleTableSegmentBinderContext) each).isContainsDBLink();
+                return isMetadataUnavailable((SimpleTableSegmentBinderContext) each);
             }
         }
         for (TableSegmentBinderContext each : outerBinderContexts) {
             if (each instanceof FunctionTableSegmentBinderContext) {
                 return true;
             }
             if (each instanceof SimpleTableSegmentBinderContext) {
-                return ((SimpleTableSegmentBinderContext) each).isContainsDBLink();
+                return isMetadataUnavailable((SimpleTableSegmentBinderContext) each);
             }
         }
         return false;
     }
 
+    private static boolean isMetadataUnavailable(final SimpleTableSegmentBinderContext tableBinderContext) {
+        return tableBinderContext.isContainsDBLink() || tableBinderContext.isContainsTableVariable();
+    }
+
     /**
      * Bind using column segment.
      *

infra/binder/core/src/main/java/org/apache/shardingsphere/infra/binder/engine/segment/dml/from/context/type/SimpleTableSegmentBinderContext.java

@@ -54,6 +54,8 @@ public final class SimpleTableSegmentBinderContext implements TableSegmentBinder
 
     private boolean containsDBLink;
 
+    private boolean containsTableVariable;
+
     public SimpleTableSegmentBinderContext(final Collection<ProjectionSegment> projectionSegments, final TableSourceType tableSourceType) {
         columnLabelProjectionSegments = new CaseInsensitiveMap<>(projectionSegments.size(), 1F);
         projectionSegments.forEach(each -> putColumnLabelProjectionSegments(each, columnLabelProjectionSegments));

infra/binder/core/src/main/java/org/apache/shardingsphere/infra/binder/engine/segment/dml/from/type/SimpleTableSegmentBinder.java

@@ -318,6 +318,9 @@ private static void checkTableExists(final SQLStatementBinderContext binderConte
         if (segment.getDbLink().isPresent()) {
             return;
         }
+        if (isVariableTable(binderContext, segment)) {
+            return;
+        }
         if (binderContext.getExternalTableBinderContexts().containsKey(CaseInsensitiveString.of(tableNameValue))) {
             return;
         }
@@ -327,6 +330,28 @@ private static void checkTableExists(final SQLStatementBinderContext binderConte
         ShardingSpherePreconditions.checkState(null != schema && schema.containsTable(tableName), () -> new TableNotFoundException(tableNameValue));
     }
 
+    private static boolean isVariableTable(final SQLStatementBinderContext binderContext, final SimpleTableSegment segment) {
+        if (!(binderContext.getSqlStatement() instanceof UpdateStatement)) {
+            return false;
+        }
+        if (segment.getOwner().isPresent()) {
+            return false;
+        }
+        IdentifierValue tableName = segment.getTableName().getIdentifier();
+        if (QuoteCharacter.NONE != tableName.getQuoteCharacter()) {
+            return false;
+        }
+        DialectDatabaseMetaData dialectDatabaseMetaData = new DatabaseTypeRegistry(binderContext.getSqlStatement().getDatabaseType()).getDialectDatabaseMetaData();
+        return dialectDatabaseMetaData.getVariableTableNamePrefix().filter(tableName.getValue()::startsWith).isPresent();
+    }
+
+    private static boolean isUpdateTargetTableVariable(final SQLStatementBinderContext binderContext, final SimpleTableSegment segment) {
+        if (!isVariableTable(binderContext, segment)) {
+            return false;
+        }
+        return ((UpdateStatement) binderContext.getSqlStatement()).getTable() == segment;
+    }
+
     private static boolean isCreateTable(final SimpleTableSegment simpleTableSegment, final String tableName) {
         return simpleTableSegment.getTableName().getIdentifier().getValue().equalsIgnoreCase(tableName);
     }
@@ -379,6 +404,11 @@ private static void checkTableMetadata(final SQLStatementBinderContext binderCon
     private static Optional<SimpleTableSegmentBinderContext> createSimpleTableBinderContext(final SimpleTableSegment segment, final ShardingSphereSchema schema, final IdentifierValue databaseName,
                                                                                             final IdentifierValue schemaName, final SQLStatementBinderContext binderContext) {
         IdentifierValue tableName = segment.getTableName().getIdentifier();
+        if (isUpdateTargetTableVariable(binderContext, segment)) {
+            SimpleTableSegmentBinderContext result = new SimpleTableSegmentBinderContext(Collections.emptyList(), TableSourceType.TEMPORARY_TABLE);
+            result.setContainsTableVariable(true);
+            return Optional.of(result);
+        }
         Optional<SimpleTableSegmentBinderContext> externalTableBinderContext = createExternalTableBinderContext(segment, tableName, binderContext);
         if (externalTableBinderContext.isPresent()) {
             return externalTableBinderContext;
@@ -392,6 +422,9 @@ private static Optional<SimpleTableSegmentBinderContext> createSimpleTableBinder
         }
         SimpleTableSegmentBinderContext result = new SimpleTableSegmentBinderContext(Collections.emptyList(), TableSourceType.TEMPORARY_TABLE);
         segment.getDbLink().ifPresent(optional -> result.setContainsDBLink(true));
+        if (isVariableTable(binderContext, segment)) {
+            result.setContainsTableVariable(true);
+        }
         return Optional.of(result);
     }
 

infra/binder/core/src/test/java/org/apache/shardingsphere/infra/binder/context/statement/type/dml/UpdateStatementContextTest.java

@@ -17,6 +17,7 @@
 
 package org.apache.shardingsphere.infra.binder.context.statement.type.dml;
 
+import org.apache.shardingsphere.database.connector.core.metadata.database.enums.QuoteCharacter;
 import org.apache.shardingsphere.database.connector.core.type.DatabaseType;
 import org.apache.shardingsphere.infra.spi.type.typed.TypedSPILoader;
 import org.apache.shardingsphere.sql.parser.statement.core.segment.dml.assignment.SetAssignmentSegment;
@@ -52,6 +53,8 @@ class UpdateStatementContextTest {
 
     private final DatabaseType databaseType = TypedSPILoader.getService(DatabaseType.class, "FIXTURE");
 
+    private final DatabaseType sqlServerDatabaseType = TypedSPILoader.getService(DatabaseType.class, "SQLServer");
+
     @Mock
     private WhereSegment whereSegment;
 
@@ -95,6 +98,35 @@ void assertGetTableNamesWithSQLServerUpdateAliasTargetExcludesAlias() {
         assertFalse(actual.getTablesContext().getTableNames().contains("sr"));
     }
 
+    @Test
+    void assertGetTableNamesWithSQLServerUpdateTableVariableTargetExcludesVariableTable() {
+        SimpleTableSegment fromTable = new SimpleTableSegment(new TableNameSegment(0, 0, new IdentifierValue("Employee")));
+        fromTable.setOwner(new OwnerSegment(0, 0, new IdentifierValue("HumanResources")));
+        UpdateStatement updateStatement = UpdateStatement.builder()
+                .databaseType(sqlServerDatabaseType)
+                .table(new SimpleTableSegment(new TableNameSegment(0, 0, new IdentifierValue("@MyTableVar"))))
+                .from(fromTable)
+                .setAssignment(new SetAssignmentSegment(0, 0, Collections.emptyList()))
+                .build();
+        UpdateStatementContext actual = new UpdateStatementContext(updateStatement);
+        assertThat(actual.getTablesContext().getTableNames(), is(Collections.singleton("Employee")));
+        assertFalse(actual.getTablesContext().getTableNames().contains("@MyTableVar"));
+        assertThat(((SimpleTableSegment) actual.getSqlStatement().getTable()).getTableName().getIdentifier().getValue(), is("@MyTableVar"));
+    }
+
+    @Test
+    void assertGetTableNamesWithSQLServerBracketDelimitedAtSignTargetIncludesPhysicalTable() {
+        SimpleTableSegment fromTable = new SimpleTableSegment(new TableNameSegment(0, 0, new IdentifierValue("Employee")));
+        UpdateStatement updateStatement = UpdateStatement.builder()
+                .databaseType(sqlServerDatabaseType)
+                .table(new SimpleTableSegment(new TableNameSegment(0, 0, new IdentifierValue("@MyTable", QuoteCharacter.BRACKETS))))
+                .from(fromTable)
+                .setAssignment(new SetAssignmentSegment(0, 0, Collections.emptyList()))
+                .build();
+        UpdateStatementContext actual = new UpdateStatementContext(updateStatement);
+        assertThat(actual.getTablesContext().getTableNames(), is(new HashSet<>(Arrays.asList("@MyTable", "Employee"))));
+    }
+
     @Test
     void assertGetTableNamesWithPostgreSQLUpdateFromClauseIncludesTargetTable() {
         SimpleTableSegment targetTable = new SimpleTableSegment(new TableNameSegment(7, 18, new IdentifierValue("ScrapReason")));

infra/binder/core/src/test/java/org/apache/shardingsphere/infra/binder/engine/segment/dml/expression/type/ColumnSegmentBinderTest.java

@@ -25,6 +25,7 @@
 import org.apache.shardingsphere.infra.binder.engine.segment.dml.from.context.TableSegmentBinderContext;
 import org.apache.shardingsphere.infra.binder.engine.segment.dml.from.context.type.SimpleTableSegmentBinderContext;
 import org.apache.shardingsphere.infra.binder.engine.statement.SQLStatementBinderContext;
+import org.apache.shardingsphere.infra.exception.kernel.metadata.ColumnNotFoundException;
 import org.apache.shardingsphere.infra.exception.kernel.syntax.AmbiguousColumnException;
 import org.apache.shardingsphere.infra.hint.HintValueContext;
 import org.apache.shardingsphere.infra.metadata.ShardingSphereMetaData;
@@ -238,4 +239,25 @@ void assertBindExcludedColumnInSetAssignment() {
         assertTrue(actual.getOwner().isPresent());
         assertThat(actual.getOwner().get().getIdentifier().getValue(), is("EXCLUDED"));
     }
+
+    @Test
+    void assertBindWithTableVariableSkipsColumnExistenceCheck() {
+        Multimap<CaseInsensitiveString, TableSegmentBinderContext> tableBinderContexts = LinkedHashMultimap.create();
+        SimpleTableSegmentBinderContext tableVariableBinderContext = new SimpleTableSegmentBinderContext(Collections.emptyList(), TableSourceType.TEMPORARY_TABLE);
+        tableVariableBinderContext.setContainsTableVariable(true);
+        tableBinderContexts.put(CaseInsensitiveString.of("@MyTableVar"), tableVariableBinderContext);
+        ColumnSegment columnSegment = new ColumnSegment(0, 0, new IdentifierValue("NewVacationHours"));
+        ColumnSegment actual = ColumnSegmentBinder.bind(columnSegment, SegmentType.SET_ASSIGNMENT, createBinderContext(), tableBinderContexts, LinkedHashMultimap.create());
+        assertThat(actual.getColumnBoundInfo().getOriginalColumn().getValue(), is("NewVacationHours"));
+    }
+
+    @Test
+    void assertBindUnknownColumnWithoutTableVariable() {
+        Multimap<CaseInsensitiveString, TableSegmentBinderContext> tableBinderContexts = LinkedHashMultimap.create();
+        tableBinderContexts.put(CaseInsensitiveString.of("t_order"),
+                new SimpleTableSegmentBinderContext(Collections.emptyList(), TableSourceType.PHYSICAL_TABLE));
+        ColumnSegment columnSegment = new ColumnSegment(0, 0, new IdentifierValue("NewVacationHours"));
+        assertThrows(ColumnNotFoundException.class,
+                () -> ColumnSegmentBinder.bind(columnSegment, SegmentType.SET_ASSIGNMENT, createBinderContext(), tableBinderContexts, LinkedHashMultimap.create()));
+    }
 }