datax-admin/app/crud/job_jdbc_datasource.py

import time
from typing import List
from sqlalchemy.orm import Session
from app.core.datasource.datasource import DataSrouceFactory
import app.schemas as schemas
import app.models as models
from app.utils import decode_user


def _decode(url, datasource, database_name):
    return url.replace('jdbc:', '').replace('hive2://', '').replace(f'{datasource}://', '').replace(f'/{database_name}',
                                                                                                    '')


def _format_datasource(db: Session, item: schemas.JobJdbcDatasourceBase, ds_id: int = 0):
    if ds_id != 0:
        item = db.query(models.JobJdbcDatasource).filter(models.JobJdbcDatasource.id == ds_id).first()
        if not item:
            raise Exception('未找到该数据源')
        item.jdbc_url = _decode(item.jdbc_url, item.datasource, item.database_name)
        item.jdbc_username, item.jdbc_password = decode_user(item.jdbc_username, item.jdbc_password)
    host, port = item.jdbc_url.split(':')
    if not host or not port:
        raise Exception('jdbc_url无效')

    ds = DataSrouceFactory.create(item.datasource, {'port': port, 'host': host, 'username': item.jdbc_username,
                                                    'password': item.jdbc_password,
                                                    'database_name': item.database_name,
                                                    'kerberos': item.kerberos,
                                                    'keytab': item.keytab,
                                                    'krb5config': item.krb5config,
                                                    'kerberos_service_name': item.kerberos_service_name,
                                                    'principal':item.principal})
    item.jdbc_url = ds.jdbc_url
    item.jdbc_username = ds.jdbc_username if item.kerberos == 0 else None
    item.jdbc_password = ds.jdbc_password if item.kerberos == 0 else None
    return ds, item


def test_datasource_connection(db: Session, item: schemas.JobJdbcDatasourceCreate):
    ds, item = _format_datasource(db, item)
    return ds.is_connect()


def get_table_schema(db: Session, ds_id: int, table_name: str):
    ds, item = _format_datasource(db, None, ds_id)
    return ds.get_table_schema(table_name)


def get_preview_data(db: Session, ds_id: int, table_name: str, limit: int = 100):
    ds, item = _format_datasource(db, None, ds_id)
    return ds.get_preview_data(table_name, limit)


def get_table_names(db: Session, ds_id: int):
    ds, item = _format_datasource(db, None, ds_id)
    return ds.list_tables()


def create_job_jdbc_datasource(db: Session, item: schemas.JobJdbcDatasourceCreate):
    ds, item = _format_datasource(db, item)

    #
    create_time: int = int(time.time())
    db_item = models.JobJdbcDatasource(**item.dict(), **{
        'status': 1,
        'create_time': create_time,
        'create_by': 'admin',
        'update_time': create_time,
        'update_by': 'admin',
        'jdbc_driver_class': ds.jdbc_driver_class
    })
    db.add(db_item)
    db.commit()
    db.refresh(db_item)
    return db_item


def get_job_jdbc_datasources(db: Session, datasource_type: str = None, skip: int = 0, limit: int = 20):
    res: List[models.JobJdbcDatasource] = []
    if datasource_type is not None and datasource_type != '':
        res = db.query(models.JobJdbcDatasource)\
            .filter(models.JobJdbcDatasource.datasource == datasource_type)\
            .filter(models.JobJdbcDatasource.status == 1).all()
    else:
        res = db.query(models.JobJdbcDatasource)\
            .filter(models.JobJdbcDatasource.status == 1).all()
    for item in res:
        item.jdbc_url = _decode(item.jdbc_url, item.datasource, item.database_name)
    return res


def update_job_jdbc_datasources(db: Session, ds_id: int, update_item: schemas.JobJdbcDatasourceUpdate):
    ds, update_item = _format_datasource(db, update_item)

    db_item = db.query(models.JobJdbcDatasource).filter(models.JobJdbcDatasource.id == ds_id).first()
    if not db_item:
        raise Exception('未找到该数据源')
    update_dict = update_item.dict(exclude_unset=True)
    for k, v in update_dict.items():
        setattr(db_item, k, v)
    db_item.jdbc_driver_class = ds.jdbc_driver_class
    db_item.update_time = int(time.time())
    db_item.update_by = 'admin1'  # TODO
    db.commit()
    db.flush()
    db.refresh(db_item)
    return db_item


def delete_job_jdbc_datasource(db: Session, ds_id: int):
    db_item = db.query(models.JobJdbcDatasource).filter(models.JobJdbcDatasource.id == ds_id).first()
    if not db_item:
        raise Exception('未找到该数据源')
    db_item.status = 0
    db.commit()
    db.flush()
    db.refresh(db_item)
    return db_item