Snowflake中基于GeoIP2实现IP解析的UDF开发问题

一、Python Snowpark UDF 问题解决 #

1. Anaconda包权限错误处理 #

错误提示需ORGADMIN角色接受Anaconda条款才能使用第三方包，执行以下操作：

• 切换到ORGADMIN角色登录Snowflake，运行ALTER ACCOUNT SET ANACONDA_ACCESS_ENABLED = TRUE;
• 在Snowflake控制台的「Admin > Accounts > Settings」中完成Anaconda条款授权

2. UDF访问Stage文件的正确方式 #

stage_location参数用于存储UDF的代码和依赖包，不是用来读取GeoLite2文件。要读取Stage中的mmdb文件，需通过Snowflake文件API将文件下载到临时目录后再读取：

def parse_ip_country(ip_str: str) -> str:
    import tempfile
    import os
    import geoip2.database
    from snowflake.snowpark.files import SnowflakeFile
    
    # 从Stage读取文件到临时路径
    with SnowflakeFile.open('@AWS_CSV_STAGE/GeoLite2-City.mmdb', 'rb') as f:
        with tempfile.NamedTemporaryFile(delete=False, suffix='.mmdb') as tmp_f:
            tmp_f.write(f.read())
            tmp_path = tmp_f.name
    
    try:
        with geoip2.database.Reader(tmp_path) as reader:
            response = reader.city(ip_str)
            return response.country.iso_code if response.country.iso_code else 'UNKNOWN'
    finally:
        os.unlink(tmp_path)

注册UDF时，stage_location指定UDF代码的存储路径（如@AWS_CSV_STAGE/udf_stages），与mmdb文件位置分离。

3. 导入geoip2包的正确方式 #

session.add_packages('geoip2')是正确的，但需先解决Anaconda权限问题。若权限无法开通，可手动上传geoip2包到Stage，再用session.add_imports指定路径。

修正后的完整Python代码 #

from snowflake.snowpark import Session
from snowflake.snowpark.types import StringType
import tempfile
import os
from snowflake.snowpark.files import SnowflakeFile

# 连接参数
cnn_params = {
    "account": '*********',
    "user": '*********',
    "password": '*********',
    "warehouse": '*********',
    "database": '*********',
    "schema": '*********',
    "role": '*********',
}

def parse_ip_country(ip_str: str) -> str:
    import geoip2.database
    
    # 从Stage读取MMDB文件到临时文件
    with SnowflakeFile.open('@AWS_CSV_STAGE/GeoLite2-City.mmdb', 'rb') as stage_file:
        with tempfile.NamedTemporaryFile(delete=False, suffix='.mmdb') as tmp_file:
            tmp_file.write(stage_file.read())
            tmp_path = tmp_file.name
    
    try:
        with geoip2.database.Reader(tmp_path) as reader:
            response = reader.city(ip_str)
            return response.country.iso_code if response.country.iso_code else 'UNKNOWN'
    finally:
        os.unlink(tmp_path)

try:
    session = Session.builder.configs(cnn_params).create()
    # 添加依赖包
    session.add_packages('geoip2')
    
    # 注册永久UDF
    session.udf.register(
        func=parse_ip_country,
        return_type=StringType(),
        input_types=[StringType()],
        is_permanent=True,
        name='PARSE_IP_COUNTRY',
        replace=True,
        stage_location='@AWS_CSV_STAGE/udf_stages'  # 存储UDF代码的Stage路径
    )
    
    # 测试UDF
    session.sql("SELECT PARSE_IP_COUNTRY('203.0.113.0') AS COUNTRY_ISO").show()
except Exception as e:
    print(f"Error: {str(e)}")
finally:
    if session:
        session.close()
        print('Connection closed.')

二、Java UDF 问题解决 #

1. 语法与逻辑修正 #

Java UDF代码块必须是纯Java语法，不能混用Python写法；导入的文件直接使用Stage中的相对路径访问即可：

修正后的Java UDF代码 #

CREATE OR REPLACE FUNCTION PARSE_IP_COUNTRY(ip_str VARCHAR)
RETURNS VARCHAR NOT NULL
LANGUAGE JAVA
IMPORTS = (
    '@AWS_CSV_STAGE/lib/geoip2-2.8.0.jar',
    '@AWS_CSV_STAGE/geodata/GeoLite2-City.mmdb'
)
HANDLER = 'IpParser.getCountryIso'
AS
$$
import com.maxmind.geoip2.DatabaseReader;
import com.maxmind.geoip2.model.CityResponse;
import com.maxmind.geoip2.record.Country;
import java.net.InetAddress;
import java.io.File;

public class IpParser {
    public static String getCountryIso(String ipStr) {
        try {
            // 导入的MMDB文件直接使用Stage中的相对路径
            File dbFile = new File("geodata/GeoLite2-City.mmdb");
            DatabaseReader reader = new DatabaseReader.Builder(dbFile).build();
            InetAddress ipAddress = InetAddress.getByName(ipStr);
            CityResponse response = reader.city(ipAddress);
            Country country = response.getCountry();
            return country.getIsoCode() != null ? country.getIsoCode() : "UNKNOWN";
        } catch (Exception e) {
            return "ERROR: " + e.getMessage();
        }
    }
}
$$;

-- 测试UDF
SELECT PARSE_IP_COUNTRY('128.101.101.101') AS COUNTRY_ISO;

注意事项 #

• 确保geoip2依赖完整，若geoip2-2.8.0.jar不包含maxmind-db.jar，需一并上传到Stage
• 增加异常处理，避免UDF执行崩溃

内容的提问来源于stack exchange，提问作者Kaido