From b42b623d758b5824c8e949575f122791242944c6 Mon Sep 17 00:00:00 2001 From: Herman Schaaf Date: Tue, 1 Aug 2023 16:47:08 +0100 Subject: [PATCH 1/2] Fix column resolver resource set --- cloudquery/sdk/scheduler/table_resolver.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/cloudquery/sdk/scheduler/table_resolver.py b/cloudquery/sdk/scheduler/table_resolver.py index 84cdc25..1fa6734 100644 --- a/cloudquery/sdk/scheduler/table_resolver.py +++ b/cloudquery/sdk/scheduler/table_resolver.py @@ -1,7 +1,8 @@ -from cloudquery.sdk.schema.table import Table -from cloudquery.sdk.schema import Resource from typing import Any, Generator +from cloudquery.sdk.schema import Resource +from cloudquery.sdk.schema.table import Table + class TableResolver: def __init__(self, table: Table, child_resolvers=[]) -> None: @@ -31,7 +32,7 @@ def resolve_column(self, client, resource: Resource, column_name: str): resource.set(column_name, resource.item[column_name]) else: if hasattr(resource.item, column_name): - resource.set(column_name, resource.item[column_name]) + resource.set(column_name, resource.item.column_name) def post_resource_resolve(self, client, resource): return From ef31b08f4b080ca4d9aca7cb21718b1e2cad107e Mon Sep 17 00:00:00 2001 From: Herman Schaaf Date: Tue, 1 Aug 2023 17:03:15 +0100 Subject: [PATCH 2/2] Fix getattr in resource set --- cloudquery/sdk/scheduler/table_resolver.py | 2 +- tests/scheduler/table_resolver.py | 34 ++++++++++++++++++++++ 2 files changed, 35 insertions(+), 1 deletion(-) create mode 100644 tests/scheduler/table_resolver.py diff --git a/cloudquery/sdk/scheduler/table_resolver.py b/cloudquery/sdk/scheduler/table_resolver.py index 1fa6734..3cde6bd 100644 --- a/cloudquery/sdk/scheduler/table_resolver.py +++ b/cloudquery/sdk/scheduler/table_resolver.py @@ -32,7 +32,7 @@ def resolve_column(self, client, resource: Resource, column_name: str): resource.set(column_name, resource.item[column_name]) else: if hasattr(resource.item, column_name): - resource.set(column_name, resource.item.column_name) + resource.set(column_name, getattr(resource.item, column_name)) def post_resource_resolve(self, client, resource): return diff --git a/tests/scheduler/table_resolver.py b/tests/scheduler/table_resolver.py new file mode 100644 index 0000000..3a11575 --- /dev/null +++ b/tests/scheduler/table_resolver.py @@ -0,0 +1,34 @@ +from dataclasses import dataclass + +import pyarrow as pa + +from cloudquery.sdk.scheduler import TableResolver +from cloudquery.sdk.schema import Table, Resource, Column + + +@dataclass +class TestItem: + test_column2: int = 1 + + +def test_table_resolver_resolve_column(): + test_table = Table( + name="test_table", + columns=[ + Column(name="test_column", type=pa.int64()), + Column(name="test_column2", type=pa.int64()), + ], + ) + resource_dict = Resource(table=test_table, parent=None, item={"test_column": 1}) + resource_obj = Resource( + table=test_table, parent=None, item=TestItem(test_column2=2) + ) + test_resolver = TableResolver(table=test_table, child_resolvers=[]) + test_resolver.resolve_column( + client=None, resource=resource_dict, column_name="test_column" + ) + test_resolver.resolve_column( + client=None, resource=resource_obj, column_name="test_column2" + ) + assert resource_dict.to_list_of_arr() == [[1], [None]] + assert resource_obj.to_list_of_arr() == [[None], [2]]