Pass ConnectorSession to PlanChecker.validate

rschlussel · rschlussel · commit 377e3fe827e7 · 2025-02-12T12:44:32.000-05:00
This allows the planchecker to pass the ConnectorSession to any
connector calls that it makes.
diff --git a/presto-main/src/main/java/com/facebook/presto/sql/planner/sanity/PlanChecker.java b/presto-main/src/main/java/com/facebook/presto/sql/planner/sanity/PlanChecker.java
@@ -87,7 +87,7 @@ public void validateFinalPlan(PlanNode planNode, Session session, Metadata metad
         checkers.get(Stage.FINAL).forEach(checker -> checker.validate(planNode, session, metadata, warningCollector));
         for (PlanCheckerProvider provider : planCheckerProviderManager.getPlanCheckerProviders()) {
             for (com.facebook.presto.spi.plan.PlanChecker checker : provider.getFinalPlanCheckers()) {
-                checker.validate(planNode, warningCollector);
+                checker.validate(planNode, warningCollector, session.toConnectorSession());
             }
         }
     }
@@ -97,7 +97,7 @@ public void validateIntermediatePlan(PlanNode planNode, Session session, Metadat
         checkers.get(Stage.INTERMEDIATE).forEach(checker -> checker.validate(planNode, session, metadata, warningCollector));
         for (PlanCheckerProvider provider : planCheckerProviderManager.getPlanCheckerProviders()) {
             for (com.facebook.presto.spi.plan.PlanChecker checker : provider.getIntermediatePlanCheckers()) {
-                checker.validate(planNode, warningCollector);
+                checker.validate(planNode, warningCollector, session.toConnectorSession());
             }
         }
     }
@@ -107,15 +107,18 @@ public void validatePlanFragment(PlanFragment planFragment, Session session, Met
         checkers.get(Stage.FRAGMENT).forEach(checker -> checker.validateFragment(planFragment, session, metadata, warningCollector));
         for (PlanCheckerProvider provider : planCheckerProviderManager.getPlanCheckerProviders()) {
             for (com.facebook.presto.spi.plan.PlanChecker checker : provider.getFragmentPlanCheckers()) {
-                checker.validateFragment(new SimplePlanFragment(
-                        planFragment.getId(),
-                        planFragment.getRoot(),
-                        planFragment.getVariables(),
-                        planFragment.getPartitioning(),
-                        planFragment.getTableScanSchedulingOrder(),
-                        planFragment.getPartitioningScheme(),
-                        planFragment.getStageExecutionDescriptor(),
-                        planFragment.isOutputTableWriterFragment()), warningCollector);
+                checker.validateFragment(
+                        new SimplePlanFragment(
+                                planFragment.getId(),
+                                planFragment.getRoot(),
+                                planFragment.getVariables(),
+                                planFragment.getPartitioning(),
+                                planFragment.getTableScanSchedulingOrder(),
+                                planFragment.getPartitioningScheme(),
+                                planFragment.getStageExecutionDescriptor(),
+                                planFragment.isOutputTableWriterFragment()),
+                        warningCollector,
+                        session.toConnectorSession());
             }
         }
     }
diff --git a/presto-native-sidecar-plugin/src/main/java/com/facebook/presto/sidecar/nativechecker/NativePlanChecker.java b/presto-native-sidecar-plugin/src/main/java/com/facebook/presto/sidecar/nativechecker/NativePlanChecker.java
@@ -16,6 +16,7 @@
 import com.facebook.airlift.json.JsonCodec;
 import com.facebook.airlift.log.Logger;
 import com.facebook.presto.spi.ConnectorId;
+import com.facebook.presto.spi.ConnectorSession;
 import com.facebook.presto.spi.NodeManager;
 import com.facebook.presto.spi.PrestoException;
 import com.facebook.presto.spi.TableHandle;
