[FLINK-32299][gateway] Upload python jar when sql contains python udf…

… jar
liuyongvs · Nov 17, 2023 · febbbf3 · febbbf3
1 parent 873bd13
commit febbbf3
Show file tree

Hide file tree

Showing 7 changed files with 38 additions and 38 deletions.
diff --git a/docs/themes/book b/docs/themes/book
diff --git a/...k-sql-client/src/main/java/org/apache/flink/table/client/gateway/DefaultContextUtils.java b/...k-sql-client/src/main/java/org/apache/flink/table/client/gateway/DefaultContextUtils.java
@@ -54,14 +54,13 @@ public static DefaultContext buildDefaultContext(CliOptions.EmbeddedCliOptions o
         }
         Configuration sessionConfig = options.getPythonConfiguration();
         sessionConfig.addAll(ConfigurationUtils.createConfiguration(options.getSessionConfig()));
-        return DefaultContext.load(sessionConfig, discoverDependencies(jars, libDirs), true, true);
+        return DefaultContext.load(sessionConfig, discoverDependencies(jars, libDirs), true);
     }
 
     public static DefaultContext buildDefaultContext(CliOptions.GatewayCliOptions options) {
         return DefaultContext.load(
                 ConfigurationUtils.createConfiguration(options.getSessionConfig()),
                 Collections.emptyList(),
-                false,
                 false);
     }
     // --------------------------------------------------------------------------------------------

