################################################################################
# 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.
################################################################################
"""
Global context for DataFrame operations.
This module provides global configuration for DataFrame operations.
"""
from collections.abc import Mapping as MappingABC
from typing import Any, Dict, List, Mapping, Optional, TYPE_CHECKING
from pyflink.common.configuration import Configuration
from pyflink.dataframe.util.artifacts import add_built_in_catalog
from pyflink.java_gateway import get_gateway
if TYPE_CHECKING:
from pyflink.table import StreamTableEnvironment
_global_t_env: Optional["StreamTableEnvironment"] = None
[docs]class DataFrameConfig:
"""
Configuration for DataFrame operations.
Stores key-value config pairs and applies them to the underlying
TableEnvironment. If a TableEnvironment is already set, configs are
applied immediately. Otherwise, they are buffered and applied when
the environment becomes available.
Example::
>>> import pyflink.dataframe as pf
>>> pf.config.set("parallelism.default", "4")
>>> pf.config.get("parallelism.default")
'4'
"""
def __init__(self):
self._buffer: Dict[str, str] = {}
[docs] def set(self, key: str, value: str) -> "DataFrameConfig":
"""
Set a configuration key-value pair.
Args:
key: The configuration key.
value: The configuration value.
Returns:
self, for chaining.
"""
self._buffer[key] = value
if _global_t_env is not None:
_global_t_env.get_config().set(key, value)
return self
[docs] def get(self, key: str, default: Optional[str] = None) -> Optional[str]:
"""
Get a configuration value.
If a TableEnvironment is set, reads from its config.
Otherwise, reads from the local buffer.
Args:
key: The configuration key.
default: Default value if key is not found.
Returns:
The configuration value, or default.
"""
if _global_t_env is not None:
table_config = _global_t_env.get_config()
if table_config.get_configuration().contains_key(key):
return table_config.get_configuration().get_string(key, default)
return self._buffer.get(key, default)
def _apply_to(self, t_env: "StreamTableEnvironment") -> None:
"""Apply all buffered configs to the given TableEnvironment."""
for key, value in self._buffer.items():
t_env.get_config().set(key, value)
def __repr__(self) -> str:
items = ", ".join(f"{k!r}: {v!r}" for k, v in self._buffer.items())
return f"DataFrameConfig({{{items}}})"
config = DataFrameConfig()
[docs]def set_table_environment(t_env: Optional["StreamTableEnvironment"]) -> None:
"""
Set the global TableEnvironment for DataFrame operations.
This is useful in testing scenarios where you want to reuse
the same TableEnvironment across multiple DataFrame creations.
Args:
t_env: The TableEnvironment to use globally. Pass None to reset.
Example::
>>> import pyflink.dataframe as pf
>>> # In test setup
>>> pf.set_table_environment(self.t_env)
>>> # Now from_records, from_dict will use this t_env
>>> df = pf.from_records([(1, 2)], schema=["a", "b"])
>>> # In test teardown
>>> pf.set_table_environment(None)
"""
global _global_t_env
_global_t_env = t_env
if t_env is not None:
config._apply_to(t_env)
[docs]def get_table_environment() -> Optional["StreamTableEnvironment"]:
"""
Get the currently set global TableEnvironment.
Returns:
The global TableEnvironment, or None if not set.
"""
return _global_t_env
[docs]def get_or_create_table_environment() -> "StreamTableEnvironment":
"""
Return the global TableEnvironment, creating and storing one if needed.
Buffered DataFrame config is applied when a new environment is created.
"""
t_env = get_table_environment()
if t_env is not None:
return t_env
from pyflink.datastream import StreamExecutionEnvironment
from pyflink.table import StreamTableEnvironment
env = StreamExecutionEnvironment.get_execution_environment()
t_env = StreamTableEnvironment.create(env)
set_table_environment(t_env)
return t_env
[docs]def create_catalog(name: str, options: Mapping[str, Any]) -> None:
"""
Create a catalog in the DataFrame default TableEnvironment.
Args:
name: Catalog name.
options: Catalog options. Keys and values are converted to strings.
Example::
>>> import pyflink.dataframe as pf
>>> pf.create_catalog("my_fluss", {
... "type": "fluss",
... "bootstrap.servers": "localhost:9123",
... },
... )
"""
if not isinstance(name, str) or not name.strip():
raise ValueError("catalog name must be a non-empty string.")
if not isinstance(options, MappingABC):
raise TypeError("options must be a mapping of catalog option key/value pairs.")
catalog_options = {
str(key): str(value)
for key, value in options.items()
if value is not None
}
t_env = get_or_create_table_environment()
catalog_type = catalog_options.get("type")
if catalog_type:
add_built_in_catalog(t_env, catalog_type)
configuration = Configuration()
for key, value in catalog_options.items():
configuration.set_string(key, value)
descriptor = (
get_gateway()
.jvm
.org
.apache
.flink
.table
.catalog
.CatalogDescriptor
.of(name, configuration._j_configuration)
)
t_env._j_tenv.createCatalog(name, descriptor)
[docs]def use_catalog(name: str) -> None:
"""
Set the current catalog for DataFrame operations.
Args:
name: Catalog name.
Example::
>>> import pyflink.dataframe as pf
>>> pf.use_catalog("my_fluss")
>>> pf.get_current_catalog()
'my_fluss'
"""
if not isinstance(name, str) or not name.strip():
raise ValueError("catalog name must be a non-empty string.")
get_or_create_table_environment().use_catalog(name)
[docs]def get_current_catalog() -> str:
"""
Get the current catalog used by DataFrame operations.
Returns:
The current catalog name.
Example::
>>> import pyflink.dataframe as pf
>>> pf.get_current_catalog()
'default_catalog'
"""
return get_or_create_table_environment().get_current_catalog()
[docs]def list_catalogs() -> List[str]:
"""
List catalogs in the DataFrame default TableEnvironment.
Returns:
Catalog names.
Example::
>>> import pyflink.dataframe as pf
>>> pf.list_catalogs()
['default_catalog']
"""
return get_or_create_table_environment().list_catalogs()
[docs]def use_database(name: str) -> None:
"""
Set the current database for DataFrame operations.
Args:
name: Database name.
Example::
>>> import pyflink.dataframe as pf
>>> pf.use_database("my_database")
>>> pf.get_current_database()
'my_database'
"""
if not isinstance(name, str) or not name.strip():
raise ValueError("database name must be a non-empty string.")
get_or_create_table_environment().use_database(name)
[docs]def get_current_database() -> str:
"""
Get the current database used by DataFrame operations.
Returns:
The current database name.
Example::
>>> import pyflink.dataframe as pf
>>> pf.get_current_database()
'default_database'
"""
return get_or_create_table_environment().get_current_database()
[docs]def list_databases() -> List[str]:
"""
List databases in the current catalog.
Returns:
Database names.
Example::
>>> import pyflink.dataframe as pf
>>> pf.list_databases()
['default_database']
"""
return get_or_create_table_environment().list_databases()