@@ -61,13 +62,13 @@ public NativePlanChecker(NodeManager nodeManager, JsonCodec<SimplePlanFragment>
     }
 
     @Override
-    public void validate(PlanNode planNode, WarningCollector warningCollector)
+    public void validate(PlanNode planNode, WarningCollector warningCollector, ConnectorSession session)
     {
         // NO-OP, only validating fragments
     }
 
     @Override
-    public void validateFragment(SimplePlanFragment planFragment, WarningCollector warningCollector)
+    public void validateFragment(SimplePlanFragment planFragment, WarningCollector warningCollector, ConnectorSession session)
     {
         if (!planFragment.getPartitioning().isCoordinatorOnly() && !isInternalSystemConnector(planFragment.getRoot())) {
             runValidation(planFragment);
diff --git a/presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar/TestPlanCheckerProvider.java b/presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar/TestPlanCheckerProvider.java
@@ -93,15 +93,15 @@ public void testNativePlanMockValidate()
             PlanConversionResponse responseOk = new PlanConversionResponse(ImmutableList.of());
             String responseOkString = PLAN_CONVERSION_RESPONSE_JSON_CODEC.toJson(responseOk);
             server.enqueue(new MockResponse().setBody(responseOkString));
-            checker.validateFragment(fragment, null);
+            checker.validateFragment(fragment, null, null);
 
             String errorMessage = "native conversion error";
             ErrorCode errorCode = StandardErrorCode.NOT_SUPPORTED.toErrorCode();
             PlanConversionResponse responseError = new PlanConversionResponse(ImmutableList.of(new PlanConversionFailureInfo("MockError", errorMessage, null, ImmutableList.of(), ImmutableList.of(), errorCode)));
             String responseErrorString = PLAN_CONVERSION_RESPONSE_JSON_CODEC.toJson(responseError);
             server.enqueue(new MockResponse().setResponseCode(500).setBody(responseErrorString));
             PrestoException error = expectThrows(PrestoException.class,
-                    () -> checker.validateFragment(fragment, null));
+                    () -> checker.validateFragment(fragment, null, null));
             assertEquals(error.getErrorCode(), errorCode);
             assertTrue(error.getMessage().contains(errorMessage));
         }
diff --git a/presto-spi/src/main/java/com/facebook/presto/spi/plan/PlanChecker.java b/presto-spi/src/main/java/com/facebook/presto/spi/plan/PlanChecker.java
@@ -14,14 +14,15 @@
 
 package com.facebook.presto.spi.plan;
 
+import com.facebook.presto.spi.ConnectorSession;
 import com.facebook.presto.spi.WarningCollector;
 
 public interface PlanChecker
 {
-    void validate(PlanNode planNode, WarningCollector warningCollector);
+    void validate(PlanNode planNode, WarningCollector warningCollector, ConnectorSession session);
 
-    default void validateFragment(SimplePlanFragment planFragment, WarningCollector warningCollector)
+    default void validateFragment(SimplePlanFragment planFragment, WarningCollector warningCollector, ConnectorSession session)
     {
-        validate(planFragment.getRoot(), warningCollector);
+        validate(planFragment.getRoot(), warningCollector, session);
     }
 }
diff --git a/presto-tests/src/test/java/com/facebook/presto/execution/TriggerFailurePlanChecker.java b/presto-tests/src/test/java/com/facebook/presto/execution/TriggerFailurePlanChecker.java
@@ -13,6 +13,7 @@
  */
 package com.facebook.presto.execution;
 
+import com.facebook.presto.spi.ConnectorSession;
 import com.facebook.presto.spi.PrestoException;
 import com.facebook.presto.spi.StandardErrorCode;
 import com.facebook.presto.spi.WarningCollector;
@@ -36,7 +37,7 @@ public TriggerFailurePlanChecker(AtomicBoolean triggerValidationFailure)
     }
 
     @Override
