feat: list virtual tables

Author: nicornk

libs/foundry-dev-tools/src/foundry_dev_tools/clients/tables.py

@@ -101,3 +101,45 @@ def create_snowflake_table(
         return self.api_create_table(
             name=name, parent_rid=parent_rid, upstream_config=upstream_config, skip_validation=skip_validation
         ).json()
+
+    def api_list_tables(
+        self, connection_rid: SourceRid, limit: int | None = 1000, page_token: str | None = None, **kwargs
+    ) -> requests.Response:
+        """List Virtual Tables for a source."""
+        return self.api_request(
+            "GET",
+            api_path="tables",
+            params={
+                "connectionRid": connection_rid,
+                **({"limit": limit} if limit else {}),
+                **({"pageToken": page_token} if page_token else {}),
+            },
+            **kwargs,
+        )
+
+    def list_tables(self, connection_rid: SourceRid) -> list[dict]:
+        """Returns all tables for a connection/source by handling pagination.
+
+        Args:
+            connection_rid: The rid of the compass Source/Connection.
+
+        Returns:
+            a list of dicts with the following format:
+                {
+                    'rid': 'ri.tables.main.table.<...>',
+                    'upstream': {
+                        'type': '<type>',
+                        '<type>': {
+                                'connectionRid': 'ri.magritte..source.<...>',
+                                'relation': {'database': 'test', 'schema': 'test', 'table': 'test2'}
+                            }
+                    }
+                }
+        """
+        tables_response = self.api_list_tables(connection_rid=connection_rid, limit=1000)
+        response_json = tables_response.json()
+        tables = response_json["tables"]
+        while next_page_token := response_json.get("nextPageToken"):
+            response_json = self.api_list_tables(connection_rid=connection_rid, page_token=next_page_token).json()
+            tables.extend(response_json["tables"])
+        return tables

tests/integration/clients/test_magritte_coordinator.py

@@ -7,6 +7,8 @@
 from tests.integration.conftest import TEST_SINGLETON
 from tests.integration.utils import INTEGRATION_TEST_COMPASS_ROOT_RID, INTEGRATION_TEST_EGRESS_POLICY_RID, MARKING_ID
 
+SNOWFLAKE_TEST_ACCOUNT_IDENTIFIER = "fdt-integration-test-account1"
+
 
 @pytest.fixture()
 def empty_s3_source() -> api_types.SourceRid:
@@ -20,7 +22,26 @@ def empty_s3_source() -> api_types.SourceRid:
 
     yield source_rid
 
-    # Delete test folder
+    # Delete test source
+    response = TEST_SINGLETON.ctx.compass.api_delete_permanently(
+        rids={source_rid}, delete_options={"DO_NOT_REQUIRE_TRASHED"}
+    )
+    assert response.status_code == 200
+
+
+@pytest.fixture()
+def empty_snowflake_source() -> api_types.SourceRid:
+    # Create a new Snowflake
+    client = TEST_SINGLETON.ctx.magritte_coordinator
+    rnd = "".join(choice(ascii_uppercase) for _ in range(5))
+    name = "fdt-test-snowflake_" + rnd
+    source_rid = client.create_snowflake_source(
+        name=name, parent_rid=INTEGRATION_TEST_COMPASS_ROOT_RID, account_identifier=SNOWFLAKE_TEST_ACCOUNT_IDENTIFIER
+    )
+
+    yield source_rid
+
+    # Delete test source
     response = TEST_SINGLETON.ctx.compass.api_delete_permanently(
         rids={source_rid}, delete_options={"DO_NOT_REQUIRE_TRASHED"}
     )
@@ -208,15 +229,9 @@ def test_export_toggles(empty_s3_source):
     }
 
 
-def test_snowflake():
+def test_snowflake(empty_snowflake_source):
     client = TEST_SINGLETON.ctx.magritte_coordinator
-    # Create a new S3 Direct source
-    rnd = "".join(choice(ascii_uppercase) for _ in range(5))
-    name = "fdt-test-snowflake_" + rnd
-    account_identifier = "account1"
-    source_rid = client.create_snowflake_source(
-        name=name, parent_rid=INTEGRATION_TEST_COMPASS_ROOT_RID, account_identifier=account_identifier
-    )
+    source_rid = empty_snowflake_source
 
     # enable OIDC
     client.enable_snowflake_external_oauth(source_rid=source_rid)
@@ -234,7 +249,7 @@ def test_snowflake():
     oidc_issuer = client.get_oidc_issuer()
     assert as_json["type"] == "cloud"
     assert as_json["cloud"]["oidcRuntime"] == {
-        "audience": f"https://{account_identifier}.snowflakecomputing.com",
+        "audience": f"https://{SNOWFLAKE_TEST_ACCOUNT_IDENTIFIER}.snowflakecomputing.com",
         "issuer": oidc_issuer,
         "subject": source_rid,
     }
@@ -255,3 +270,41 @@ def test_snowflake():
     }
 
     TEST_SINGLETON.ctx.compass.api_delete_permanently(rids={source_rid}, delete_options={"DO_NOT_REQUIRE_TRASHED"})
+
+
+def test_snowflake_tables(empty_snowflake_source):
+    client = TEST_SINGLETON.ctx.magritte_coordinator
+    source_rid = empty_snowflake_source
+
+    client.enable_snowflake_external_oauth(source_rid=source_rid)
+
+    tables_client = TEST_SINGLETON.ctx.tables
+
+    tables = tables_client.list_tables(connection_rid=source_rid)
+    assert tables == []
+
+    table1 = tables_client.create_snowflake_table(
+        parent_rid=INTEGRATION_TEST_COMPASS_ROOT_RID,
+        connection_rid=source_rid,
+        name="fdt-test-snowflake_table" + "".join(choice(ascii_uppercase) for _ in range(5)),
+        database="test",
+        schema="test",
+        table="test1",
+        skip_validation=True,
+    )
+
+    table2 = tables_client.create_snowflake_table(
+        parent_rid=INTEGRATION_TEST_COMPASS_ROOT_RID,
+        connection_rid=source_rid,
+        name="fdt-test-snowflake_table" + "".join(choice(ascii_uppercase) for _ in range(5)),
+        database="test",
+        schema="test",
+        table="test2",
+        skip_validation=True,
+    )
+
+    tables = tables_client.list_tables(connection_rid=source_rid)
+    assert len(tables) == 2
+
+    TEST_SINGLETON.ctx.compass.api_delete_permanently(rids={table1, table2}, delete_options={"DO_NOT_REQUIRE_TRASHED"})
+    TEST_SINGLETON.ctx.compass.api_delete_permanently(rids={source_rid}, delete_options={"DO_NOT_REQUIRE_TRASHED"})