|
1 | 1 | """Test DuckDB backend steps.""" |
2 | 2 |
|
3 | 3 | # pylint: disable=redefined-outer-name,unused-import,line-too-long |
| 4 | +import datetime |
| 5 | +from io import StringIO |
| 6 | +import json |
4 | 7 | import tempfile |
5 | 8 | from pathlib import Path |
6 | 9 | from typing import Iterator, List, Optional, Set, Tuple, Type |
@@ -884,3 +887,112 @@ def test_read_and_write_nested_parquet(nested_typecast_parquet): |
884 | 887 | "datetimefield": "TIMESTAMP", |
885 | 888 | "subfield": "STRUCT(id BIGINT, substrfield VARCHAR, subarrayfield DATE[])[]", |
886 | 889 | } |
| 890 | + |
| 891 | +def test_cache_management(): |
| 892 | + conn = DUCKDB_STEP_BACKEND.connection |
| 893 | + with tempfile.NamedTemporaryFile(mode="w") as tf1, tempfile.NamedTemporaryFile(mode="w") as tf2: |
| 894 | + td1 = [ |
| 895 | + {"greeting": "hi", "num_one": 2, "num_two": 4, "test_date": datetime.date(2020,5,1), "active": True}, |
| 896 | + {"greeting": "bonjour", "num_one": 3, "num_two": 9, "test_date": datetime.date(2025,7,4), "active": False}, |
| 897 | + ] |
| 898 | + tf1.write(json.dumps(td1, default=str)) |
| 899 | + |
| 900 | + tf1.seek(0) |
| 901 | + |
| 902 | + td1_schema = {"greeting": "STRING", "num_one": "BIGINT", "num_two" :"BIGINT", "test_date": "DATE", "active": "BOOLEAN"} |
| 903 | + |
| 904 | + td2 = [ |
| 905 | + {"farewell": "aurevoir", "lots_of_nums": [1,2,3], "nested_field": {"nested_str": "test1", "nested_timestamp": datetime.datetime(2024,3,5,1,2,3)}}, |
| 906 | + {"farewell": "bye", "lots_of_nums": [4,5,6,7,8], "nested_field": {"nested_str": "test2", "nested_timestamp": datetime.datetime(2022,8,12,10,12,14)}}, |
| 907 | + ] |
| 908 | + |
| 909 | + tf2.write(json.dumps(td2, default=str)) |
| 910 | + |
| 911 | + tf2.seek(0) |
| 912 | + |
| 913 | + td2_schema = {"farewell": "STRING", "lots_of_nums": "BIGINT[]", "nested_field": "STRUCT(nested_str STRING, nested_timestamp TIMESTAMP)"} |
| 914 | + |
| 915 | + em = EntityManager({}) |
| 916 | + em.entities["test_one"] = conn.read_json(tf1.name, columns=td1_schema) |
| 917 | + em.entities["test_two"] = conn.read_json(tf2.name, columns=td2_schema) |
| 918 | + |
| 919 | + DUCKDB_STEP_BACKEND.cache_entity("test_one", em.entities) |
| 920 | + |
| 921 | + cached_tables = [rw["table_name"] for rw in conn.sql("SELECT table_name from duckdb_tables() WHERE database_name = 'temp'").pl().to_dicts()] |
| 922 | + |
| 923 | + test_one_temp_name = list(filter(lambda x: x.startswith("test_one_"), cached_tables))[0] |
| 924 | + |
| 925 | + assert "test_one" in DUCKDB_STEP_BACKEND.entity_cache_tracker |
| 926 | + assert DUCKDB_STEP_BACKEND.entity_cache_tracker["test_one"] == test_one_temp_name |
| 927 | + assert sorted(em.entities["test_one"].pl().to_dicts(), key=lambda x: x.get("test_date")) == td1 |
| 928 | + assert sorted(conn.table(test_one_temp_name).pl().to_dicts(), key=lambda x: x.get("test_date")) == td1 |
| 929 | + assert "test_two" not in DUCKDB_STEP_BACKEND.entity_cache_tracker |
| 930 | + |
| 931 | + DUCKDB_STEP_BACKEND.cache_entity("test_two", em.entities) |
| 932 | + DUCKDB_STEP_BACKEND._remove_cached_artifact("test_one") |
| 933 | + |
| 934 | + cached_tables = [rw["table_name"] for rw in conn.sql("SELECT table_name from duckdb_tables() WHERE database_name = 'temp'").pl().to_dicts()] |
| 935 | + test_two_temp_name = list(filter(lambda x: x.startswith("test_two_"), cached_tables))[0] |
| 936 | + |
| 937 | + assert "test_two" in DUCKDB_STEP_BACKEND.entity_cache_tracker |
| 938 | + assert DUCKDB_STEP_BACKEND.entity_cache_tracker["test_two"] in cached_tables |
| 939 | + assert sorted(em.entities["test_two"].pl().to_dicts(), key=lambda x: x.get("farewell")) == td2 |
| 940 | + assert sorted(conn.table(test_two_temp_name).pl().to_dicts(), key=lambda x: x.get("farewell")) == td2 |
| 941 | + assert "test_one" not in DUCKDB_STEP_BACKEND.entity_cache_tracker |
| 942 | + |
| 943 | + DUCKDB_STEP_BACKEND.clear_entity_cache() |
| 944 | + |
| 945 | + cached_tables = [rw["table_name"] for rw in conn.sql("SELECT table_name from duckdb_tables() WHERE database_name = 'temp'").pl().to_dicts()] |
| 946 | + |
| 947 | + assert not DUCKDB_STEP_BACKEND.entity_cache_tracker |
| 948 | + assert not any(tbl.startswith("test_one_") for tbl in cached_tables) |
| 949 | + assert not any(tbl.startswith("test_two_") for tbl in cached_tables) |
| 950 | + |
| 951 | +def test_cache_management_with_update(): |
| 952 | + conn = DUCKDB_STEP_BACKEND.connection |
| 953 | + with tempfile.NamedTemporaryFile(mode="w") as tf, tempfile.NamedTemporaryFile(mode="w") as ef: |
| 954 | + td = [ |
| 955 | + {"idx": 1, "lots_of_nums": [1,2,3], "nested_field": {"nested_str": "test1", "nested_timestamp": datetime.datetime(2024,3,5,1,2,3)}}, |
| 956 | + {"idx": 2, "lots_of_nums": [4,5,6,7,8], "nested_field": {"nested_str": "test2", "nested_timestamp": datetime.datetime(2022,8,12,10,12,14)}}, |
| 957 | + ] |
| 958 | + |
| 959 | + tf.write(json.dumps(td, default=str)) |
| 960 | + tf.seek(0) |
| 961 | + td_schema = {"idx": "BIGINT", |
| 962 | + "lots_of_nums": "BIGINT[]", |
| 963 | + "nested_field": "STRUCT(nested_str STRING, nested_timestamp TIMESTAMP)"} |
| 964 | + |
| 965 | + |
| 966 | + em = EntityManager({}) |
| 967 | + em.entities["test_df"] = conn.read_json(tf.name, columns=td_schema) |
| 968 | + |
| 969 | + DUCKDB_STEP_BACKEND.cache_entity("test_df", em.entities) |
| 970 | + |
| 971 | + cached_tables = [rw["table_name"] for rw in conn.sql("SELECT table_name from duckdb_tables() WHERE database_name = 'temp'").pl().to_dicts()] |
| 972 | + |
| 973 | + first_temp_name = list(filter(lambda x: x.startswith("test_df_"), cached_tables))[0] |
| 974 | + |
| 975 | + assert "test_df" in DUCKDB_STEP_BACKEND.entity_cache_tracker |
| 976 | + assert DUCKDB_STEP_BACKEND.entity_cache_tracker["test_df"] == first_temp_name |
| 977 | + assert sorted(em.entities["test_df"].pl().to_dicts(), key=lambda x: x.get("idx")) == td |
| 978 | + |
| 979 | + extra_data = [{"idx": 3, "lots_of_nums": [9], "nested_field": {"nested_str": "test3", "nested_timestamp": datetime.datetime(2024,1,9,3,2,1)}}] |
| 980 | + ef.write(json.dumps(extra_data, default=str)) |
| 981 | + ef.seek(0) |
| 982 | + em.entities["test_df"] = em.entities["test_df"].union(conn.read_json(ef.name, columns=td_schema)) |
| 983 | + DUCKDB_STEP_BACKEND.cache_entity("test_df", em.entities) |
| 984 | + |
| 985 | + cached_tables = [rw["table_name"] for rw in conn.sql("SELECT table_name from duckdb_tables() WHERE database_name = 'temp'").pl().to_dicts()] |
| 986 | + tables_of_interest = list(filter(lambda x: x.startswith("test_df_"), cached_tables)) |
| 987 | + assert len(tables_of_interest) == 1 |
| 988 | + assert first_temp_name not in tables_of_interest |
| 989 | + second_temp_name = tables_of_interest[0] |
| 990 | + assert sorted(em.entities["test_df"].pl().to_dicts(), key=lambda x: x.get("idx")) == td + extra_data |
| 991 | + assert sorted(conn.table(second_temp_name).pl().to_dicts(), key=lambda x: x.get("idx")) == td + extra_data |
| 992 | + |
| 993 | + DUCKDB_STEP_BACKEND.clear_entity_cache() |
| 994 | + |
| 995 | + cached_tables = [rw["table_name"] for rw in conn.sql("SELECT table_name from duckdb_tables() WHERE database_name = 'temp'").pl().to_dicts()] |
| 996 | + |
| 997 | + assert not DUCKDB_STEP_BACKEND.entity_cache_tracker |
| 998 | + assert not any(tbl.startswith("test_df_") for tbl in cached_tables) |
0 commit comments