-    public void validate(PlanNode planNode, WarningCollector warningCollector)
+    public void validate(PlanNode planNode, WarningCollector warningCollector, ConnectorSession session)
     {
         if (triggerValidationFailure.get()) {
             throw new PrestoException(FAILURE_ERROR_CODE, FAILURE_MESSAGE);

Original file line number	Diff line number	Diff line change
`@@ -87,7 +87,7 @@ public void validateFinalPlan(PlanNode planNode, Session session, Metadata metad`
`87`	`87`	`checkers.get(Stage.FINAL).forEach(checker -> checker.validate(planNode, session, metadata, warningCollector));`
`88`	`88`	`for (PlanCheckerProvider provider : planCheckerProviderManager.getPlanCheckerProviders()) {`
`89`	`89`	`for (com.facebook.presto.spi.plan.PlanChecker checker : provider.getFinalPlanCheckers()) {`
`90`		`- checker.validate(planNode, warningCollector);`
	`90`	`+ checker.validate(planNode, warningCollector, session.toConnectorSession());`
`91`	`91`	`}`
`92`	`92`	`}`
`93`	`93`	`}`
`@@ -97,7 +97,7 @@ public void validateIntermediatePlan(PlanNode planNode, Session session, Metadat`
`97`	`97`	`checkers.get(Stage.INTERMEDIATE).forEach(checker -> checker.validate(planNode, session, metadata, warningCollector));`
`98`	`98`	`for (PlanCheckerProvider provider : planCheckerProviderManager.getPlanCheckerProviders()) {`
`99`	`99`	`for (com.facebook.presto.spi.plan.PlanChecker checker : provider.getIntermediatePlanCheckers()) {`
`100`		`- checker.validate(planNode, warningCollector);`
	`100`	`+ checker.validate(planNode, warningCollector, session.toConnectorSession());`
`101`	`101`	`}`
`102`	`102`	`}`
`103`	`103`	`}`
`@@ -107,15 +107,18 @@ public void validatePlanFragment(PlanFragment planFragment, Session session, Met`
`107`	`107`	`checkers.get(Stage.FRAGMENT).forEach(checker -> checker.validateFragment(planFragment, session, metadata, warningCollector));`
`108`	`108`	`for (PlanCheckerProvider provider : planCheckerProviderManager.getPlanCheckerProviders()) {`
`109`	`109`	`for (com.facebook.presto.spi.plan.PlanChecker checker : provider.getFragmentPlanCheckers()) {`
`110`		`- checker.validateFragment(new SimplePlanFragment(`
`111`		`- planFragment.getId(),`
`112`		`- planFragment.getRoot(),`
`113`		`- planFragment.getVariables(),`
`114`		`- planFragment.getPartitioning(),`
`115`		`- planFragment.getTableScanSchedulingOrder(),`
`116`		`- planFragment.getPartitioningScheme(),`
`117`		`- planFragment.getStageExecutionDescriptor(),`
`118`		`- planFragment.isOutputTableWriterFragment()), warningCollector);`
	`110`	`+ checker.validateFragment(`
	`111`	`+ new SimplePlanFragment(`
	`112`	`+ planFragment.getId(),`
	`113`	`+ planFragment.getRoot(),`
	`114`	`+ planFragment.getVariables(),`
	`115`	`+ planFragment.getPartitioning(),`
	`116`	`+ planFragment.getTableScanSchedulingOrder(),`
	`117`	`+ planFragment.getPartitioningScheme(),`
	`118`	`+ planFragment.getStageExecutionDescriptor(),`
	`119`	`+ planFragment.isOutputTableWriterFragment()),`
	`120`	`+ warningCollector,`
	`121`	`+ session.toConnectorSession());`
`119`	`122`	`}`
`120`	`123`	`}`
`121`	`124`	`}`
Original file line number	Diff line number	Diff line change
`@@ -16,6 +16,7 @@`
`16`	`16`	`import com.facebook.airlift.json.JsonCodec;`
`17`	`17`	`import com.facebook.airlift.log.Logger;`
`18`	`18`	`import com.facebook.presto.spi.ConnectorId;`
	`19`	`+import com.facebook.presto.spi.ConnectorSession;`
`19`	`20`	`import com.facebook.presto.spi.NodeManager;`
`20`	`21`	`import com.facebook.presto.spi.PrestoException;`
`21`	`22`	`import com.facebook.presto.spi.TableHandle;`
`@@ -61,13 +62,13 @@ public NativePlanChecker(NodeManager nodeManager, JsonCodec<SimplePlanFragment>`
`61`	`62`	`}`
`62`	`63`
`63`	`64`	`@Override`
`64`		`- public void validate(PlanNode planNode, WarningCollector warningCollector)`
	`65`	`+ public void validate(PlanNode planNode, WarningCollector warningCollector, ConnectorSession session)`
`65`	`66`	`{`
`66`	`67`	`// NO-OP, only validating fragments`
`67`	`68`	`}`
`68`	`69`
`69`	`70`	`@Override`
`70`		`- public void validateFragment(SimplePlanFragment planFragment, WarningCollector warningCollector)`
	`71`	`+ public void validateFragment(SimplePlanFragment planFragment, WarningCollector warningCollector, ConnectorSession session)`
`71`	`72`	`{`
`72`	`73`	`if (!planFragment.getPartitioning().isCoordinatorOnly() && !isInternalSystemConnector(planFragment.getRoot())) {`
`73`	`74`	`runValidation(planFragment);`
Original file line number	Diff line number	Diff line change
`@@ -14,14 +14,15 @@`
`14`	`14`
`15`	`15`	`package com.facebook.presto.spi.plan;`
`16`	`16`
	`17`	`+import com.facebook.presto.spi.ConnectorSession;`
`17`	`18`	`import com.facebook.presto.spi.WarningCollector;`
`18`	`19`
`19`	`20`	`public interface PlanChecker`
`20`	`21`	`{`
`21`		`- void validate(PlanNode planNode, WarningCollector warningCollector);`
	`22`	`+ void validate(PlanNode planNode, WarningCollector warningCollector, ConnectorSession session);`
`22`	`23`
`23`		`- default void validateFragment(SimplePlanFragment planFragment, WarningCollector warningCollector)`
	`24`	`+ default void validateFragment(SimplePlanFragment planFragment, WarningCollector warningCollector, ConnectorSession session)`
`24`	`25`	`{`
`25`		`- validate(planFragment.getRoot(), warningCollector);`
	`26`	`+ validate(planFragment.getRoot(), warningCollector, session);`
`26`	`27`	`}`
`27`	`28`	`}`
Original file line number	Diff line number	Diff line change
`@@ -13,6 +13,7 @@`
`13`	`13`	`*/`
`14`	`14`	`package com.facebook.presto.execution;`
`15`	`15`
	`16`	`+import com.facebook.presto.spi.ConnectorSession;`
`16`	`17`	`import com.facebook.presto.spi.PrestoException;`
`17`	`18`	`import com.facebook.presto.spi.StandardErrorCode;`
`18`	`19`	`import com.facebook.presto.spi.WarningCollector;`
`@@ -36,7 +37,7 @@ public TriggerFailurePlanChecker(AtomicBoolean triggerValidationFailure)`
`36`	`37`	`}`
`37`	`38`
`38`	`39`	`@Override`
`39`		`- public void validate(PlanNode planNode, WarningCollector warningCollector)`
	`40`	`+ public void validate(PlanNode planNode, WarningCollector warningCollector, ConnectorSession session)`
`40`	`41`	`{`
`41`	`42`	`if (triggerValidationFailure.get()) {`
`42`	`43`	`throw new PrestoException(FAILURE_ERROR_CODE, FAILURE_MESSAGE);`