diff --git a/flink-table/flink-sql-gateway/src/main/java/org/apache/flink/table/gateway/SqlGateway.java b/flink-table/flink-sql-gateway/src/main/java/org/apache/flink/table/gateway/SqlGateway.java
@@ -113,7 +113,6 @@ static void startSqlGateway(PrintStream stream, String[] args) {
                 DefaultContext.load(
                         ConfigurationUtils.createConfiguration(cliOptions.getDynamicConfigs()),
                         Collections.emptyList(),
-                        true,
                         true);
         SqlGateway gateway =
                 new SqlGateway(

diff --git a/...-gateway/src/main/java/org/apache/flink/table/gateway/service/context/DefaultContext.java b/...-gateway/src/main/java/org/apache/flink/table/gateway/service/context/DefaultContext.java
@@ -36,11 +36,7 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-import java.net.URISyntaxException;
 import java.net.URL;
-import java.nio.file.Paths;
-import java.util.ArrayList;
-import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.stream.Collectors;
@@ -137,13 +133,9 @@ private static CustomCommandLine findActiveCommandLine(
      * @param dynamicConfig user specified configuration.
      * @param dependencies user specified jars
      * @param discoverExecutionConfig flag whether to load the execution configuration
-     * @param discoverPythonJar flag whetehr to load the python jar
      */
     public static DefaultContext load(
-            Configuration dynamicConfig,
-            List<URL> dependencies,
-            boolean discoverExecutionConfig,
-            boolean discoverPythonJar) {
+            Configuration dynamicConfig, List<URL> dependencies, boolean discoverExecutionConfig) {
         // 1. find the configuration directory
         String flinkConfigDir = CliFrontend.getConfigurationDirectoryFromEnv();
 
@@ -159,11 +151,6 @@ public static DefaultContext load(
         FileSystem.initialize(
                 configuration, PluginUtils.createPluginManagerFromRootFolder(configuration));
 
-        if (discoverPythonJar) {
-            dependencies = new ArrayList<>(dependencies);
-            dependencies.addAll(discoverPythonDependencies());
-        }
-
         if (discoverExecutionConfig) {
             Options commandLineOptions = collectCommandLineOptions(commandLines);
 
@@ -184,23 +171,4 @@ public static DefaultContext load(
 
         return new DefaultContext(configuration, dependencies);
     }
-
-    private static List<URL> discoverPythonDependencies() {
-        try {
-            URL location =
-                    Class.forName(
-                                    "org.apache.flink.python.PythonFunctionRunner",
-                                    false,
-                                    Thread.currentThread().getContextClassLoader())
-                            .getProtectionDomain()
-                            .getCodeSource()
-                            .getLocation();
-            if (Paths.get(location.toURI()).toFile().isFile()) {
-                return Collections.singletonList(location);
-            }
-        } catch (URISyntaxException | ClassNotFoundException e) {
-            LOG.warn("Failed to find flink-python jar." + e);
-        }
-        return Collections.emptyList();
-    }
 }
diff --git a/...rc/test/java/org/apache/flink/table/gateway/service/utils/SqlGatewayServiceExtension.java b/...rc/test/java/org/apache/flink/table/gateway/service/utils/SqlGatewayServiceExtension.java
@@ -89,7 +89,7 @@ public void beforeAll(ExtensionContext context) throws Exception {
             sessionManager =
                     sessionManagerCreator.apply(
                             DefaultContext.load(
-                                    new Configuration(), Collections.emptyList(), true, false));
+                                    new Configuration(), Collections.emptyList(), true));
         } finally {
             CommonTestUtils.setEnv(originalEnv);
         }

diff --git a/...le/flink-table-api-java/src/main/java/org/apache/flink/table/catalog/FunctionCatalog.java b/...le/flink-table-api-java/src/main/java/org/apache/flink/table/catalog/FunctionCatalog.java
@@ -771,6 +771,9 @@ private FunctionDefinition getFunctionDefinition(String name, CatalogFunction fu
             // directly.
             return ((InlineCatalogFunction) function).getDefinition();
         }
+        if (function.getFunctionLanguage() == FunctionLanguage.PYTHON) {
+            resourceManager.registerPythonResources();
+        }
         // If the jar resource of UDF used is not empty, register it to classloader before
         // validate.
         registerFunctionJarResources(name, function.getFunctionResources());

diff --git a/...e/flink-table-api-java/src/main/java/org/apache/flink/table/resource/ResourceManager.java b/...e/flink-table-api-java/src/main/java/org/apache/flink/table/resource/ResourceManager.java
@@ -42,8 +42,10 @@
 import java.io.Closeable;
 import java.io.FileNotFoundException;
 import java.io.IOException;
+import java.net.URISyntaxException;
 import java.net.URL;
 import java.net.URLClassLoader;
+import java.nio.file.Paths;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
@@ -78,6 +80,8 @@ public class ResourceManager implements Closeable {
     protected final Map<ResourceUri, URL> resourceInfos;
     protected final MutableURLClassLoader userClassLoader;
 
+    private boolean containPython = false;
+
     public static ResourceManager createResourceManager(
             URL[] urls, ClassLoader parent, ReadableConfig config) {
         MutableURLClassLoader mutableURLClassLoader =
@@ -197,6 +201,13 @@ public void unregisterFunctionResources(List<ResourceUri> resourceUris) {
         }
     }
 
+    public void registerPythonResources() {
+        if (!containPython) {
+            registerResources(discoverPythonDependencies(), true);
+            containPython = true;
+        }
+    }
+
     public URLClassLoader getUserClassLoader() {
         return userClassLoader;
     }
@@ -425,6 +436,26 @@ private Path getResourceLocalPath(Path remotePath) {
         return new Path(localResourceDir, fileNameWithUUID);
     }
 
+    private static Map<ResourceUri, URL> discoverPythonDependencies() {
+        try {
+            URL location =
+                    Class.forName(
+                                    "org.apache.flink.python.PythonFunctionRunner",
+                                    false,
+                                    Thread.currentThread().getContextClassLoader())
+                            .getProtectionDomain()
+                            .getCodeSource()
+                            .getLocation();
+            if (Paths.get(location.toURI()).toFile().isFile()) {
+                return Collections.singletonMap(
+                        new ResourceUri(ResourceType.JAR, location.getPath()), location);
+            }
+        } catch (URISyntaxException | ClassNotFoundException e) {
+            LOG.warn("Failed to find flink-python jar." + e);
+        }
+        return Collections.emptyMap();
+    }
+
     private void checkResources(Collection<ResourceUri> resourceUris, ResourceType expectedType)
             throws IOException {
         // check the resource type