use fakesnow for better transpilation

buremba · Aug 4, 2024 · 147731f · 147731f
1 parent c683c28
commit 147731f
Show file tree

Hide file tree

Showing 12 changed files with 463 additions and 337 deletions.
diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml
@@ -9,6 +9,7 @@ jobs:
   pip:
     name: Release PyPI
     runs-on: ubuntu-latest
+    environment: deploy
     steps:
       - name: Check out the repository
         uses: actions/checkout@v3

diff --git a/.gitignore b/.gitignore
@@ -103,4 +103,5 @@ celerybeat.pid
 .metabase/*
 .clickhouse/*
 .DS_STORE
-.certs/*
+.certs/*
+.rill/*
diff --git a/README.md b/README.md
@@ -143,40 +143,6 @@ It gives you free https connection to your local server and it's the default hos
 For Catalog, [Snowflake](https://docs.snowflake.com/en/sql-reference/sql/create-iceberg-table-snowflake) and [Object Store](https://docs.snowflake.com/en/sql-reference/sql/create-iceberg-table-iceberg-files) catalogs are supported at the moment.
 For Data lake, S3 and GCS supported.
 
-## Can't query all Snowflake types locally
-
-Here is a Markdown table of some Snowflake data types with a "Supported" column. The checkbox indicates whether the type is supported or not. Please replace the checkboxes with the correct values according to your project's support for each data type.
-
-| Snowflake Data Type | Supported                      |
-| --- |--------------------------------|
-| NUMBER | &check;                        |
-| DECIMAL | &check;                        |
-| INT | &check;                        |
-| BIGINT | &check;                        |
-| SMALLINT | &check;                        |
-| TINYINT | &check;                        |
-| FLOAT | &check;                        |
-| DOUBLE | &check;                        |
-| VARCHAR | &check;                        |
-| CHAR | &check;                        |
-| STRING | &check;                        |
-| TEXT | &check;                        |
-| BOOLEAN | &check;                        |
-| DATE | &check;                        |
-| DATETIME | &check;                        |
-| TIME | &check;                        |
-| TIMESTAMP | &check;                        |
-| TIMESTAMP_LTZ | &cross; ¹                       |
-| TIMESTAMP_NTZ | &cross; ¹ |
-| TIMESTAMP_TZ | &cross;¹                       |
-| VARIANT | &check;                     |
-| OBJECT | &check;                          |
-| ARRAY | &check;                           |
-| GEOGRAPHY | &cross; ¹                         |
-| VECTOR | &cross; ¹                          |
-
-¹: No Support in DuckDB yet.
-
 ## Can't query native Snowflake tables locally
 
 UniverSQL doesn't support querying native Snowflake tables as they're not accessible from outside of Snowflake. If you try to query a Snowflake table directly, it will return an error.

diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -17,7 +17,6 @@ uvicorn = "^0.30.1"
 snowflake-connector-python = {extras = ["pandas", "secure-local-storage"], version = "^3.11.0"}
 eval-type-backport = "^0.2.0"
 pip-system-certs = "^4.0"
-sqlglot = {extras = ["rs"], version = "^25.5.1"}
 chdb = "^1.4.1"
 fsspec = "^2024.6.1"
 gcsfs = "^2024.6.1"
@@ -29,6 +28,7 @@ pyiceberg = "^0.7.0"
 sqlalchemy = "^2.0.31"
 fastapi-utils = "^0.7.0"
 
+fakesnow = "^0.9.20"
 [tool.poetry.dev-dependencies]
 pylint = ">=2.11.1"
 

diff --git a/tests/sqlglot_tests.py b/tests/sqlglot_tests.py
@@ -1,5 +1,11 @@
+import time
+import pyarrow as pa
+
 import duckdb
 import sqlglot
+from fakesnow.fakes import FakeSnowflakeCursor, FakeSnowflakeConnection
+
+from universql.warehouse.duckdb import fix_snowflake_to_duckdb_types
 
 # queries = sqlglot.parse("""
 # SET tables = (SHOW TABLES);
@@ -8,17 +14,38 @@
 # """, read="snowflake")
 
 
-query = sqlglot.parse_one("""
-SET stmt = $$
-    SELECT PI();
-$$;
-
-SELECT *, 1 FROM $stmt;
-""", dialect="snowflake")
-
-zsql = query.sql(dialect="duckdb")
-
-query = duckdb.sql(query.sql())
-
-
-print(sql)
+# query = sqlglot.parse_one("""
+# SET stmt = $$
+#     SELECT PI();
+# $$;
+#
+# SELECT *, 1 FROM $stmt;
+# """, dialect="snowflake")
+
+fields = [
+    pa.field("epoch", nullable=False, type=pa.int64()),
+    pa.field("fraction", nullable=False, type=pa.int32()),
+    pa.field("timezone", nullable=False, type=pa.int32()),
+]
+pa_type = pa.struct(fields)
+pa.StructArray.from_arrays(arrays=[pa.array([1, 2, 3], type=pa.int64()), pa.array([1, 2, 3], type=pa.int32()), pa.array([1, 2, 3], type=pa.int32())], fields=fields)
+
+query = """
+SELECT
+  CAST('2023-01-01 10:34:56 +00:00' AS TIMESTAMPLTZ) AS sample_timestamp_ltz,
+  CAST('2023-01-01 11:34:56' AS TIMESTAMP) AS sample_timestamp_ntz,
+  CAST('2023-01-01 12:34:56 +00:00' AS TIMESTAMPTZ) AS sample_timestamp_tz,
+  CAST(JSON('{"key":"value"}') /* Semi-structured data types */ AS VARIANT) AS sample_variant,
+"""
+
+start = time.time()
+for i in range(10):
+    con = duckdb.connect(f"ali/{i}")
+    con.execute("CREATE TABLE test (a int, b int)")
+print(time.time() - start)
+ast = sqlglot.parse_one(query, dialect="duckdb")
+transformed_ast = ast.transform(fix_snowflake_to_duckdb_types)
+query = transformed_ast.sql(dialect="duckdb", pretty=True)
+print(query)
+response = duckdb.sql(query)
+print(response.show())
diff --git a/universql/catalog/__init__.py b/universql/catalog/__init__.py
@@ -6,7 +6,6 @@
 from snowflake.connector.options import pyarrow
 
 from universql import util
-from universql.lake.cloud import CACHE_DIRECTORY_KEY
 from universql.util import Catalog
 
 
@@ -25,7 +24,7 @@ def get_catalog(context: dict, query_id: str, credentials: dict):
 
 class Cursor(ABC):
     @abstractmethod
-    def execute(self, ast: sqlglot.exp.Expression) -> None:
+    def execute(self, ast: typing.Optional[sqlglot.exp.Expression], raw_query : str) -> None:
         pass
 
     @abstractmethod