Class DataSourceContext

app/data_source/api/context.py:26–151

This class is responsible for loading data sources and caching them. It dynamically loads data source types from the data_source/sources directory. It loads data sources from the database and caches them.

class DataSourceContext:
    """
    This class is responsible for loading data sources and caching them.
    It dynamically loads data source types from the data_source/sources directory.
    It loads data sources from the database and caches them.
    """
    _initialized = False
    _data_source_cache: Dict[int, CachedDataSource] = {}
    _data_source_classes: Dict[str, BaseDataSource] = {}

    @classmethod
    def get_data_source_instance(cls, data_source_id: int) -> BaseDataSource:
        if not cls._initialized:
            cls.init()
            cls._initialized = True

        return cls._data_source_cache[data_source_id].instance

    @classmethod
    def get_data_source_class(cls, data_source_name: str) -> BaseDataSource:
        if not cls._initialized:
            cls.init()
            cls._initialized = True

        return cls._data_source_classes[data_source_name]

    @classmethod
    def get_data_source_classes(cls) -> Dict[str, BaseDataSource]:
        if not cls._initialized:
            cls.init()
            cls._initialized = True

        return cls._data_source_classes

    @classmethod
    async def create_data_source(cls, name: str, config: dict) -> BaseDataSource:
        async with async_session() as session:
            data_source_type = await session.execute(
                select(DataSourceType).filter_by(name=name)
            )
            data_source_type = data_source_type.scalar_one_or_none()
            if data_source_type is None:
                raise KnownException(message=f"Data source type {name} does not exist")

            data_source_class = DynamicLoader.get_data_source_class(name)
            logger.info(f"validating config for data source {name}")
            await data_source_class.validate_config(config)
            config_str = json.dumps(config)

            data_source_row = DataSource(type_id=data_source_type.id, config=config_str, created_at=get_utc_time_now())
            session.add(data_source_row)
            await session.commit()

            data_source = data_source_class(config=config, data_source_id=data_source_row.id)
            cls._data_source_cache[data_source_row.id] = CachedDataSource(indexed_docs=0, failed_tasks=0,
                                                                          instance=data_source)

            return data_source