from snowflake.snowpark.session import Session
from snowflake.snowpark.types import StringType
import base64
from hashlib import md5

def create_get_md5_hashkey(session: Session):
    def get_md5_hashkey(key: str) -> str:
        return base64.b64encode(md5(key.encode("utf-8")).digest()).decode("utf-8")

    session.udf.register(
        func=get_md5_hashkey,
        return_type=StringType(),
        input_types=[StringType()],
        name="get_md5_id",
        is_permanent=True,
        stage_location="@SUNO_CODE",
        replace=True,
    )
    print ("get_md5_id function created")