datax-admin/app/crud/job_jdbc_datasource.py

import base64
import time
from typing import List
from sqlalchemy.orm import Session
from app.core.ds.datasource import DataSrouceFactory
import app.schemas as schemas
import  app.models as models
from app.utils import decode_user
from configs.logging import logger


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('JobJdbcDatasource not found')
        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 is invalid')

    ds = DataSrouceFactory.create(item.datasource, {'port': port, 'host': host, 'username': item.jdbc_username, 'password': item.jdbc_password, 'database_name': item.database_name})
    item.jdbc_url = ds.jdbc_url
    item.jdbc_username = ds.jdbc_username
    item.jdbc_password = ds.jdbc_password
    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, skip: int = 0, limit: int = 20):

    res: List[models.JobJdbcDatasource] = 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

    # return db.query(models.JobJdbcDatasource).offset(skip).limit(limit).all()


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('JobJdbcDatasource not found')
    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('JobJdbcDatasource not found')
    db_item.status = 0
    db.commit()
    db.flush()
    db.refresh(db_item)
    return db_item