flink-python/pyflink/dataframe/context.py

################################################################################
#  Licensed to the Apache Software Foundation (ASF) under one
#  or more contributor license agreements.  See the NOTICE file
#  distributed with this work for additional information
#  regarding copyright ownership.  The ASF licenses this file
#  to you 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.
################################################################################

from typing import Optional

from pyflink.table import StreamTableEnvironment, TableEnvironment
from pyflink.util.api_stability_decorators import PublicEvolving

__all__ = [
    "set_table_environment",
    "get_table_environment",
    "get_or_create_table_environment",
]

_global_table_environment: Optional[TableEnvironment] = None


@PublicEvolving()
def set_table_environment(t_env: Optional[TableEnvironment]) -> None:
    """
    Set the environment used by DataFrame operations.

    :param t_env: Environment to use, or ``None`` to clear it.
    :raises TypeError: If ``t_env`` is neither a :class:`TableEnvironment` nor ``None``.

    Example::

        >>> import pyflink.dataframe as pf
        >>> pf.set_table_environment(None)
        >>> pf.get_table_environment() is None
        True

    .. versionadded:: 2.4.0
    """
    global _global_table_environment
    if t_env is not None and not isinstance(t_env, TableEnvironment):
        raise TypeError("t_env must be a TableEnvironment or None")
    _global_table_environment = t_env


@PublicEvolving()
def get_table_environment() -> Optional[TableEnvironment]:
    """
    Return the environment used by DataFrame operations, if one is configured.

    :return: The configured environment, or ``None``.

    Example::

        >>> import pyflink.dataframe as pf
        >>> pf.set_table_environment(None)
        >>> pf.get_table_environment() is None
        True

    .. versionadded:: 2.4.0
    """
    return _global_table_environment


@PublicEvolving()
def get_or_create_table_environment() -> TableEnvironment:
    """
    Return the configured environment, creating one when necessary.

    The created environment is retained for subsequent DataFrame operations and calls to
    :func:`get_table_environment`.

    :return: The configured or newly created environment.

    Example::

        >>> import pyflink.dataframe as pf
        >>> pf.set_table_environment(None)
        >>> environment = pf.get_or_create_table_environment()
        >>> pf.get_table_environment() is environment
        True

    .. versionadded:: 2.4.0
    """
    global _global_table_environment

    if _global_table_environment is None:
        from pyflink.datastream import StreamExecutionEnvironment

        stream_environment = StreamExecutionEnvironment.get_execution_environment()
        _global_table_environment = StreamTableEnvironment.create(stream_environment)

    return _global_table_environment