fix iceberg tests

buremba · Jan 13, 2025 · 451dc8f · 451dc8f
1 parent edd433f
commit 451dc8f
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 5 deletions.
diff --git a/tests/integration/load.py b/tests/integration/load.py
@@ -14,12 +14,13 @@ def test_create_iceberg_table(self):
             AS {SIMPLE_QUERY}
             """)
             universql_result = execute_query(conn, f"SELECT * FROM test_iceberg_table LIMIT 1")
-            print(universql_result)
+            assert universql_result.num_rows == 1
 
     def test_create_temp_table(self):
         with universql_connection(warehouse=None) as conn:
             execute_query(conn, f"CREATE TEMP TABLE test_table AS {SIMPLE_QUERY}")
-            execute_query(conn, "SELECT * FROM TEST_TABLE")
+            universql_result = execute_query(conn, "SELECT * FROM TEST_TABLE")
+            assert universql_result.num_rows == 1
 
     @pytest.mark.skip(reason="not implemented")
     def test_create_stage(self):

diff --git a/universql/protocol/session.py b/universql/protocol/session.py
@@ -10,7 +10,7 @@
 import sentry_sdk
 import sqlglot
 from pyiceberg.catalog import PY_CATALOG_IMPL, load_catalog, TYPE
-from pyiceberg.exceptions import NoSuchTableError, TableAlreadyExistsError, NoSuchNamespaceError
+from pyiceberg.exceptions import TableAlreadyExistsError, NoSuchNamespaceError
 from pyiceberg.io import PY_IO_IMPL
 from sqlglot import ParseError
 from sqlglot.expressions import Create, Identifier, DDL, Query, Use
@@ -69,7 +69,8 @@ def _get_iceberg_catalog(self):
             catalog_props[TYPE] = "glue"
             catalog = load_catalog(catalog_name, **catalog_props)
         else:
-            database_path = Template(self.context.get('database_path') or f"_{self.session_id}_universql_session").substitute(
+            database_path = Template(
+                self.context.get('database_path') or f"_{self.session_id}_universql_session").substitute(
                 {"session_id": self.session_id}) + ".sqlite"
             catalog_name = "duckdb"
             self.metadata_db = tempfile.NamedTemporaryFile(delete=False, suffix=database_path)
@@ -219,7 +220,7 @@ def get_table_paths_from_catalog(self, alternative_catalog: ICatalog, tables: li
         for table in tables:
             full_qualifier_ = full_qualifier(table, self.credentials)
             table_path = alternative_catalog.get_table_paths([full_qualifier_]).get(full_qualifier_, False)
-            if table_path is not False:
+            if table_path is None or isinstance(table_path, pyarrow.Table):
                 continue
             # try:
             #     namespace = self.iceberg_catalog.properties.get('namespace')