singlestoredb 1.15.0__cp38-abi3-win_amd64.whl → 1.15.2__cp38-abi3-win_amd64.whl

singlestoredb/functions/ext/asgi.py

@@ -73,6 +73,8 @@ from ..signature import signature_to_sql
 from ..typing import Masked
 from ..typing import Table
 from .timer import Timer
+from singlestoredb.docstring.parser import parse
+from singlestoredb.functions.dtypes import escape_name
 
 try:
     import cloudpickle
@@ -89,7 +91,6 @@ except ImportError:
 
 logger = utils.get_logger('singlestoredb.functions.ext.asgi')
 
-
 # If a number of processes is specified, create a pool of workers
 num_processes = max(0, int(os.environ.get('SINGLESTOREDB_EXT_NUM_PROCESSES', 0)))
 if num_processes > 1:
@@ -538,6 +539,8 @@ def make_func(
         Name of the function to create
     func : Callable
         The function to call as the endpoint
+    database : str, optional
+        The database to use for the function definition
 
     Returns
     -------
@@ -615,7 +618,7 @@ async def cancel_on_disconnect(
     """Cancel request if client disconnects."""
     while True:
         message = await receive()
-        if message['type'] == 'http.disconnect':
+        if message.get('type', '') == 'http.disconnect':
             raise asyncio.CancelledError(
                 'Function call was cancelled by client',
             )
@@ -674,6 +677,24 @@ class Application(object):
     link_credentials : Dict[str, Any], optional
         The CREDENTIALS section of a LINK definition. This dictionary gets
         converted to JSON for the CREATE LINK call.
+    name_prefix : str, optional
+        Prefix to add to function names when registering with the database
+    name_suffix : str, optional
+        Suffix to add to function names when registering with the database
+    function_database : str, optional
+        The database to use for external function definitions.
+    log_file : str, optional
+        File path to write logs to instead of console. If None, logs are
+        written to console. When specified, application logger handlers
+        are replaced with a file handler.
+    log_level : str, optional
+        Logging level for the application logger. Valid values are 'info',
+        'debug', 'warning', 'error'. Defaults to 'info'.
+    disable_metrics : bool, optional
+        Disable logging of function call metrics. Defaults to False.
+    app_name : str, optional
+        Name for the application instance. Used to create a logger-specific
+        name. If not provided, a random name will be generated.
 
     """
 
@@ -816,6 +837,7 @@ class Application(object):
     invoke_path = ('invoke',)
     show_create_function_path = ('show', 'create_function')
     show_function_info_path = ('show', 'function_info')
+    status = ('status',)
 
     def __init__(
         self,
@@ -838,6 +860,11 @@ class Application(object):
         link_credentials: Optional[Dict[str, Any]] = None,
         name_prefix: str = get_option('external_function.name_prefix'),
         name_suffix: str = get_option('external_function.name_suffix'),
+        function_database: Optional[str] = None,
+        log_file: Optional[str] = get_option('external_function.log_file'),
+        log_level: str = get_option('external_function.log_level'),
+        disable_metrics: bool = get_option('external_function.disable_metrics'),
+        app_name: Optional[str] = get_option('external_function.app_name'),
     ) -> None:
         if link_name and (link_config or link_credentials):
             raise ValueError(
@@ -854,6 +881,15 @@ class Application(object):
                 get_option('external_function.link_credentials') or '{}',
             ) or None
 
+        # Generate application name if not provided
+        if app_name is None:
+            app_name = f'udf_app_{secrets.token_hex(4)}'
+
+        self.name = app_name
+
+        # Create logger instance specific to this application
+        self.logger = utils.get_logger(f'singlestoredb.functions.ext.asgi.{self.name}')
+
         # List of functions specs
         specs: List[Union[str, Callable[..., Any], ModuleType]] = []
 
@@ -944,6 +980,98 @@ class Application(object):
         self.link_credentials = link_credentials
         self.endpoints = endpoints
         self.external_functions = external_functions
+        self.function_database = function_database
+        self.log_file = log_file
+        self.log_level = log_level
+        self.disable_metrics = disable_metrics
+
+        # Configure logging
+        self._configure_logging()
+
+    def _configure_logging(self) -> None:
+        """Configure logging based on the log_file settings."""
+        # Set logger level
+        self.logger.setLevel(getattr(logging, self.log_level.upper()))
+
+        # Remove all existing handlers to ensure clean configuration
+        self.logger.handlers.clear()
+
+        # Configure log file if specified
+        if self.log_file:
+            # Create file handler
+            file_handler = logging.FileHandler(self.log_file)
+            file_handler.setLevel(getattr(logging, self.log_level.upper()))
+
+            # Use JSON formatter for file logging
+            formatter = utils.JSONFormatter()
+            file_handler.setFormatter(formatter)
+
+            # Add the handler to the logger
+            self.logger.addHandler(file_handler)
+        else:
+            # For console logging, create a new stream handler with JSON formatter
+            console_handler = logging.StreamHandler()
+            console_handler.setLevel(getattr(logging, self.log_level.upper()))
+            console_handler.setFormatter(utils.JSONFormatter())
+            self.logger.addHandler(console_handler)
+
+        # Prevent propagation to avoid duplicate or differently formatted messages
+        self.logger.propagate = False
+
+    def get_uvicorn_log_config(self) -> Dict[str, Any]:
+        """
+        Create uvicorn log config that matches the Application's logging format.
+
+        This method returns the log configuration used by uvicorn, allowing external
+        users to match the logging format of the Application class.
+
+        Returns
+        -------
+        Dict[str, Any]
+            Log configuration dictionary compatible with uvicorn's log_config parameter
+
+        """
+        log_config = {
+            'version': 1,
+            'disable_existing_loggers': False,
+            'formatters': {
+                'json': {
+                    '()': 'singlestoredb.functions.ext.utils.JSONFormatter',
+                },
+            },
+            'handlers': {
+                'default': {
+                    'class': (
+                        'logging.FileHandler' if self.log_file
+                        else 'logging.StreamHandler'
+                    ),
+                    'formatter': 'json',
+                },
+            },
+            'loggers': {
+                'uvicorn': {
+                    'handlers': ['default'],
+                    'level': self.log_level.upper(),
+                    'propagate': False,
+                },
+                'uvicorn.error': {
+                    'handlers': ['default'],
+                    'level': self.log_level.upper(),
+                    'propagate': False,
+                },
+                'uvicorn.access': {
+                    'handlers': ['default'],
+                    'level': self.log_level.upper(),
+                    'propagate': False,
+                },
+            },
+        }
+
+        # Add filename to file handler if log file is specified
+        if self.log_file:
+            log_config['handlers']['default']['filename'] = self.log_file  # type: ignore
+
+        return log_config
 
     async def __call__(
         self,
@@ -967,19 +1095,22 @@ class Application(object):
         request_id = str(uuid.uuid4())
 
         timer = Timer(
+            app_name=self.name,
             id=request_id,
             timestamp=datetime.datetime.now(
                 datetime.timezone.utc,
             ).strftime('%Y-%m-%dT%H:%M:%S.%fZ'),
         )
         call_timer = Timer(
+            app_name=self.name,
             id=request_id,
             timestamp=datetime.datetime.now(
                 datetime.timezone.utc,
             ).strftime('%Y-%m-%dT%H:%M:%S.%fZ'),
         )
 
-        assert scope['type'] == 'http'
+        if scope['type'] != 'http':
+            raise ValueError(f"Expected HTTP scope, got {scope['type']}")
 
         method = scope['method']
         path = tuple(x for x in scope['path'].split('/') if x)
@@ -992,6 +1123,7 @@ class Application(object):
         accepts = headers.get(b'accepts', content_type)
         func_name = headers.get(b's2-ef-name', b'')
         func_endpoint = self.endpoints.get(func_name)
+        ignore_cancel = headers.get(b's2-ef-ignore-cancel', b'false') == b'true'
 
         timer.metadata['function'] = func_name.decode('utf-8') if func_name else ''
         call_timer.metadata['function'] = timer.metadata['function']
@@ -1004,14 +1136,15 @@ class Application(object):
         # Call the endpoint
         if method == 'POST' and func is not None and path == self.invoke_path:
 
-            logger.info(
-                json.dumps({
-                    'type': 'function_call',
-                    'id': request_id,
-                    'name': func_name.decode('utf-8'),
+            self.logger.info(
+                'Function call initiated',
+                extra={
+                    'app_name': self.name,
+                    'request_id': request_id,
+                    'function_name': func_name.decode('utf-8'),
                     'content_type': content_type.decode('utf-8'),
                     'accepts': accepts.decode('utf-8'),
-                }),
+                },
             )
 
             args_data_format = func_info['args_data_format']
@@ -1021,7 +1154,7 @@ class Application(object):
             with timer('receive_data'):
                 while more_body:
                     request = await receive()
-                    if request['type'] == 'http.disconnect':
+                    if request.get('type', '') == 'http.disconnect':
                         raise RuntimeError('client disconnected')
                     data.append(request['body'])
                     more_body = request.get('more_body', False)
@@ -1051,7 +1184,8 @@ class Application(object):
                     ),
                 )
                 disconnect_task = asyncio.create_task(
-                    cancel_on_disconnect(receive),
+                    asyncio.sleep(int(1e9))
+                    if ignore_cancel else cancel_on_disconnect(receive),
                 )
                 timeout_task = asyncio.create_task(
                     cancel_on_timeout(func_info['timeout']),
@@ -1090,8 +1224,14 @@ class Application(object):
                 await send(output_handler['response'])
 
             except asyncio.TimeoutError:
-                logging.exception(
-                    'Timeout in function call: ' + func_name.decode('utf-8'),
+                self.logger.exception(
+                    'Function call timeout',
+                    extra={
+                        'app_name': self.name,
+                        'request_id': request_id,
+                        'function_name': func_name.decode('utf-8'),
+                        'timeout': func_info['timeout'],
+                    },
                 )
                 body = (
                     '[TimeoutError] Function call timed out after ' +
@@ -1101,15 +1241,26 @@ class Application(object):
                 await send(self.error_response_dict)
 
             except asyncio.CancelledError:
-                logging.exception(
-                    'Function call cancelled: ' + func_name.decode('utf-8'),
+                self.logger.exception(
+                    'Function call cancelled',
+                    extra={
+                        'app_name': self.name,
+                        'request_id': request_id,
+                        'function_name': func_name.decode('utf-8'),
+                    },
                 )
                 body = b'[CancelledError] Function call was cancelled'
                 await send(self.error_response_dict)
 
             except Exception as e:
-                logging.exception(
-                    'Error in function call: ' + func_name.decode('utf-8'),
+                self.logger.exception(
+                    'Function call error',
+                    extra={
+                        'app_name': self.name,
+                        'request_id': request_id,
+                        'function_name': func_name.decode('utf-8'),
+                        'exception_type': type(e).__name__,
+                    },
                 )
                 body = f'[{type(e).__name__}] {str(e).strip()}'.encode('utf-8')
                 await send(self.error_response_dict)
@@ -1130,6 +1281,7 @@ class Application(object):
                             endpoint_info['signature'],
                             url=self.url or reflected_url,
                             data_format=self.data_format,
+                            database=self.function_database or None,
                         ),
                     )
             body = '\n'.join(syntax).encode('utf-8')
@@ -1142,6 +1294,11 @@ class Application(object):
             body = json.dumps(dict(functions=functions)).encode('utf-8')
             await send(self.text_response_dict)
 
+        # Return status
+        elif method == 'GET' and path == self.status:
+            body = json.dumps(dict(status='ok')).encode('utf-8')
+            await send(self.text_response_dict)
+
         # Path not found
         else:
             body = b''
@@ -1156,7 +1313,17 @@ class Application(object):
         for k, v in call_timer.metrics.items():
             timer.metrics[k] = v
 
-        timer.finish()
+        if not self.disable_metrics:
+            metrics = timer.finish()
+            self.logger.info(
+                'Function call metrics',
+                extra={
+                    'app_name': self.name,
+                    'request_id': request_id,
+                    'function_name': timer.metadata.get('function', ''),
+                    'metrics': metrics,
+                },
+            )
 
     def _create_link(
         self,
@@ -1184,20 +1351,27 @@ class Application(object):
     def _locate_app_functions(self, cur: Any) -> Tuple[Set[str], Set[str]]:
         """Locate all current functions and links belonging to this app."""
         funcs, links = set(), set()
-        cur.execute('SHOW FUNCTIONS')
+        if self.function_database:
+            database_prefix = escape_name(self.function_database) + '.'
+            cur.execute(f'SHOW FUNCTIONS IN {escape_name(self.function_database)}')
+        else:
+            database_prefix = ''
+            cur.execute('SHOW FUNCTIONS')
+
         for row in list(cur):
             name, ftype, link = row[0], row[1], row[-1]
             # Only look at external functions
             if 'external' not in ftype.lower():
                 continue
             # See if function URL matches url
-            cur.execute(f'SHOW CREATE FUNCTION `{name}`')
+            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:
-                    funcs.add(fname)
+                    funcs.add(f'{database_prefix}{escape_name(fname)}')
                     if link and re.match(r'^py_ext_func_link_\S{14}$', link):
                         links.add(link)
+
         return funcs, links
 
     def get_function_info(
@@ -1206,9 +1380,11 @@ class Application(object):
     ) -> Dict[str, Any]:
         """
         Return the functions and function signature information.
+
         Returns
         -------
         Dict[str, Any]
+
         """
         functions = {}
         no_default = object()
@@ -1220,20 +1396,71 @@ class Application(object):
             sig = info['signature']
             sql_map[sig['name']] = sql
 
-        for key, (_, info) in self.endpoints.items():
+        for key, (func, info) in self.endpoints.items():
+            # Get info from docstring
+            doc_summary = ''
+            doc_long_description = ''
+            doc_params = {}
+            doc_returns = None
+            doc_examples = []
+            if func.__doc__:
+                try:
+                    docs = parse(func.__doc__)
+                    doc_params = {p.arg_name: p for p in docs.params}
+                    doc_returns = docs.returns
+                    if not docs.short_description and docs.long_description:
+                        doc_summary = docs.long_description or ''
+                    else:
+                        doc_summary = docs.short_description or ''
+                        doc_long_description = docs.long_description or ''
+                    for ex in docs.examples:
+                        ex_dict: Dict[str, Any] = {
+                            'description': None,
+                            'code': None,
+                            'output': None,
+                        }
+                        if ex.description:
+                            ex_dict['description'] = ex.description
+                        if ex.snippet:
+                            code, output = [], []
+                            for line in ex.snippet.split('\n'):
+                                line = line.rstrip()
+                                if re.match(r'^(\w+>|>>>|\.\.\.)', line):
+                                    code.append(line)
+                                else:
+                                    output.append(line)
+                            ex_dict['code'] = '\n'.join(code) or None
+                            ex_dict['output'] = '\n'.join(output) or None
+                        if ex.post_snippet:
+                            ex_dict['postscript'] = ex.post_snippet
+                        doc_examples.append(ex_dict)
+
+                except Exception as e:
+                    self.logger.warning(
+                        'Could not parse docstring for function',
+                        extra={
+                            'app_name': self.name,
+                            'function_name': key.decode('utf-8'),
+                            'error': str(e),
+                        },
+                    )
+
             if not func_name or key == func_name:
                 sig = info['signature']
                 args = []
 
                 # Function arguments
-                for a in sig.get('args', []):
+                for i, a in enumerate(sig.get('args', [])):
+                    name = a['name']
                     dtype = a['dtype']
                     nullable = '?' in dtype
                     args.append(
                         dict(
-                            name=a['name'],
+                            name=name,
                             dtype=dtype.replace('?', ''),
                             nullable=nullable,
+                            description=(doc_params[name].description or '')
+                            if name in doc_params else '',
                         ),
                     )
                     if a.get('default', no_default) is not no_default:
@@ -1250,6 +1477,8 @@ class Application(object):
                         dict(
                             dtype=dtype.replace('?', ''),
                             nullable=nullable,
+                            description=doc_returns.description
+                            if doc_returns else '',
                         ),
                     )
                     if a.get('name', None):
@@ -1263,6 +1492,9 @@ class Application(object):
                     returns=returns,
                     function_type=info['function_type'],
                     sql_statement=sql,
+                    summary=doc_summary,
+                    long_description=doc_long_description,
+                    examples=doc_examples,
                 )
 
         return functions
@@ -1303,6 +1535,7 @@ class Application(object):
                     app_mode=self.app_mode,
                     replace=replace,
                     link=link or None,
+                    database=self.function_database or None,
                 ),
             )
 
@@ -1332,7 +1565,7 @@ class Application(object):
                 if replace:
                     funcs, links = self._locate_app_functions(cur)
                     for fname in funcs:
-                        cur.execute(f'DROP FUNCTION IF EXISTS `{fname}`')
+                        cur.execute(f'DROP FUNCTION IF EXISTS {fname}')
                     for link in links:
                         cur.execute(f'DROP LINK {link}')
                 for func in self.get_create_functions(replace=replace):
@@ -1358,7 +1591,7 @@ class Application(object):
             with conn.cursor() as cur:
                 funcs, links = self._locate_app_functions(cur)
                 for fname in funcs:
-                    cur.execute(f'DROP FUNCTION IF EXISTS `{fname}`')
+                    cur.execute(f'DROP FUNCTION IF EXISTS {fname}')
                 for link in links:
                     cur.execute(f'DROP LINK {link}')
 
@@ -1415,6 +1648,7 @@ class Application(object):
                 b'accepts': accepts[data_format.lower()],
                 b's2-ef-name': name.encode('utf-8'),
                 b's2-ef-version': data_version.encode('utf-8'),
+                b's2-ef-ignore-cancel': b'true',
             },
         )
 
@@ -1663,6 +1897,22 @@ def main(argv: Optional[List[str]] = None) -> None:
             ),
             help='logging level',
         )
+        parser.add_argument(
+            '--log-file', metavar='filepath',
+            default=defaults.get(
+                'log_file',
+                get_option('external_function.log_file'),
+            ),
+            help='File path to write logs to instead of console',
+        )
+        parser.add_argument(
+            '--disable-metrics', action='store_true',
+            default=defaults.get(
+                'disable_metrics',
+                get_option('external_function.disable_metrics'),
+            ),
+            help='Disable logging of function call metrics',
+        )
         parser.add_argument(
             '--name-prefix', metavar='name_prefix',
             default=defaults.get(
@@ -1679,6 +1929,22 @@ def main(argv: Optional[List[str]] = None) -> None:
             ),
             help='Suffix to add to function names',
         )
+        parser.add_argument(
+            '--function-database', metavar='function_database',
+            default=defaults.get(
+                'function_database',
+                get_option('external_function.function_database'),
+            ),
+            help='Database to use for the function definition',
+        )
+        parser.add_argument(
+            '--app-name', metavar='app_name',
+            default=defaults.get(
+                'app_name',
+                get_option('external_function.app_name'),
+            ),
+            help='Name for the application instance',
+        )
         parser.add_argument(
             'functions', metavar='module.or.func.path', nargs='*',
             help='functions or modules to export in UDF server',
@@ -1686,8 +1952,6 @@ def main(argv: Optional[List[str]] = None) -> None:
 
         args = parser.parse_args(argv)
 
-        logger.setLevel(getattr(logging, args.log_level.upper()))
-
         if i > 0:
             break
 
@@ -1778,6 +2042,11 @@ def main(argv: Optional[List[str]] = None) -> None:
         app_mode='remote',
         name_prefix=args.name_prefix,
         name_suffix=args.name_suffix,
+        function_database=args.function_database or None,
+        log_file=args.log_file,
+        log_level=args.log_level,
+        disable_metrics=args.disable_metrics,
+        app_name=args.app_name,
     )
 
     funcs = app.get_create_functions(replace=args.replace_existing)
@@ -1785,11 +2054,11 @@ def main(argv: Optional[List[str]] = None) -> None:
         raise RuntimeError('no functions specified')
 
     for f in funcs:
-        logger.info(f)
+        app.logger.info(f)
 
     try:
         if args.db:
-            logger.info('registering functions with database')
+            app.logger.info('Registering functions with database')
             app.register_functions(
                 args.db,
                 replace=args.replace_existing,
@@ -1804,6 +2073,9 @@ def main(argv: Optional[List[str]] = None) -> None:
             ).items() if v is not None
         }
 
+        # Configure uvicorn logging to use JSON format matching Application's format
+        app_args['log_config'] = app.get_uvicorn_log_config()
+
         if use_async:
             asyncio.create_task(_run_uvicorn(uvicorn, app, app_args, db=args.db))
         else:
@@ -1811,7 +2083,7 @@ def main(argv: Optional[List[str]] = None) -> None:
 
     finally:
         if not use_async and args.db:
-            logger.info('dropping functions from database')
+            app.logger.info('Dropping functions from database')
             app.drop_functions(args.db)
 
 
@@ -1824,7 +2096,7 @@ async def _run_uvicorn(
     """Run uvicorn server and clean up functions after shutdown."""
     await uvicorn.Server(uvicorn.Config(app, **app_args)).serve()
     if db:
-        logger.info('dropping functions from database')
+        app.logger.info('Dropping functions from database')
         app.drop_functions(db)
 
 

singlestoredb/functions/ext/timer.py

@@ -4,10 +4,6 @@ from typing import Any
 from typing import Dict
 from typing import Optional
 
-from . import utils
-
-logger = utils.get_logger('singlestoredb.functions.ext.metrics')
-
 
 class RoundedFloatEncoder(json.JSONEncoder):
 
@@ -87,12 +83,7 @@ class Timer:
         self.entries.clear()
         self._current_key = None
 
-    def finish(self) -> None:
+    def finish(self) -> Dict[str, Any]:
         """Finish the current timing context and store the elapsed time."""
         self.metrics['total'] = time.perf_counter() - self.start_time
-        self.log_metrics()
-
-    def log_metrics(self) -> None:
-        if self.metadata.get('function'):
-            result = dict(type='function_metrics', **self.metadata, **self.metrics)
-            logger.info(json.dumps(result, cls=RoundedFloatEncoder))
+        return dict(type='function_metrics', **self.metadata, **self.metrics)