trino-module/trino-module.py at main · fybrik/trino-module · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
import time
import base64
import json
import trino
import yaml

data_dict = {}


def execute_query(cur, query):
    cur.execute(query)
    # res = None
    try:
        res = cur.fetchall()
        return res
    except Exception as e:
        print(e)
    # return res

def connect_user(username, catalog):
    conn = trino.dbapi.connect(host='host.docker.internal', port=8080, user=username, catalog=catalog)
    # conn = trino.dbapi.connect(host='localhost', port=8080, user=username, catalog=catalog)
    cur = conn.cursor()
    return cur

def get_details_from_conf():
    """ Parse the configuration and get the data details and policies """
    with open("/etc/conf/conf.yaml", 'r') as stream:
    # with open("sample-conf.yaml", 'r') as stream:
        content = yaml.safe_load(stream)
        for key, val in content.items():
            if "data" in key:
                for data in val:
                    dataset_id = data["name"]
                    name = dataset_id.split("/")[1]
                    endpoint_url = data["connection"]["s3"]["endpoint_url"]
                    transformations = base64.b64decode(data["transformations"])
                    transformations_json = json.loads(transformations.decode('utf-8'))
                    transformation = transformations_json[0]['name']
                    transformation_cols = transformations_json[0][transformation]["columns"]
                    data_dict[name] = {'format': data["format"], 'endpoint_url': endpoint_url, 'path': data["path"], 'transformation': transformation,
                     'transformation_cols': transformation_cols}
    print(data_dict[name])
    return data_dict[name]

def get_policy_query(transformation_cols, sql_path, col_names):
    request_cols = [col for col in col_names if col not in transformation_cols]
    if len(request_cols) < 1:
        return ""
    requested_cols_string = request_cols[0]
    for col in request_cols[1:]:
        add_col = ", " + col
        requested_cols_string += add_col
    sql_vds = "select " + requested_cols_string + " from " + sql_path
    return sql_vds

if __name__ == "__main__":
    print("show catalogs")
    cur = connect_user("admin", "iceberg")
    res = execute_query(cur, "SHOW CATALOGS")
    print(res)

    print("create schema")
    schema_query = "create schema hive.icebergtrino with (location = 's3a://iceberg/')"
    # schema_query = "create schema hive.icebergtrino with (location = 's3a://fybric-objectstorage-iceberg-demo/warehouse/db/')"

    res = execute_query(cur, schema_query)
    print(res)

    create_table_query = "create table iceberg.icebergtrino.logs (\
        a DOUBLE,\
        b DOUBLE,\
        c DOUBLE,\
        d DOUBLE\
        )\
        with (format = 'ORC')"

    res = execute_query(cur, create_table_query)
    print(res)

    print("insert rows to table")
    insert_query = "INSERT INTO iceberg.icebergtrino.logs VALUES\
                    (\
                        1,\
                        2,\
                        3,\
                        4\
                    )"

    res = execute_query(cur, insert_query)
    print(res)

    # Create a sql query for transformations from fybrik
    parse_conf = get_details_from_conf()
    transformation = parse_conf['transformation']
    transformation_cols = parse_conf['transformation_cols']

    # Get the columns of the new source
    sql_path = "iceberg.icebergtrino.logs"
    col_names = ["a", "b", "c", "d"]

    # Get the sql query from the policies
    sql_view = get_policy_query(transformation_cols, sql_path, col_names)


    print("create view")
    # view_query = "create view iceberg.icebergtrino.view1 as select event_time, message from iceberg.icebergtrino.logs"
    view_query = "create view iceberg.icebergtrino.view1 as " + sql_view
    print(view_query)
    res = execute_query(cur, view_query)
    print(res)

    print("select from the table")
    select_query = 'select * from iceberg.icebergtrino.logs'
    res = execute_query(cur, select_query)
    print(res)

    print("user1 select")
    cur = connect_user("user1", "iceberg")
    res = execute_query(cur, select_query)
    print(res)

    print("user1 select from the view")
    select_query = 'select * from iceberg.icebergtrino.view1'
    res = execute_query(cur, select_query)
    print(res)