from typing import Optional, Dict

from graphql import GraphQLArgument, GraphQLNonNull
from sqlalchemy.ext.declarative import DeclarativeMeta

from .types import get_graphql_type_from_column


def make_pk_args(model: DeclarativeMeta) -> Optional[Dict[str, GraphQLArgument]]:
    primary_key = model.__table__.primary_key  # type: ignore

    if not primary_key:
        return None

    args = {}
    for column in primary_key.columns:
        graphql_type = get_graphql_type_from_column(column)
        args[column.name] = GraphQLArgument(GraphQLNonNull(graphql_type))

    return args
