from collections.abc import Sequence

from assets import deals, new_deals, works
from dagster_duckdb import DuckDBIOManager
from dagster_duckdb.io_manager import DbTypeHandler
from dagster_duckdb_pandas import DuckDBPandasTypeHandler
from dagster_polars import PolarsParquetIOManager
from jobs import check_partititions_job, deals_job, musicbrainz_lookup_job
from schedules import deals_schedule
from sensors import musicbrainz_lookup_sensor

from dagster import Definitions


class PandasDuckDBIOManager(DuckDBIOManager):
    @staticmethod
    def type_handlers() -> Sequence[DbTypeHandler]:
        return [DuckDBPandasTypeHandler()]


definitions = Definitions(
    assets=[deals, new_deals, works],
    resources={
        "polars_parquet_io_manager": PolarsParquetIOManager(),
        "duckdb_io_manager": PandasDuckDBIOManager(database="vinyl"),
    },
    jobs=[deals_job, check_partititions_job, musicbrainz_lookup_job],
    schedules=[deals_schedule],
    sensors=[musicbrainz_lookup_sensor],
)