1# Licensed to the Apache Software Foundation (ASF) under one
2# or more contributor license agreements. See the NOTICE file
3# distributed with this work for additional information
4# regarding copyright ownership. The ASF licenses this file
5# to you under the Apache License, Version 2.0 (the
6# "License"); you may not use this file except in compliance
7# with the License. You may obtain a copy of the License at
8#
9# http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing,
12# software distributed under the License is distributed on an
13# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14# KIND, either express or implied. See the License for the
15# specific language governing permissions and limitations
16# under the License.
17
18from __future__ import annotations
19
20import contextlib
21import logging
22import traceback
23from typing import TYPE_CHECKING
24
25from airflow import settings
26from airflow.api_internal.internal_api_call import InternalApiConfig
27from airflow.settings import TracebackSession
28
29if TYPE_CHECKING:
30 from sqlalchemy.orm import Session
31
32__current_task_instance_session: Session | None = None
33
34log = logging.getLogger(__name__)
35
36
37def get_current_task_instance_session() -> Session:
38 global __current_task_instance_session
39 if not __current_task_instance_session:
40 if InternalApiConfig.get_use_internal_api():
41 __current_task_instance_session = TracebackSession()
42 return __current_task_instance_session
43 log.warning("No task session set for this task. Continuing but this likely causes a resource leak.")
44 log.warning("Please report this and stacktrace below to https://github.com/apache/airflow/issues")
45 for filename, line_number, name, line in traceback.extract_stack():
46 log.warning('File: "%s", %s , in %s', filename, line_number, name)
47 if line:
48 log.warning(" %s", line.strip())
49 __current_task_instance_session = settings.Session()
50 return __current_task_instance_session
51
52
53@contextlib.contextmanager
54def set_current_task_instance_session(session: Session):
55 global __current_task_instance_session
56 if __current_task_instance_session:
57 raise RuntimeError(
58 "Session already set for this task. "
59 "You can only have one 'set_current_task_session' context manager active at a time."
60 )
61 __current_task_instance_session = session
62 try:
63 yield
64 finally:
65 __current_task_instance_session = None