odps/ipython/magics.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Copyright 1999-2025 Alibaba Group Holding Ltd.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import logging
import time

from .. import ODPS, options
from .. import types as odps_types
from ..compat import StringIO, six
from ..df import DataFrame, Delay, NullScalar, Scalar
from ..df.backends.frame import ResultFrame
from ..df.backends.odpssql.types import odps_schema_to_df_schema, odps_type_to_df_type
from ..inter import enter, list_rooms, setup, teardown
from ..models import TableSchema
from ..ui.common import html_notify
from ..ui.progress import (
    create_instance_group,
    fetch_instance_group,
    reload_instance_status,
)
from ..utils import (
    init_progress_ui,
    replace_sql_parameters,
    split_backquoted,
    strip_backquotes,
)

logger = logging.getLogger(__name__)

try:
    import numpy as np

    np_int_types = map(np.dtype, [np.int_, np.int8, np.int16, np.int32, np.int64])
    np_float_types = map(np.dtype, [np.float_, np.float16, np.float32, np.float64])
    np_to_odps_types = dict(
        [(t, odps_types.bigint) for t in np_int_types]
        + [(t, odps_types.double) for t in np_float_types]
    )
except ImportError:
    pass

try:
    from IPython.core.magic import Magics, line_cell_magic, line_magic, magics_class
except ImportError:
    # skipped for ci
    ODPSSql = None
    pass
