Misc changes

Author: psachdeva-ss

singlestoredb/apps/_config.py

@@ -64,3 +64,52 @@ def token(self) -> Optional[str]:
             return self.app_token
         else:
             return self.user_token
+
+
+@dataclass
+class PythonUdfAppConfig:
+    listen_port: int
+    base_url: str
+    base_path: str
+    running_interactively: bool
+    is_gateway_enabled: bool
+
+    @staticmethod
+    def _read_variable(name: str) -> str:
+        value = os.environ.get(name)
+        if value is None:
+            raise RuntimeError(
+                f'Missing {name} environment variable. '
+                'Is the code running outside SingleStoreDB notebook environment?',
+            )
+        return value
+
+    @classmethod
+    def from_env(cls) -> 'AppConfig':
+        port = cls._read_variable('SINGLESTOREDB_APP_LISTEN_PORT')
+        base_url = cls._read_variable('SINGLESTOREDB_APP_BASE_URL')
+        base_path = cls._read_variable('SINGLESTOREDB_APP_BASE_PATH')
+
+        workload_type = os.environ.get('SINGLESTOREDB_WORKLOAD_TYPE')
+        running_interactively = workload_type == 'InteractiveNotebook'
+
+        is_gateway_enabled = 'SINGLESTOREDB_NOVA_GATEWAY_ENDPOINT' in os.environ
+
+        if running_interactively:
+            if is_gateway_enabled:
+                base_url = cls._read_variable('SINGLESTOREDB_PYTHON_UDF_BASE_URL')
+                base_path = cls._read_variable('SINGLESTOREDB_PYTHON_UDF_BASE_PATH')
+                assert base_url is not None
+                assert base_path is not None
+            else:
+                raise RuntimeError(
+                    'Running Python UDFs in interactive mode without nova-gateway enabled is not supported'
+                )
+
+        return cls(
+            listen_port=int(port),
+            base_url=base_url,
+            base_path=base_path,
+            running_interactively=running_interactively,
+            is_gateway_enabled=is_gateway_enabled,
+        )

singlestoredb/apps/_connection_info.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass
-from typing import Optional
+from typing import Optional, Dict, Any
 
 
 @dataclass
@@ -8,3 +8,8 @@ class ConnectionInfo:
 
     # Only present in interactive mode
     token: Optional[str]
+
+@dataclass
+class PythonUdfConnectionInfo:
+    url: str
+    functions: Dict[str, Any] 

singlestoredb/apps/_python_udfs.py

@@ -1,9 +1,10 @@
 import asyncio
 import textwrap
 import typing
+import os
 
-from ._config import AppConfig
-from ._connection_info import ConnectionInfo
+from ._config import PythonUdfAppConfig
+from ._connection_info import ConnectionInfo, PythonUdfConnectionInfo
 from ._process import kill_process_by_port
 from ..functions.ext.asgi import Application
 
@@ -15,9 +16,8 @@
 
 
 async def run_udf_app(
-    app: Application,
     log_level: str = 'error',
-    kill_existing_app_server: bool = False,
+    kill_existing_app_server: bool = True,
 ) -> ConnectionInfo:
     global _running_server
     from ._uvicorn_util import AwaitableUvicornServer
@@ -27,7 +27,7 @@ async def run_udf_app(
     except ImportError:
         raise ImportError('package uvicorn is required to run python udfs')
 
-    app_config = AppConfig.from_env()
+    app_config = PythonUdfAppConfig.from_env()
 
     if kill_existing_app_server:
         # Shutdown the server gracefully if it was started by us.
@@ -40,10 +40,9 @@ async def run_udf_app(
         # Kill if any other process is occupying the port
         kill_process_by_port(app_config.listen_port)
 
+    app = Application()
     app.root_path = app_config.base_path
 
-    print("Listening on port", app_config.listen_port)
-
     config = uvicorn.Config(
         app,
         host='0.0.0.0',
@@ -52,14 +51,14 @@ async def run_udf_app(
     )
     _running_server = AwaitableUvicornServer(config)
 
+    # In interactive mode this should be set to true
+    replace = app_config.running_interactively
     app.register_functions(replace=True)
+
     asyncio.create_task(_running_server.serve())
     await _running_server.wait_for_startup()
 
-    connection_info = ConnectionInfo(app_config.base_url, app_config.token)
+    connection_info = PythonUdfConnectionInfo(app_config.base_url, app.get_function_info())
 
-    print(
-        'Following Python UDFs are available: ', app.get_function_info()
-    )
 
     return connection_info

singlestoredb/functions/ext/asgi.py

@@ -500,7 +500,7 @@ def __init__(
             ]
         ] = None,
         app_mode: str = get_option('external_function.app_mode'),
-        url: str = get_option('external_function.url') + "invoke",
+        url: str = get_option('external_function.url'),
         data_format: str = get_option('external_function.data_format'),
         data_version: str = get_option('external_function.data_version'),
         link_name: Optional[str] = get_option('external_function.link_name'),