Reuse universql servers in e2e

buremba · Jan 8, 2025 · 6a0e391 · 6a0e391
1 parent 93e4f66
commit 6a0e391
Show file tree

Hide file tree

Showing 4 changed files with 40 additions and 30 deletions.
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
@@ -1,4 +1,4 @@
-name: Release
+name: Test
 
 on:
   push:

diff --git a/tests/integration/extract.py b/tests/integration/extract.py
@@ -39,6 +39,7 @@ def test_simple_select(self):
             universql_result = execute_query(conn, SIMPLE_QUERY)
             print(universql_result)
 
+    @pytest.mark.skip(reason="Stages are not implemented yet")
     def test_from_stage(self):
         with universql_connection() as conn:
             universql_result = execute_query(conn, "select count(*) from @stage/iceberg_stage")
@@ -67,9 +68,9 @@ def test_switch_schema(self):
             universql_result = execute_query(conn, "SHOW SCHEMAS")
             assert universql_result.num_rows > 0, f"The query did not return any rows!"
 
-    def test_in_schema(self):
-        with universql_connection(schema="public", warehouse="local()") as conn:
-            universql_result = execute_query(conn, "select count(*) from table_in_public_schema")
+    def test_in_database(self):
+        with universql_connection(database="public") as conn:
+            universql_result = execute_query(conn, "select * from information_schema.columns")
             print(universql_result)
 
     def test_qualifiers(self):

diff --git a/tests/integration/utils.py b/tests/integration/utils.py
@@ -85,6 +85,7 @@ def snowflake_connection(**properties) -> Generator:
     yield conn
     conn.close()
 
+server_cache = {}
 
 @contextmanager
 def universql_connection(**properties) -> SnowflakeConnection:
@@ -94,36 +95,45 @@ def universql_connection(**properties) -> SnowflakeConnection:
     if SNOWFLAKE_CONNECTION_NAME not in connections:
         raise pytest.fail(f"Snowflake connection '{SNOWFLAKE_CONNECTION_NAME}' not found in config")
     connection = connections[SNOWFLAKE_CONNECTION_NAME]
-    from universql.main import snowflake
-    with socketserver.TCPServer(("localhost", 0), None) as s:
-        free_port = s.server_address[1]
-
-    def start_universql():
-        runner = CliRunner()
-        try:
-            invoke = runner.invoke(snowflake,
-                                   [
-                                       '--account', connection.get('account'),
-                                       '--port', free_port, '--catalog', 'snowflake',
-                                       # AWS_DEFAULT_PROFILE env can be used to pass AWS profile
-                                   ],
-                                   )
-        except Exception as e:
-            pytest.fail(e)
-
-        if invoke.exit_code != 0:
-            pytest.fail("Unable to start Universql")
-
-    thread = threading.Thread(target=start_universql)
-    thread.daemon = True
-    thread.start()
+    account = connection.get('account')
+
+    if account in server_cache:
+        uni_string = {"host": LOCALHOSTCOMPUTING_COM, "port": server_cache[account]} | properties
+    else:
+        from universql.main import snowflake
+        with socketserver.TCPServer(("localhost", 0), None) as s:
+            free_port = s.server_address[1]
+        print(f"Reusing existing server running on port {free_port} for account {account}")
+
+        def start_universql():
+            runner = CliRunner()
+            try:
+                invoke = runner.invoke(snowflake,
+                                       [
+                                           '--account', account,
+                                           '--port', free_port, '--catalog', 'snowflake',
+                                           # AWS_DEFAULT_PROFILE env can be used to pass AWS profile
+                                       ],
+                                       )
+            except Exception as e:
+                pytest.fail(e)
+
+            if invoke.exit_code != 0:
+                pytest.fail("Unable to start Universql")
+
+
+        print(f"Starting running on port {free_port} for account {account}")
+        thread = threading.Thread(target=start_universql)
+        thread.daemon = True
+        thread.start()
+        server_cache[account] = free_port
     # with runner.isolated_filesystem():
-    uni_string = {"host": LOCALHOSTCOMPUTING_COM, "port": free_port} | properties
+        uni_string = {"host": LOCALHOSTCOMPUTING_COM, "port": free_port} | properties
 
     try:
         connect = snowflake_connect(connection_name=SNOWFLAKE_CONNECTION_NAME, **uni_string)
         yield connect
-    finally:  # Force stop the thread
+    finally:
         connect.close()
 
 

diff --git a/tests/snow_client.py b/tests/snow_client.py
@@ -35,7 +35,6 @@ def measure_all(cur, query):
 def tt(max_iteration):
     connection = snowflake.connector.connect(connection_name="ryan_snowflake",
                                              port='8084', host='localhostcomputing.com',
-                                             warehouse='local()',
                                              database='test',
                                              # host='4ho74nvv4nyxhqxmcxrnpiid2m0bpcet.lambda-url.us-east-1.on.aws',
                                              )