44 lines
1.3 KiB
Python
44 lines
1.3 KiB
Python
|
import logging
|
||
|
|
from app.core.config import settings
|
||
|
|
|
||
|
|
logger = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
def save_connection_to_metadb(
|
||
|
|
conn_id: str,
|
||
|
|
conn_type: str,
|
||
|
|
host: str,
|
||
|
|
port: int,
|
||
|
|
schema: str,
|
||
|
|
login: str,
|
||
|
|
password: str = ""
|
||
|
|
) -> bool:
|
||
|
|
"""Persist custom connection to Airflow metadata database if Airflow & SQLAlchemy are available."""
|
||
|
|
try:
|
||
|
|
from sqlalchemy import create_engine
|
||
|
|
from sqlalchemy.orm import sessionmaker
|
||
|
|
|
||
|
|
engine = create_engine(settings.AIRFLOW_DATABASE_URL)
|
||
|
|
Session = sessionmaker(bind=engine)
|
||
|
|
session = Session()
|
||
|
|
|
||
|
|
from airflow.models import Connection
|
||
|
|
existing = session.query(Connection).filter(Connection.conn_id == conn_id).first()
|
||
|
|
if not existing:
|
||
|
|
c = Connection(
|
||
|
|
conn_id=conn_id,
|
||
|
|
conn_type=conn_type,
|
||
|
|
host=host,
|
||
|
|
port=port,
|
||
|
|
schema=schema,
|
||
|
|
login=login,
|
||
|
|
password=password
|
||
|
|
)
|
||
|
|
session.add(c)
|
||
|
|
session.commit()
|
||
|
|
logger.info("Connection '%s' saved to Airflow MetaDB", conn_id)
|
||
|
|
session.close()
|
||
|
|
return True
|
||
|
|
except Exception as e:
|
||
|
|
logger.warning("MetaDB Connection Save Note: %s", e)
|
||
|
|
return False
|