Merge branch 'main' of https://github.com/singlestore-labs/singlestoredb-python into users/kaushik/udf-idle-start-fix

Author: kkampli-singlestore

singlestoredb/apps/_config.py

@@ -1,6 +1,7 @@
 import os
 from dataclasses import dataclass
 from typing import Optional
+from singlestoredb.connection import build_params
 
 
 @dataclass
@@ -9,6 +10,8 @@ class AppConfig:
     base_url: str
     base_path: str
     notebook_server_id: str
+    workspace_group_id: str
+    database_name: str
     app_token: Optional[str]
     user_token: Optional[str]
     running_interactively: bool
@@ -41,6 +44,15 @@ def from_env(cls) -> 'AppConfig':
         app_token = os.environ.get('SINGLESTOREDB_APP_TOKEN')
         user_token = os.environ.get('SINGLESTOREDB_USER_TOKEN')
 
+        if 'SINGLESTOREDB_URL' in os.environ:
+            dbname = build_params(host=os.environ['SINGLESTOREDB_URL']).get('database')
+        elif 'SINGLESTOREDB_HOST' in os.environ:
+            dbname = build_params(host=os.environ['SINGLESTOREDB_HOST']).get('database')
+        elif 'SINGLESTOREDB_DATABASE' in os.environ:
+            dbname = os.environ['SINGLESTOREDB_DATBASE']
+        
+        workspace_group_id = os.environ.get('SINGLESTOREDB_WORKSPACE_GROUP')
+
         # Make sure the required variables are present
         # and present useful error message if not
         if running_interactively:
@@ -54,6 +66,8 @@ def from_env(cls) -> 'AppConfig':
             base_url=base_url,
             base_path=base_path,
             notebook_server_id=notebook_server_id,
+            workspace_group_id=workspace_group_id,
+            database_name=dbname,
             app_token=app_token,
             user_token=user_token,
             running_interactively=running_interactively,

singlestoredb/apps/_python_udfs.py

@@ -6,6 +6,7 @@
 from ._config import AppConfig
 from ._connection_info import UdfConnectionInfo
 from ._process import kill_process_by_port
+from singlestoredb.connection import build_params
 
 if typing.TYPE_CHECKING:
     from ._uvicorn_util import AwaitableUvicornServer
@@ -43,7 +44,7 @@ async def run_udf_app(
 
     udf_suffix = ''
     if app_config.running_interactively:
-        udf_suffix = '_test'
+        udf_suffix = "_" +  app_config.notebook_server_id
     app = Application(url=base_url, app_mode='managed', name_suffix=udf_suffix)
 
     config = uvicorn.Config(

singlestoredb/functions/ext/asgi.py

@@ -53,6 +53,7 @@
 from typing import Set
 from typing import Tuple
 from typing import Union
+from ...apps._config import AppConfig
 
 from . import arrow
 from . import json as jdata
@@ -67,8 +68,6 @@
 from ..typing import Masked
 from ..typing import Table
 from ...config import get_option
-from singlestoredb.connection import build_params
-
 
 try:
     import cloudpickle
@@ -905,8 +904,18 @@ async def __call__(
 
         # Return function info
         elif method == 'GET' and (path == self.show_function_info_path or not path):
+            app_config = AppConfig.from_env()
             functions = self.get_function_info()
-            body = json.dumps(dict(functions=functions)).encode('utf-8')
+            workspace = app_config.workspace_group_id
+            database = app_config.database_name
+
+            body = json.dumps({
+                    "functions": functions,
+                    "cluster_id": workspace,
+                    "database": database,
+                }).encode('utf-8')
+            
+            print(body)
             await send(self.text_response_dict)
 
         # Path not found
@@ -981,18 +990,7 @@ def get_function_info(
             sig = info['signature']
             sql_map[sig['name']] = sql
 
-        if 'SINGLESTOREDB_URL' in os.environ:
-            dbname = build_params(host=os.environ['SINGLESTOREDB_URL']).get('database')
-        elif 'SINGLESTOREDB_HOST' in os.environ:
-            dbname = build_params(host=os.environ['SINGLESTOREDB_HOST']).get('database')
-        elif 'SINGLESTOREDB_DATABASE' in os.environ:
-            dbname = os.environ['SINGLESTOREDB_DATBASE']
 
-        connection_info = {}
-        workspace_group_id = os.environ.get('SINGLESTOREDB_WORKSPACE_GROUP')
-        connection_info['database_name'] = dbname
-        connection_info['workspace_group_id'] = workspace_group_id
-
 
         for key, (_, info) in self.endpoints.items():
             if not func_name or key == func_name:
@@ -1039,10 +1037,7 @@ def get_function_info(
                     sql_statement=sql,
                 )
 
-        return {
-            'functions': functions,
-            'connection_info': connection_info
-        }
+        return functions
 
     def get_create_functions(
         self,