Skip to content

Commit a3983fb

Browse files
pkundurthyPraveen Kundurthy
pkundurthy
and
Praveen Kundurthy
authored
[Compile] Update compile.py table_display option to also display output for stating queries (#864)
Cleanup commits Co-authored-by: Praveen Kundurthy <[email protected]>
1 parent 9c35e39 commit a3983fb

File tree

4 files changed

+59
-7
lines changed

4 files changed

+59
-7
lines changed

‎api/py/ai/chronon/repo/compile.py

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -370,12 +370,14 @@ def _handle_deprecation_warning(
370370
)
371371

372372

373-
def _print_tables(obj: Union[Join, GroupBy], obj_class: Type[Union[Join, GroupBy]]) -> None:
373+
def _print_tables(obj: utils.ChrononJobTypes, obj_class: Type[utils.ChrononJobTypes]) -> None:
374374
tables = utils.get_modes_tables(obj)
375375
if obj_class is Join:
376376
_print_modes_tables("Output Join Tables", tables)
377377
if obj_class is GroupBy:
378378
_print_modes_tables("Output GroupBy Tables", tables)
379+
if obj_class is StagingQuery:
380+
_print_modes_tables("Output StagingQuery Tables", tables)
379381

380382

381383
def _handle_extra_conf_objects_to_materialize(
@@ -515,7 +517,10 @@ def _print_features_names(left, right):
515517

516518
def _print_modes_tables(left, right):
517519
text = textwrap.indent(json.dumps(right, indent=2), " " * 27)
518-
print(f"{left:>25} - \u001b[32m\n{text}\u001b[0m")
520+
json_start = "json.start"
521+
json_end = "json.end\n"
522+
523+
print(f"{left:>25} - \n{json_start:>25} \u001b[32m\n{text}\u001b[0m \n{json_end:>25}")
519524

520525

521526
def _print_error(left, right):

‎api/py/test/sample/production/group_bys/sample_team/sample_group_by_group_by.v1

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
"production": 0,
55
"customJson": "{\"lag\": 0, \"groupby_tags\": null, \"column_tags\": {}}",
66
"dependencies": [
7-
"{\"name\": \"wait_for_sample_namespace.sample_team_sample_group_by_require_backfill_ds\", \"spec\": \"sample_namespace.sample_team_sample_group_by_require_backfill/ds={{ ds }}\", \"start\": \"2021-04-09\", \"end\": null}"
7+
"{\"name\": \"wait_for_sample_namespace.sample_team_sample_group_by_group_by_require_backfill_ds\", \"spec\": \"sample_namespace.sample_team_sample_group_by_group_by_require_backfill/ds={{ ds }}\", \"start\": \"2021-04-09\", \"end\": null}"
88
],
99
"tableProperties": {
1010
"sample_config_json": "{\"sample_key\": \"sample_value\"}",
@@ -17,7 +17,7 @@
1717
"sources": [
1818
{
1919
"events": {
20-
"table": "sample_namespace.sample_team_sample_group_by_require_backfill",
20+
"table": "sample_namespace.sample_team_sample_group_by_group_by_require_backfill",
2121
"query": {
2222
"selects": {
2323
"event": "event_expr",

‎api/py/test/sample/production/joins/sample_team/sample_join.group_by_of_group_by

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66
"customJson": "{\"check_consistency\": false, \"lag\": 0, \"join_tags\": null, \"join_part_tags\": {}}",
77
"dependencies": [
88
"{\"name\": \"wait_for_sample_namespace.sample_team_sample_staging_query_v1_ds\", \"spec\": \"sample_namespace.sample_team_sample_staging_query_v1/ds={{ ds }}\", \"start\": \"2021-03-01\", \"end\": null}",
9-
"{\"name\": \"wait_for_sample_namespace.sample_team_sample_group_by_require_backfill_ds\", \"spec\": \"sample_namespace.sample_team_sample_group_by_require_backfill/ds={{ ds }}\", \"start\": \"2021-04-09\", \"end\": null}"
9+
"{\"name\": \"wait_for_sample_namespace.sample_team_sample_group_by_group_by_require_backfill_ds\", \"spec\": \"sample_namespace.sample_team_sample_group_by_group_by_require_backfill/ds={{ ds }}\", \"start\": \"2021-04-09\", \"end\": null}"
1010
],
1111
"tableProperties": {
1212
"source": "chronon"
@@ -39,7 +39,7 @@
3939
"production": 0,
4040
"customJson": "{\"lag\": 0, \"groupby_tags\": null, \"column_tags\": {}}",
4141
"dependencies": [
42-
"{\"name\": \"wait_for_sample_namespace.sample_team_sample_group_by_require_backfill_ds\", \"spec\": \"sample_namespace.sample_team_sample_group_by_require_backfill/ds={{ ds }}\", \"start\": \"2021-04-09\", \"end\": null}"
42+
"{\"name\": \"wait_for_sample_namespace.sample_team_sample_group_by_group_by_require_backfill_ds\", \"spec\": \"sample_namespace.sample_team_sample_group_by_group_by_require_backfill/ds={{ ds }}\", \"start\": \"2021-04-09\", \"end\": null}"
4343
],
4444
"tableProperties": {
4545
"sample_config_json": "{\"sample_key\": \"sample_value\"}",
@@ -52,7 +52,7 @@
5252
"sources": [
5353
{
5454
"events": {
55-
"table": "sample_namespace.sample_team_sample_group_by_require_backfill",
55+
"table": "sample_namespace.sample_team_sample_group_by_group_by_require_backfill",
5656
"query": {
5757
"selects": {
5858
"event": "event_expr",

‎api/py/test/test_compile.py

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,9 @@
1616
# See the License for the specific language governing permissions and
1717
# limitations under the License.
1818

19+
import json
1920
import os
21+
import re
2022

2123
import pytest
2224
from ai.chronon.repo.compile import extract_and_convert
@@ -53,6 +55,17 @@ def _assert_file_exists(full_file_path, message):
5355
assert os.path.isfile(full_file_path), message
5456

5557

58+
def _extract_display_output_json_block(output, block_name):
59+
"""Extract a JSON block from the output of a CLI command."""
60+
61+
field_index = output.find(block_name)
62+
subset = output[field_index:]
63+
start_index = subset.find("json.start")
64+
end_index = subset.find("json.end")
65+
s = re.sub(r"\x1b\[[0-9;]*m", "", subset[start_index + 10 : end_index].strip())
66+
return json.loads(s)
67+
68+
5669
@pytest.fixture
5770
def specific_setup():
5871
# This setup code will only run for tests that request this fixture
@@ -281,6 +294,18 @@ def test_compile_table_display():
281294
result = _invoke_cli_with_params(runner, input_path, ["--table-display"])
282295

283296
assert "Output Join Tables" in result.output
297+
output_json_dict = _extract_display_output_json_block(result.output, "Output Join Tables")
298+
expected_json_dict = {
299+
"backfill": ["chronon_db.sample_team_sample_join_with_derivations_on_external_parts_v1"],
300+
"stats-summary": ["chronon_db.sample_team_sample_join_with_derivations_on_external_parts_v1_daily_stats"],
301+
"log-flattener": ["chronon_db.sample_team_sample_join_with_derivations_on_external_parts_v1_logged"],
302+
"bootstrap": ["chronon_db.sample_team_sample_join_with_derivations_on_external_parts_v1_bootstrap"],
303+
"join_parts": [
304+
"sample_team_sample_join_with_derivations_on_external_parts_v1_sample_team_event_sample_group_by_v1",
305+
"sample_team_sample_join_with_derivations_on_external_parts_v1_sample_team_entity_sample_group_by_from_module_v1",
306+
],
307+
}
308+
assert json.dumps(output_json_dict, sort_keys=True) == json.dumps(expected_json_dict, sort_keys=True)
284309
assert result.exit_code == 0
285310

286311

@@ -307,6 +332,28 @@ def test_compile_feature_display():
307332
assert result.exit_code == 0
308333

309334

335+
def test_table_display_staging_query():
336+
"""
337+
Test a staging query compile produces related table
338+
"""
339+
runner = CliRunner()
340+
input_path = f"staging_queries/sample_team/sample_staging_query.py"
341+
result = runner.invoke(
342+
extract_and_convert,
343+
[
344+
"--chronon_root=test/sample",
345+
f"--input_path={input_path}",
346+
"--table-display",
347+
],
348+
)
349+
350+
assert "Output StagingQuery Tables" in result.output
351+
output_json_dict = _extract_display_output_json_block(result.output, "Output StagingQuery Tables")
352+
expected_json_dict = {"backfill": ["sample_namespace.sample_team_sample_staging_query_v1"]}
353+
assert json.dumps(output_json_dict, sort_keys=True) == json.dumps(expected_json_dict, sort_keys=True)
354+
assert result.exit_code == 0
355+
356+
310357
def test_compile_dependency_staging_query():
311358
"""
312359
Test that compiling a staging query does not error out

0 commit comments

Comments
 (0)