From 19051af70f3d943d3137fc064ad27b0504c03a8a Mon Sep 17 00:00:00 2001 From: Karish Date: Tue, 29 Sep 2026 15:00:04 +0530 Subject: [PATCH 1/5] [backend:fix] fix interactive and portal publish PythonUDF MCDB-100138 --- singlestoredb/apps/_python_udfs.py | 8 +- singlestoredb/functions/ext/asgi.py | 66 ++++++++- singlestoredb/functions/ext/function_url.py | 73 ++++++++++ singlestoredb/tests/test_function_url.py | 144 ++++++++++++++++++++ 4 files changed, 288 insertions(+), 3 deletions(-) create mode 100644 singlestoredb/functions/ext/function_url.py create mode 100644 singlestoredb/tests/test_function_url.py diff --git a/singlestoredb/apps/_python_udfs.py b/singlestoredb/apps/_python_udfs.py index 30295f32b..ae63f54bd 100644 --- a/singlestoredb/apps/_python_udfs.py +++ b/singlestoredb/apps/_python_udfs.py @@ -78,13 +78,19 @@ async def run_udf_app( # Register the functions only if the app is running interactively. if app_config.running_interactively: - app.register_functions(replace=True) + app.register_interactive_functions() _running_server = AwaitableUvicornServer(config) asyncio.create_task(_running_server.serve()) await _running_server.wait_for_startup() print(f'Python UDF registered at {base_url}') + if app_config.running_interactively: + sql_names = [ + info['signature']['name'] + for _func, info in app.endpoints.values() + ] + print(f'Registered SQL functions: {", ".join(sql_names)}') return UdfConnectionInfo(base_url, app.get_function_info()) diff --git a/singlestoredb/functions/ext/asgi.py b/singlestoredb/functions/ext/asgi.py index 876cf9eab..d40fdb109 100755 --- a/singlestoredb/functions/ext/asgi.py +++ b/singlestoredb/functions/ext/asgi.py @@ -67,6 +67,10 @@ from . import json as jdata from . import rowdat_1 from . import utils +from .function_url import classify_interactive_registration +from .function_url import extract_service_url +from .function_url import is_function_not_defined +from .function_url import urls_equal from ... import connection from ...config import get_option from ...management.stage import get_stage @@ -1607,8 +1611,7 @@ def _locate_app_functions(self, cur: Any) -> Tuple[Set[str], Set[str]]: # See if function URL matches url cur.execute(f'SHOW CREATE FUNCTION {database_prefix}{escape_name(name)}') for fname, _, code, *_ in list(cur): - m = re.search(r" (?:\w+) (?:SERVICE|MANAGED) '([^']+)'", code) - if m and m.group(1) == self.url: + if urls_equal(extract_service_url(code), self.url): funcs.add(f'{database_prefix}{escape_name(fname)}') if link and re.match(r'^py_ext_func_link_\S{14}$', link): links.add(link) @@ -1812,6 +1815,65 @@ def register_functions( for func in self.get_create_functions(replace=replace): cur.execute(func) + def _show_create_service_url(self, cur: Any, sql_name: str) -> Optional[str]: + if self.function_database: + qualified = ( + f'{escape_name(self.function_database)}.{escape_name(sql_name)}' + ) + else: + qualified = escape_name(sql_name) + try: + cur.execute(f'SHOW CREATE FUNCTION {qualified}') + except Exception as exc: + if is_function_not_defined(exc): + return None + raise + rows = list(cur) + if not rows: + return None + code = rows[0][2] + if isinstance(code, bytes): + code = code.decode('utf-8') + return extract_service_url(code) + + def register_interactive_functions( + self, + *connection_args: Any, + **connection_kwargs: Any, + ) -> None: + """Register functions for an interactive notebook session. + + Creates or replaces a name only when it is missing or already + points at this session's interactive URL. Published and + other-session functions are left unchanged. + """ + with connection.connect(*connection_args, **connection_kwargs) as conn: + with conn.cursor() as cur: + for _key, (_endpoint, info) in self.endpoints.items(): + sig = info['signature'] + sql_name = sig['name'] + existing = self._show_create_service_url(cur, sql_name) + try: + action = classify_interactive_registration( + existing, self.url, + ) + except ValueError as exc: + raise RuntimeError( + f'Cannot register SQL function `{sql_name}`: {exc}', + ) from exc + create_sqls = [ + signature_to_sql( + sig, + url=self.url, + data_format=self.data_format, + app_mode=self.app_mode, + replace=(action == 'replace'), + database=self.function_database or None, + ), + ] + for stmt in create_sqls: + cur.execute(stmt) + def drop_functions( self, *connection_args: Any, diff --git a/singlestoredb/functions/ext/function_url.py b/singlestoredb/functions/ext/function_url.py new file mode 100644 index 000000000..59f50089e --- /dev/null +++ b/singlestoredb/functions/ext/function_url.py @@ -0,0 +1,73 @@ +"""Parse and classify external function service URLs from SHOW CREATE.""" + +from __future__ import annotations + +import re +from typing import Optional + +from ...mysql.constants import ER + + +# Managed Python UDFs emit: AS MANAGED SERVICE "https://..." +# Older / remote forms may use SERVICE 'https://...' +_SERVICE_URL_RE = re.compile( + r'(?:MANAGED\s+)?SERVICE\s+[\'"]([^\'"]+)[\'"]', + re.IGNORECASE, +) + + +def extract_service_url(create_sql: str) -> Optional[str]: + """Return the MANAGED/REMOTE SERVICE URL from SHOW CREATE FUNCTION text.""" + if not create_sql: + return None + match = _SERVICE_URL_RE.search(create_sql) + if match is None: + return None + return match.group(1) + + +def normalize_service_url(url: str) -> str: + return url.rstrip('/') + + +def urls_equal(left: Optional[str], right: Optional[str]) -> bool: + if not left or not right: + return False + return normalize_service_url(left) == normalize_service_url(right) + + +def is_interactive_service_url(url: Optional[str]) -> bool: + if not url: + return False + return '/interactive' in normalize_service_url(url).lower() + + +def classify_interactive_registration( + existing_url: Optional[str], + this_session_url: str, +) -> str: + """Return 'create', 'replace', or raise ValueError if the name is not ours. + + Interactive registration may only create a missing name or replace a + function that already points at this notebook session's /interactive/ URL. + """ + if not existing_url: + return 'create' + if urls_equal(existing_url, this_session_url): + return 'replace' + raise ValueError( + f'Cannot register over existing function pointing at {existing_url} ' + f'(this session is {this_session_url}). ' + 'Interactive registration will not replace a published or ' + 'other-session function.', + ) + + +def is_function_not_defined(exc: BaseException) -> bool: + errno = getattr(exc, 'errno', None) + if errno == ER.FUNCTION_NOT_DEFINED: + return True + args = getattr(exc, 'args', ()) + if args and args[0] == ER.FUNCTION_NOT_DEFINED: + return True + return False diff --git a/singlestoredb/tests/test_function_url.py b/singlestoredb/tests/test_function_url.py new file mode 100644 index 000000000..1b2bca458 --- /dev/null +++ b/singlestoredb/tests/test_function_url.py @@ -0,0 +1,144 @@ +import unittest + +from singlestoredb.functions.ext.function_url import classify_interactive_registration +from singlestoredb.functions.ext.function_url import extract_service_url +from singlestoredb.functions.ext.function_url import is_function_not_defined +from singlestoredb.functions.ext.function_url import is_interactive_service_url +from singlestoredb.mysql.constants import ER + + +class TestFunctionURL(unittest.TestCase): + + def test_extract_managed_double_quoted(self): + sql = ( + 'CREATE OR REPLACE EXTERNAL FUNCTION `shi_test`(v BIGINT) ' + 'RETURNS TEXT AS MANAGED SERVICE ' + '"https://apps.us-east-1.cloud.singlestore.com/pythonudfs/abc/" ' + 'FORMAT ROWDAT_1;' + ) + self.assertEqual( + extract_service_url(sql), + 'https://apps.us-east-1.cloud.singlestore.com/pythonudfs/abc/', + ) + + def test_extract_service_single_quoted(self): + sql = ( + "CREATE EXTERNAL FUNCTION foo() RETURNS INT AS REMOTE SERVICE " + "'https://example.com/invoke' FORMAT ROWDAT_1;" + ) + self.assertEqual(extract_service_url(sql), 'https://example.com/invoke') + + def test_extract_none(self): + self.assertIsNone(extract_service_url('CREATE FUNCTION foo() RETURNS INT')) + + def test_classify_missing_create(self): + self.assertEqual( + classify_interactive_registration( + None, + 'https://gw/pythonudfs/sess/interactive/', + ), + 'create', + ) + + def test_classify_this_session_replace(self): + this = 'https://gw/pythonudfs/sess/interactive/' + self.assertEqual( + classify_interactive_registration(this, this), + 'replace', + ) + self.assertEqual( + classify_interactive_registration(this.rstrip('/'), this), + 'replace', + ) + + def test_classify_published_refuses(self): + with self.assertRaises(ValueError) as ctx: + classify_interactive_registration( + 'https://gw/pythonudfs/published-id/', + 'https://gw/pythonudfs/sess/interactive/', + ) + self.assertIn('will not replace', str(ctx.exception)) + self.assertIn('this session is', str(ctx.exception)) + + def test_classify_other_session_refuses(self): + with self.assertRaises(ValueError): + classify_interactive_registration( + 'https://gw/pythonudfs/other-sess/interactive/', + 'https://gw/pythonudfs/sess/interactive/', + ) + + def test_is_interactive(self): + self.assertTrue( + is_interactive_service_url( + 'https://gw/pythonudfs/sess/interactive/', + ), + ) + self.assertFalse( + is_interactive_service_url( + 'https://gw/pythonudfs/published-id/', + ), + ) + + def test_function_not_defined(self): + exc = type('E', (Exception,), {'errno': ER.FUNCTION_NOT_DEFINED})() + self.assertTrue(is_function_not_defined(exc)) + self.assertFalse(is_function_not_defined(ValueError('nope'))) + + def test_fake_cursor_ownership(self): + this_session = 'https://gw/pythonudfs/sess/interactive/' + published_sql = ( + 'CREATE OR REPLACE EXTERNAL FUNCTION `shi_test`(v BIGINT) ' + 'RETURNS BIGINT AS MANAGED SERVICE ' + '"https://gw/pythonudfs/published-id/" FORMAT ROWDAT_1;' + ) + this_sql = ( + 'CREATE OR REPLACE EXTERNAL FUNCTION `shi_test`(v BIGINT) ' + 'RETURNS BIGINT AS MANAGED SERVICE ' + f'"{this_session}" FORMAT ROWDAT_1;' + ) + + class _Cursor: + def __init__(self, rows=None, exc=None): + self._rows = rows or [] + self._exc = exc + + def execute(self, sql): + if self._exc is not None: + raise self._exc + + def __iter__(self): + return iter(self._rows) + + def show_url(cur): + try: + cur.execute('SHOW CREATE FUNCTION `shi_test`') + except Exception as exc: + if is_function_not_defined(exc): + return None + raise + rows = list(cur) + if not rows: + return None + return extract_service_url(rows[0][2]) + + missing = type('E', (Exception,), {'errno': ER.FUNCTION_NOT_DEFINED})() + self.assertEqual( + classify_interactive_registration(show_url(_Cursor(exc=missing)), this_session), + 'create', + ) + self.assertEqual( + classify_interactive_registration( + show_url(_Cursor(rows=[('shi_test', '', this_sql)])), + this_session, + ), + 'replace', + ) + with self.assertRaises(ValueError): + classify_interactive_registration( + show_url(_Cursor(rows=[('shi_test', '', published_sql)])), + this_session, + ) + + +if __name__ == '__main__': + unittest.main() From 18a98b12698e09f40b73dba60f37b06cc4f45147 Mon Sep 17 00:00:00 2001 From: Karish Date: Tue, 29 Sep 2026 15:16:36 +0530 Subject: [PATCH 2/5] [backend:fix] pre-commit check fixes MCDB-100138 --- singlestoredb/functions/ext/asgi.py | 8 ++-- singlestoredb/functions/ext/function_url.py | 1 - singlestoredb/tests/test_function_url.py | 45 +++++++++++++-------- 3 files changed, 33 insertions(+), 21 deletions(-) diff --git a/singlestoredb/functions/ext/asgi.py b/singlestoredb/functions/ext/asgi.py index d40fdb109..b9bebf602 100755 --- a/singlestoredb/functions/ext/asgi.py +++ b/singlestoredb/functions/ext/asgi.py @@ -67,10 +67,6 @@ from . import json as jdata from . import rowdat_1 from . import utils -from .function_url import classify_interactive_registration -from .function_url import extract_service_url -from .function_url import is_function_not_defined -from .function_url import urls_equal from ... import connection from ...config import get_option from ...management.stage import get_stage @@ -79,6 +75,10 @@ from ..signature import signature_to_sql from ..typing import Masked from ..typing import Table +from .function_url import classify_interactive_registration +from .function_url import extract_service_url +from .function_url import is_function_not_defined +from .function_url import urls_equal from .timer import Timer from singlestoredb.docstring.parser import parse from singlestoredb.functions.dtypes import escape_name diff --git a/singlestoredb/functions/ext/function_url.py b/singlestoredb/functions/ext/function_url.py index 59f50089e..08bc431cf 100644 --- a/singlestoredb/functions/ext/function_url.py +++ b/singlestoredb/functions/ext/function_url.py @@ -1,5 +1,4 @@ """Parse and classify external function service URLs from SHOW CREATE.""" - from __future__ import annotations import re diff --git a/singlestoredb/tests/test_function_url.py b/singlestoredb/tests/test_function_url.py index 1b2bca458..2513acfff 100644 --- a/singlestoredb/tests/test_function_url.py +++ b/singlestoredb/tests/test_function_url.py @@ -1,4 +1,10 @@ +# type: ignore import unittest +from collections.abc import Iterator +from typing import Any +from typing import List +from typing import Optional +from typing import Tuple from singlestoredb.functions.ext.function_url import classify_interactive_registration from singlestoredb.functions.ext.function_url import extract_service_url @@ -9,7 +15,7 @@ class TestFunctionURL(unittest.TestCase): - def test_extract_managed_double_quoted(self): + def test_extract_managed_double_quoted(self) -> None: sql = ( 'CREATE OR REPLACE EXTERNAL FUNCTION `shi_test`(v BIGINT) ' 'RETURNS TEXT AS MANAGED SERVICE ' @@ -21,17 +27,17 @@ def test_extract_managed_double_quoted(self): 'https://apps.us-east-1.cloud.singlestore.com/pythonudfs/abc/', ) - def test_extract_service_single_quoted(self): + def test_extract_service_single_quoted(self) -> None: sql = ( - "CREATE EXTERNAL FUNCTION foo() RETURNS INT AS REMOTE SERVICE " + 'CREATE EXTERNAL FUNCTION foo() RETURNS INT AS REMOTE SERVICE ' "'https://example.com/invoke' FORMAT ROWDAT_1;" ) self.assertEqual(extract_service_url(sql), 'https://example.com/invoke') - def test_extract_none(self): + def test_extract_none(self) -> None: self.assertIsNone(extract_service_url('CREATE FUNCTION foo() RETURNS INT')) - def test_classify_missing_create(self): + def test_classify_missing_create(self) -> None: self.assertEqual( classify_interactive_registration( None, @@ -40,7 +46,7 @@ def test_classify_missing_create(self): 'create', ) - def test_classify_this_session_replace(self): + def test_classify_this_session_replace(self) -> None: this = 'https://gw/pythonudfs/sess/interactive/' self.assertEqual( classify_interactive_registration(this, this), @@ -51,7 +57,7 @@ def test_classify_this_session_replace(self): 'replace', ) - def test_classify_published_refuses(self): + def test_classify_published_refuses(self) -> None: with self.assertRaises(ValueError) as ctx: classify_interactive_registration( 'https://gw/pythonudfs/published-id/', @@ -60,14 +66,14 @@ def test_classify_published_refuses(self): self.assertIn('will not replace', str(ctx.exception)) self.assertIn('this session is', str(ctx.exception)) - def test_classify_other_session_refuses(self): + def test_classify_other_session_refuses(self) -> None: with self.assertRaises(ValueError): classify_interactive_registration( 'https://gw/pythonudfs/other-sess/interactive/', 'https://gw/pythonudfs/sess/interactive/', ) - def test_is_interactive(self): + def test_is_interactive(self) -> None: self.assertTrue( is_interactive_service_url( 'https://gw/pythonudfs/sess/interactive/', @@ -79,12 +85,12 @@ def test_is_interactive(self): ), ) - def test_function_not_defined(self): + def test_function_not_defined(self) -> None: exc = type('E', (Exception,), {'errno': ER.FUNCTION_NOT_DEFINED})() self.assertTrue(is_function_not_defined(exc)) self.assertFalse(is_function_not_defined(ValueError('nope'))) - def test_fake_cursor_ownership(self): + def test_fake_cursor_ownership(self) -> None: this_session = 'https://gw/pythonudfs/sess/interactive/' published_sql = ( 'CREATE OR REPLACE EXTERNAL FUNCTION `shi_test`(v BIGINT) ' @@ -98,18 +104,22 @@ def test_fake_cursor_ownership(self): ) class _Cursor: - def __init__(self, rows=None, exc=None): + def __init__( + self, + rows: Optional[List[Tuple[str, str, str]]] = None, + exc: Optional[BaseException] = None, + ) -> None: self._rows = rows or [] self._exc = exc - def execute(self, sql): + def execute(self, sql: str) -> None: if self._exc is not None: raise self._exc - def __iter__(self): + def __iter__(self) -> Iterator[Tuple[str, str, str]]: return iter(self._rows) - def show_url(cur): + def show_url(cur: Any) -> Optional[str]: try: cur.execute('SHOW CREATE FUNCTION `shi_test`') except Exception as exc: @@ -123,7 +133,10 @@ def show_url(cur): missing = type('E', (Exception,), {'errno': ER.FUNCTION_NOT_DEFINED})() self.assertEqual( - classify_interactive_registration(show_url(_Cursor(exc=missing)), this_session), + classify_interactive_registration( + show_url(_Cursor(exc=missing)), + this_session, + ), 'create', ) self.assertEqual( From ab113bec402d352bb8adf7e298e23a056649e5d0 Mon Sep 17 00:00:00 2001 From: Karish Date: Tue, 29 Sep 2026 15:47:20 +0530 Subject: [PATCH 3/5] [backend:fix] PR comments MCDB-100138 --- singlestoredb/functions/ext/asgi.py | 23 ++++- singlestoredb/tests/test_function_url.py | 105 +++++++++++++++++++++++ 2 files changed, 124 insertions(+), 4 deletions(-) diff --git a/singlestoredb/functions/ext/asgi.py b/singlestoredb/functions/ext/asgi.py index b9bebf602..f682828de 100755 --- a/singlestoredb/functions/ext/asgi.py +++ b/singlestoredb/functions/ext/asgi.py @@ -1849,9 +1849,18 @@ def register_interactive_functions( """ with connection.connect(*connection_args, **connection_kwargs) as conn: with conn.cursor() as cur: + planned = [] + if self.function_database: + database_prefix = escape_name(self.function_database) + '.' + else: + database_prefix = '' + current_names = set() for _key, (_endpoint, info) in self.endpoints.items(): sig = info['signature'] sql_name = sig['name'] + current_names.add( + f'{database_prefix}{escape_name(sql_name)}', + ) existing = self._show_create_service_url(cur, sql_name) try: action = classify_interactive_registration( @@ -1861,7 +1870,7 @@ def register_interactive_functions( raise RuntimeError( f'Cannot register SQL function `{sql_name}`: {exc}', ) from exc - create_sqls = [ + planned.append( signature_to_sql( sig, url=self.url, @@ -1870,9 +1879,15 @@ def register_interactive_functions( replace=(action == 'replace'), database=self.function_database or None, ), - ] - for stmt in create_sqls: - cur.execute(stmt) + ) + + funcs, _links = self._locate_app_functions(cur) + for fname in funcs: + if fname not in current_names: + cur.execute(f'DROP FUNCTION IF EXISTS {fname}') + + for stmt in planned: + cur.execute(stmt) def drop_functions( self, diff --git a/singlestoredb/tests/test_function_url.py b/singlestoredb/tests/test_function_url.py index 2513acfff..5a845fd32 100644 --- a/singlestoredb/tests/test_function_url.py +++ b/singlestoredb/tests/test_function_url.py @@ -4,7 +4,9 @@ from typing import Any from typing import List from typing import Optional +from typing import Set from typing import Tuple +from unittest.mock import patch from singlestoredb.functions.ext.function_url import classify_interactive_registration from singlestoredb.functions.ext.function_url import extract_service_url @@ -152,6 +154,109 @@ def show_url(cur: Any) -> Optional[str]: this_session, ) + def _interactive_app(self, names: List[str]) -> Any: + from singlestoredb.functions.ext.asgi import Application + + app = Application.__new__(Application) + app.function_database = None + app.url = 'https://gw/pythonudfs/sess/interactive/' + app.data_format = 'rowdat_1' + app.app_mode = 'managed' + app.endpoints = {} + for name in names: + app.endpoints[name.encode('utf-8')] = ( + None, + { + 'signature': { + 'name': name, + 'args': [{'name': 'v', 'sql': 'BIGINT'}], + 'returns': [{'name': '', 'sql': 'BIGINT'}], + }, + }, + ) + return app + + def test_register_interactive_preflight(self) -> None: + app = self._interactive_app(['keep_test', 'steal_test']) + executed: List[str] = [] + + def show(_cur: Any, sql_name: str) -> Optional[str]: + if sql_name == 'steal_test': + return 'https://gw/pythonudfs/published-id/' + return None + + app._show_create_service_url = show # type: ignore[method-assign] + + class _Cursor: + def execute(self, sql: str) -> None: + executed.append(sql) + + class _CM: + def __init__(self, inner: Any) -> None: + self.inner = inner + + def __enter__(self) -> Any: + return self.inner + + def __exit__(self, *args: Any) -> bool: + return False + + class _Conn: + def cursor(self) -> _CM: + return _CM(_Cursor()) + + with patch( + 'singlestoredb.functions.ext.asgi.connection.connect', + return_value=_CM(_Conn()), + ): + with self.assertRaises(RuntimeError): + app.register_interactive_functions() + self.assertEqual(executed, []) + + def test_register_interactive_drops_stale(self) -> None: + app = self._interactive_app(['keep_test']) + executed: List[str] = [] + + def show(_cur: Any, _sql_name: str) -> Optional[str]: + return None + + def locate(_cur: Any) -> Tuple[Set[str], Set[str]]: + return {'`keep_test`', '`stale_test`'}, set() + + app._show_create_service_url = show # type: ignore[method-assign] + app._locate_app_functions = locate # type: ignore[method-assign] + + class _Cursor: + def execute(self, sql: str) -> None: + executed.append(sql) + + class _CM: + def __init__(self, inner: Any) -> None: + self.inner = inner + + def __enter__(self) -> Any: + return self.inner + + def __exit__(self, *args: Any) -> bool: + return False + + class _Conn: + def cursor(self) -> _CM: + return _CM(_Cursor()) + + with patch( + 'singlestoredb.functions.ext.asgi.connection.connect', + return_value=_CM(_Conn()), + ): + app.register_interactive_functions() + self.assertIn('DROP FUNCTION IF EXISTS `stale_test`', executed) + self.assertTrue( + any('CREATE' in sql and 'keep_test' in sql for sql in executed), + ) + self.assertFalse( + any('DROP' in sql and 'keep_test' in sql for sql in executed), + ) + if __name__ == '__main__': unittest.main() From 71cd0e39e075bcbfb6a54549d3c6fd8059cbc5c2 Mon Sep 17 00:00:00 2001 From: Karish Date: Thu, 1 Oct 2026 13:35:35 +0530 Subject: [PATCH 4/5] [backend:fix] PR comments MCDB-100138 --- singlestoredb/functions/ext/function_url.py | 4 ++-- singlestoredb/tests/test_function_url.py | 7 +++++-- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/singlestoredb/functions/ext/function_url.py b/singlestoredb/functions/ext/function_url.py index 08bc431cf..4d40ab90a 100644 --- a/singlestoredb/functions/ext/function_url.py +++ b/singlestoredb/functions/ext/function_url.py @@ -64,9 +64,9 @@ def classify_interactive_registration( def is_function_not_defined(exc: BaseException) -> bool: errno = getattr(exc, 'errno', None) - if errno == ER.FUNCTION_NOT_DEFINED: + if errno in (ER.FUNCTION_NOT_DEFINED, ER.SP_DOES_NOT_EXIST): return True args = getattr(exc, 'args', ()) - if args and args[0] == ER.FUNCTION_NOT_DEFINED: + if args and args[0] in (ER.FUNCTION_NOT_DEFINED, ER.SP_DOES_NOT_EXIST): return True return False diff --git a/singlestoredb/tests/test_function_url.py b/singlestoredb/tests/test_function_url.py index 5a845fd32..395a1deb4 100644 --- a/singlestoredb/tests/test_function_url.py +++ b/singlestoredb/tests/test_function_url.py @@ -88,8 +88,11 @@ def test_is_interactive(self) -> None: ) def test_function_not_defined(self) -> None: - exc = type('E', (Exception,), {'errno': ER.FUNCTION_NOT_DEFINED})() - self.assertTrue(is_function_not_defined(exc)) + missing = type('E', (Exception,), {'errno': ER.FUNCTION_NOT_DEFINED})() + self.assertTrue(is_function_not_defined(missing)) + sp_missing = type('E', (Exception,), {'errno': ER.SP_DOES_NOT_EXIST})() + self.assertTrue(is_function_not_defined(sp_missing)) + self.assertTrue(is_function_not_defined(Exception(ER.SP_DOES_NOT_EXIST))) self.assertFalse(is_function_not_defined(ValueError('nope'))) def test_fake_cursor_ownership(self) -> None: From 0288dd7ddf64064a78e9a4b0b470ccd0620772a0 Mon Sep 17 00:00:00 2001 From: Karish Date: Thu, 1 Oct 2026 15:18:34 +0530 Subject: [PATCH 5/5] [backend:fix] solve concurrent write and refusal killing live server fix MCDB-100138 --- singlestoredb/apps/_python_udfs.py | 26 +-- singlestoredb/functions/ext/asgi.py | 83 +++++--- singlestoredb/tests/test_function_url.py | 236 +++++++++++++++++++---- 3 files changed, 272 insertions(+), 73 deletions(-) diff --git a/singlestoredb/apps/_python_udfs.py b/singlestoredb/apps/_python_udfs.py index ae63f54bd..c061abc82 100644 --- a/singlestoredb/apps/_python_udfs.py +++ b/singlestoredb/apps/_python_udfs.py @@ -38,17 +38,6 @@ async def run_udf_app( app_config = AppConfig.from_env() - if kill_existing_app_server: - # Shutdown the server gracefully if it was started by us. - # Since the uvicorn server doesn't start a new subprocess - # killing the process would result in kernel dying. - if _running_server is not None: - await _running_server.shutdown() - _running_server = None - - # Kill if any other process is occupying the port - kill_process_by_port(app_config.listen_port) - base_url = generate_base_url(app_config) udf_suffix = '' @@ -68,6 +57,21 @@ async def run_udf_app( f'You can only define a maximum of {MAX_UDFS_LIMIT} functions.', ) + # Prove every name is allowed before killing a live interactive server. + if app_config.running_interactively: + app.preflight_interactive_functions() + + if kill_existing_app_server: + # Shutdown the server gracefully if it was started by us. + # Since the uvicorn server doesn't start a new subprocess + # killing the process would result in kernel dying. + if _running_server is not None: + await _running_server.shutdown() + _running_server = None + + # Kill if any other process is occupying the port + kill_process_by_port(app_config.listen_port) + config = uvicorn.Config( app, host='0.0.0.0', diff --git a/singlestoredb/functions/ext/asgi.py b/singlestoredb/functions/ext/asgi.py index f682828de..8cbd15982 100755 --- a/singlestoredb/functions/ext/asgi.py +++ b/singlestoredb/functions/ext/asgi.py @@ -1815,13 +1815,11 @@ def register_functions( for func in self.get_create_functions(replace=replace): cur.execute(func) - def _show_create_service_url(self, cur: Any, sql_name: str) -> Optional[str]: - if self.function_database: - qualified = ( - f'{escape_name(self.function_database)}.{escape_name(sql_name)}' - ) - else: - qualified = escape_name(sql_name) + def _service_url_for_qualified( + self, + cur: Any, + qualified: str, + ) -> Optional[str]: try: cur.execute(f'SHOW CREATE FUNCTION {qualified}') except Exception as exc: @@ -1836,6 +1834,39 @@ def _show_create_service_url(self, cur: Any, sql_name: str) -> Optional[str]: code = code.decode('utf-8') return extract_service_url(code) + def _show_create_service_url(self, cur: Any, sql_name: str) -> Optional[str]: + if self.function_database: + qualified = ( + f'{escape_name(self.function_database)}.{escape_name(sql_name)}' + ) + else: + qualified = escape_name(sql_name) + return self._service_url_for_qualified(cur, qualified) + + def _classify_interactive_name(self, cur: Any, sql_name: str) -> str: + existing = self._show_create_service_url(cur, sql_name) + try: + return classify_interactive_registration(existing, self.url) + except ValueError as exc: + raise RuntimeError( + f'Cannot register SQL function `{sql_name}`: {exc}', + ) from exc + + def preflight_interactive_functions( + self, + *connection_args: Any, + **connection_kwargs: Any, + ) -> None: + """Raise if any current name is published or owned by another session. + + Read-only: does not CREATE or DROP functions. + """ + with connection.connect(*connection_args, **connection_kwargs) as conn: + with conn.cursor() as cur: + for _key, (_endpoint, info) in self.endpoints.items(): + sql_name = info['signature']['name'] + self._classify_interactive_name(cur, sql_name) + def register_interactive_functions( self, *connection_args: Any, @@ -1845,32 +1876,34 @@ def register_interactive_functions( Creates or replaces a name only when it is missing or already points at this session's interactive URL. Published and - other-session functions are left unchanged. + other-session functions are left unchanged. Ownership is + re-checked immediately before each write or drop. """ with connection.connect(*connection_args, **connection_kwargs) as conn: with conn.cursor() as cur: - planned = [] if self.function_database: database_prefix = escape_name(self.function_database) + '.' else: database_prefix = '' current_names = set() for _key, (_endpoint, info) in self.endpoints.items(): - sig = info['signature'] - sql_name = sig['name'] + sql_name = info['signature']['name'] current_names.add( f'{database_prefix}{escape_name(sql_name)}', ) - existing = self._show_create_service_url(cur, sql_name) - try: - action = classify_interactive_registration( - existing, self.url, - ) - except ValueError as exc: - raise RuntimeError( - f'Cannot register SQL function `{sql_name}`: {exc}', - ) from exc - planned.append( + self._classify_interactive_name(cur, sql_name) + + funcs, _links = self._locate_app_functions(cur) + for fname in funcs: + if fname not in current_names: + existing = self._service_url_for_qualified(cur, fname) + if urls_equal(existing, self.url): + cur.execute(f'DROP FUNCTION IF EXISTS {fname}') + + for _key, (_endpoint, info) in self.endpoints.items(): + sig = info['signature'] + action = self._classify_interactive_name(cur, sig['name']) + cur.execute( signature_to_sql( sig, url=self.url, @@ -1881,14 +1914,6 @@ def register_interactive_functions( ), ) - funcs, _links = self._locate_app_functions(cur) - for fname in funcs: - if fname not in current_names: - cur.execute(f'DROP FUNCTION IF EXISTS {fname}') - - for stmt in planned: - cur.execute(stmt) - def drop_functions( self, *connection_args: Any, diff --git a/singlestoredb/tests/test_function_url.py b/singlestoredb/tests/test_function_url.py index 395a1deb4..771f7f977 100644 --- a/singlestoredb/tests/test_function_url.py +++ b/singlestoredb/tests/test_function_url.py @@ -1,4 +1,5 @@ # type: ignore +import sys import unittest from collections.abc import Iterator from typing import Any @@ -6,8 +7,10 @@ from typing import Optional from typing import Set from typing import Tuple +from unittest.mock import MagicMock from unittest.mock import patch +from singlestoredb.apps._config import AppConfig from singlestoredb.functions.ext.function_url import classify_interactive_registration from singlestoredb.functions.ext.function_url import extract_service_url from singlestoredb.functions.ext.function_url import is_function_not_defined @@ -179,17 +182,7 @@ def _interactive_app(self, names: List[str]) -> Any: ) return app - def test_register_interactive_preflight(self) -> None: - app = self._interactive_app(['keep_test', 'steal_test']) - executed: List[str] = [] - - def show(_cur: Any, sql_name: str) -> Optional[str]: - if sql_name == 'steal_test': - return 'https://gw/pythonudfs/published-id/' - return None - - app._show_create_service_url = show # type: ignore[method-assign] - + def _patch_connect(self, executed: List[str]) -> Any: class _Cursor: def execute(self, sql: str) -> None: executed.append(sql) @@ -208,49 +201,69 @@ class _Conn: def cursor(self) -> _CM: return _CM(_Cursor()) - with patch( + return patch( 'singlestoredb.functions.ext.asgi.connection.connect', return_value=_CM(_Conn()), - ): + ) + + def test_register_interactive_preflight(self) -> None: + app = self._interactive_app(['keep_test', 'steal_test']) + executed: List[str] = [] + + def show(_cur: Any, sql_name: str) -> Optional[str]: + if sql_name == 'steal_test': + return 'https://gw/pythonudfs/published-id/' + return None + + app._show_create_service_url = show # type: ignore[method-assign] + + with self._patch_connect(executed): with self.assertRaises(RuntimeError): app.register_interactive_functions() self.assertEqual(executed, []) - def test_register_interactive_drops_stale(self) -> None: + def test_register_interactive_rechecks_at_write(self) -> None: app = self._interactive_app(['keep_test']) executed: List[str] = [] + seen = {'n': 0} def show(_cur: Any, _sql_name: str) -> Optional[str]: - return None + seen['n'] += 1 + if seen['n'] == 1: + return app.url + return 'https://gw/pythonudfs/published-id/' def locate(_cur: Any) -> Tuple[Set[str], Set[str]]: - return {'`keep_test`', '`stale_test`'}, set() + return set(), set() app._show_create_service_url = show # type: ignore[method-assign] app._locate_app_functions = locate # type: ignore[method-assign] - class _Cursor: - def execute(self, sql: str) -> None: - executed.append(sql) + with self._patch_connect(executed): + with self.assertRaises(RuntimeError): + app.register_interactive_functions() + self.assertEqual(executed, []) - class _CM: - def __init__(self, inner: Any) -> None: - self.inner = inner + def test_register_interactive_drops_stale(self) -> None: + app = self._interactive_app(['keep_test']) + executed: List[str] = [] - def __enter__(self) -> Any: - return self.inner + def show(_cur: Any, _sql_name: str) -> Optional[str]: + return None - def __exit__(self, *args: Any) -> bool: - return False + def locate(_cur: Any) -> Tuple[Set[str], Set[str]]: + return {'`keep_test`', '`stale_test`'}, set() - class _Conn: - def cursor(self) -> _CM: - return _CM(_Cursor()) + def owned(_cur: Any, qualified: str) -> Optional[str]: + if 'stale_test' in qualified: + return app.url + return None - with patch( - 'singlestoredb.functions.ext.asgi.connection.connect', - return_value=_CM(_Conn()), - ): + app._show_create_service_url = show # type: ignore[method-assign] + app._locate_app_functions = locate # type: ignore[method-assign] + app._service_url_for_qualified = owned # type: ignore[method-assign] + + with self._patch_connect(executed): app.register_interactive_functions() self.assertIn('DROP FUNCTION IF EXISTS `stale_test`', executed) self.assertTrue( @@ -260,6 +273,163 @@ def cursor(self) -> _CM: any('DROP' in sql and 'keep_test' in sql for sql in executed), ) + def test_register_interactive_skips_stolen_stale_drop(self) -> None: + app = self._interactive_app(['keep_test']) + executed: List[str] = [] + + def show(_cur: Any, _sql_name: str) -> Optional[str]: + return None + + def locate(_cur: Any) -> Tuple[Set[str], Set[str]]: + return {'`keep_test`', '`stale_test`'}, set() + + def owned(_cur: Any, qualified: str) -> Optional[str]: + if 'stale_test' in qualified: + return 'https://gw/pythonudfs/published-id/' + return None + + app._show_create_service_url = show # type: ignore[method-assign] + app._locate_app_functions = locate # type: ignore[method-assign] + app._service_url_for_qualified = owned # type: ignore[method-assign] + + with self._patch_connect(executed): + app.register_interactive_functions() + self.assertFalse(any('DROP' in sql for sql in executed)) + self.assertTrue( + any('CREATE' in sql and 'keep_test' in sql for sql in executed), + ) + + +class TestRunUdfAppInteractiveOrder(unittest.IsolatedAsyncioTestCase): + + def tearDown(self) -> None: + from singlestoredb.apps import _python_udfs + _python_udfs._running_server = None + + def _interactive_config(self) -> AppConfig: + return AppConfig( + listen_port=8000, + base_url='https://example/', + base_path='/', + notebook_server_id='nb1', + app_token='t', + user_token=None, + running_interactively=True, + is_gateway_enabled=True, + is_local_dev=False, + ) + + def _mock_app(self, order: List[str], *, refuse: bool) -> MagicMock: + app = MagicMock() + app.endpoints = { + b'foo': (None, {'signature': {'name': 'foo_test'}}), + } + app.get_uvicorn_log_config.return_value = {} + app.get_function_info.return_value = {} + + def preflight(*_a: Any, **_k: Any) -> None: + order.append('preflight') + if refuse: + raise RuntimeError('stolen') + + def register(*_a: Any, **_k: Any) -> None: + order.append('register') + + app.preflight_interactive_functions.side_effect = preflight + app.register_interactive_functions.side_effect = register + return app + + def _uvicorn_modules(self, server_cls: Any = None) -> Any: + fake_util = MagicMock() + if server_cls is not None: + fake_util.AwaitableUvicornServer = server_cls + return patch.dict( + sys.modules, + { + 'uvicorn': MagicMock(), + 'singlestoredb.apps._uvicorn_util': fake_util, + }, + ) + + async def test_preflight_failure_leaves_server_running(self) -> None: + from singlestoredb.apps import _python_udfs + from singlestoredb.apps._python_udfs import run_udf_app + + order: List[str] = [] + existing = MagicMock() + + async def shutdown() -> None: + order.append('shutdown') + + existing.shutdown.side_effect = shutdown + _python_udfs._running_server = existing + app = self._mock_app(order, refuse=True) + + with self._uvicorn_modules(): + with patch.object( + AppConfig, 'from_env', return_value=self._interactive_config(), + ): + with patch.object( + _python_udfs, + 'generate_base_url', + return_value='https://gw/pythonudfs/sess/interactive/', + ): + with patch.object(_python_udfs, 'Application', return_value=app): + with patch.object( + _python_udfs, + 'kill_process_by_port', + side_effect=lambda _port: order.append('kill'), + ): + with self.assertRaises(RuntimeError): + await run_udf_app() + self.assertEqual(order, ['preflight']) + self.assertIs(_python_udfs._running_server, existing) + + async def test_preflight_then_shutdown_then_register(self) -> None: + from singlestoredb.apps import _python_udfs + from singlestoredb.apps._python_udfs import run_udf_app + + order: List[str] = [] + existing = MagicMock() + + async def shutdown() -> None: + order.append('shutdown') + + existing.shutdown.side_effect = shutdown + _python_udfs._running_server = existing + app = self._mock_app(order, refuse=False) + + class FakeServer: + def __init__(self, _config: Any) -> None: + pass + + async def serve(self) -> None: + return None + + async def wait_for_startup(self) -> None: + order.append('started') + + with self._uvicorn_modules(FakeServer): + with patch.object( + AppConfig, 'from_env', return_value=self._interactive_config(), + ): + with patch.object( + _python_udfs, + 'generate_base_url', + return_value='https://gw/pythonudfs/sess/interactive/', + ): + with patch.object(_python_udfs, 'Application', return_value=app): + with patch.object( + _python_udfs, + 'kill_process_by_port', + side_effect=lambda _port: order.append('kill'), + ): + await run_udf_app() + self.assertEqual( + order, + ['preflight', 'shutdown', 'kill', 'register', 'started'], + ) + if __name__ == '__main__': unittest.main()