else:

    @magics_class
    class ODPSSql(Magics):
        _odps = None

        def _set_odps(self):
            if self._odps is not None:
                return

            if options.account is not None and options.default_project is not None:
                self._odps = ODPS._from_account(
                    options.account,
                    options.default_project,
                    endpoint=options.endpoint,
                    tunnel_endpoint=options.tunnel.endpoint,
                )
            else:
                self._odps = enter().odps

        @line_magic("enter")
        def enter(self, line):
            room = line.strip()
            if room:
                r = enter(room)
                self._odps = r.odps
            else:
                r = enter()
                self._odps = r.odps

            if "o" not in self.shell.user_ns:
                self.shell.user_ns["o"] = self._odps
                self.shell.user_ns["odps"] = self._odps

            return r

        @line_magic("setup")
        def setup(self, line):
            args = line.strip().split()
            name, args = args[0], args[1:]
            setup(*args, room=name)
            html_notify("Setup succeeded")

        @line_magic("teardown")
        def teardown(self, line):
            name = line.strip()
            teardown(name)
            html_notify("Teardown succeeded")

        @line_magic("list_rooms")
        def list_rooms(self, line):
            return list_rooms()

        @line_magic("stores")
        def list_stores(self, line):
            line = line.strip()

            if line:
                room = enter(line)
            else:
                room = enter()

            return room.display()

        @staticmethod
        def _get_task_percent(task_progress):
            if len(task_progress.stages) > 0:
                all_percent = sum(
                    (float(stage.terminated_workers) / stage.total_workers)
                    for stage in task_progress.stages
                    if stage.total_workers > 0
                )
                return all_percent / len(task_progress.stages)
            else:
                return 0

        def _to_stdout(cls, msg):
            print(msg)

        @line_magic("set")
        def set_hint(self, line):
            if "=" not in line:
                raise ValueError("Hint for sql is not allowed")

            key, val = line.strip().strip(";").split("=", 1)
            key, val = key.strip(), val.strip()

            settings = options.sql.settings
            if settings is None:
                options.sql.settings = {key: val}
            else:
                options.sql.settings[key] = val

        @line_cell_magic("sql")
        def execute(self, line, cell=""):
            self._set_odps()

            content = line + "\n" + cell
            content = content.strip()

            sql = None
            hints = dict()

            splits = content.split(";")
            for s in splits:
                stripped = s.strip()
                if stripped.lower().startswith("set "):
                    hint = stripped.split(" ", 1)[1]
                    k, v = hint.split("=", 1)
                    k, v = k.strip(), v.strip()
                    hints[k] = v
                elif len(stripped) == 0:
                    continue
                else:
                    if sql is None:
                        sql = s
                    else:
                        sql = "%s;%s" % (sql, s)

            # replace user defined parameters
            sql = replace_sql_parameters(sql, self.shell.user_ns)

            if sql:
                progress_ui = init_progress_ui()
                group_id = create_instance_group("SQL Query")
                progress_ui.add_keys(group_id)

                instance = self._odps.run_sql(sql, hints=hints)
                if logger.getEffectiveLevel() <= logging.INFO:
                    logger.info(
                        "Instance ID: %s\n  Log view: %s",
                        instance.id,
                        instance.get_logview_address(),
                    )
                reload_instance_status(self._odps, group_id, instance.id)
                progress_ui.status("Executing")

                percent = 0
                while not instance.is_terminated(retry=True):
                    last_percent = percent

                    reload_instance_status(self._odps, group_id, instance.id)
                    inst_progress = fetch_instance_group(group_id).instances.get(
                        instance.id
                    )

                    if inst_progress is not None and len(inst_progress.tasks) > 0:
                        percent = sum(
                            self._get_task_percent(task)
                            for task in six.itervalues(inst_progress.tasks)
                        ) / len(inst_progress.tasks)
                    else:
                        percent = 0

                    percent = min(1, max(percent, last_percent))

                    progress_ui.update(percent)
                    progress_ui.update_group()

                    time.sleep(1)

                instance.wait_for_success()
                progress_ui.update(1)

                try:
                    with instance.open_reader() as reader:
                        try:
                            import pandas as pd

                            try:
                                from pandas.io.parsers import (
                                    ParserError as CParserError,
                                )
                            except ImportError:
                                pass
                            try:
                                from pandas.parser import CParserError  # noqa
                            except ImportError:
                                CParserError = ValueError  # noqa

                            if not hasattr(reader, "raw"):
                                res = ResultFrame(
                                    [rec.values for rec in reader],
                                    schema=odps_schema_to_df_schema(reader._schema),
                                )
                            else:
                                try:
                                    res = pd.read_csv(StringIO(reader.raw))
                                    if len(res.values) > 0:
                                        schema = DataFrame(res).schema
                                    else:
                                        cols = res.columns.tolist()
                                        schema = odps_schema_to_df_schema(
                                            TableSchema.from_lists(
                                                cols, ["string" for _ in cols]
                                            )
                                        )
                                    res = ResultFrame(res.values, schema=schema)
                                except (ValueError, CParserError):
                                    res = reader.raw
                        except (ImportError, ValueError):
                            if not hasattr(reader, "raw"):
                                res = ResultFrame(
                                    [rec.values for rec in reader],
                                    schema=odps_schema_to_df_schema(reader._schema),
                                )
                            else:
                                try:
                                    columns = [
                                        odps_types.Column(
                                            name=col.name,
                                            typo=odps_type_to_df_type(col.type),
                                        )
                                        for col in reader._columns
                                    ]
                                    res = ResultFrame(list(reader), columns=columns)
                                except TypeError:
                                    res = reader.raw

                    html_notify("SQL execution succeeded")
                    return res
                finally:
                    progress_ui.close()

        @line_magic("persist")
        def persist(self, line):
            try:
                import pandas as pd

                has_pandas = True
            except (ImportError, ValueError):
                has_pandas = False

            self._set_odps()

            line = line.strip().strip(";")

            frame_name, table_name = line.split(None, 1)

            if "." in table_name:
                parts = split_backquoted(table_name, ".")
                if len(parts) == 3:
                    project_name, schema_name, table_name = parts
                else:
                    project_name, table_name = parts
                    schema_name = None
                table_name = strip_backquotes(table_name)
            else:
                project_name = schema_name = None

            frame = self.shell.user_ns[frame_name]
            if self._odps.exist_table(
                table_name, project=project_name, schema=schema_name
            ):
                raise TypeError("%s already exists" % table_name)

            if isinstance(frame, DataFrame):
                frame.persist(
                    name=table_name,
                    project=project_name,
                    schema=schema_name,
                    notify=False,
                )
            elif has_pandas and isinstance(frame, pd.DataFrame):
                frame = DataFrame(frame)
                frame.persist(
                    name=table_name,
                    project=project_name,
                    schema=schema_name,
                    notify=False,
                )
            html_notify("Persist succeeded")

    def load_ipython_extension(ipython):
        ipython.register_magics(ODPSSql)

        # Do global import when load extension
        ipython.user_ns["DataFrame"] = DataFrame
        ipython.user_ns["Scalar"] = Scalar
        ipython.user_ns["NullScalar"] = NullScalar
        ipython.user_ns["options"] = options
        ipython.user_ns["TableSchema"] = TableSchema
        ipython.user_ns["Delay"] = Delay