diff --git a/cloudquery/sdk/schema/arrow.py b/cloudquery/sdk/schema/arrow.py index 2f21bc1..23deed2 100644 --- a/cloudquery/sdk/schema/arrow.py +++ b/cloudquery/sdk/schema/arrow.py @@ -7,3 +7,5 @@ METADATA_FALSE = b"false" METADATA_TABLE_NAME = b"cq:table_name" METADATA_TABLE_DESCRIPTION = b"cq:table_description" +METADATA_TABLE_TITLE = b"cq:table_title" +METADATA_TABLE_DEPENDS_ON = b"cq:table_depends_on" diff --git a/cloudquery/sdk/schema/table.py b/cloudquery/sdk/schema/table.py index 63c6228..7198d63 100644 --- a/cloudquery/sdk/schema/table.py +++ b/cloudquery/sdk/schema/table.py @@ -57,12 +57,24 @@ def from_arrow_schema(cls, schema: pa.Schema) -> Table: columns = [] for field in schema: columns.append(Column.from_arrow_field(field)) + parent = None + if arrow.METADATA_TABLE_DEPENDS_ON in schema.metadata: + parent = Table( + name=schema.metadata[arrow.METADATA_TABLE_DEPENDS_ON].decode("utf-8"), + columns=[], + ) return cls( - name=schema.metadata[arrow.METADATA_TABLE_NAME].decode("utf-8"), + name=schema.metadata.get(arrow.METADATA_TABLE_NAME, b"").decode("utf-8"), + title=schema.metadata.get(arrow.METADATA_TABLE_TITLE, b"").decode("utf-8"), columns=columns, description=schema.metadata.get(arrow.METADATA_TABLE_DESCRIPTION).decode( "utf-8" ), + is_incremental=schema.metadata.get( + arrow.METADATA_INCREMENTAL, arrow.METADATA_FALSE + ) + == arrow.METADATA_TRUE, + parent=parent, ) def to_arrow_schema(self): @@ -70,7 +82,11 @@ def to_arrow_schema(self): md = { arrow.METADATA_TABLE_NAME: self.name, arrow.METADATA_TABLE_DESCRIPTION: self.description, - # arrow.METADATA_CONSTRAINT_NAME: + arrow.METADATA_TABLE_TITLE: self.title, + arrow.METADATA_TABLE_DEPENDS_ON: self.parent.name if self.parent else "", + arrow.METADATA_INCREMENTAL: arrow.METADATA_TRUE + if self.is_incremental + else arrow.METADATA_FALSE, } for column in self.columns: fields.append(column.to_arrow_field()) diff --git a/tests/schema/test_table.py b/tests/schema/test_table.py index e7bb073..235c110 100644 --- a/tests/schema/test_table.py +++ b/tests/schema/test_table.py @@ -6,8 +6,22 @@ def test_table(): - table = Table("test_table", [Column("test_column", pa.int32())]) - table.to_arrow_schema() + table = Table( + name="test_table", + columns=[Column("test_column", pa.int32())], + title="Test Table", + description="Test description", + parent=Table(name="parent_table", columns=[]), + relations=[], + is_incremental=True, + ) + sch = table.to_arrow_schema() + got = Table.from_arrow_schema(sch) + assert got.name == table.name + assert got.title == table.title + assert got.description == table.description + assert got.is_incremental == table.is_incremental + assert got.parent.name == table.parent.name def test_filter_dfs_warns_no_matches():