Improve implementation of storage of transpiled files

src/databricks/labs/remorph/transpiler/execute.py

@@ -33,7 +33,7 @@
 logger = logging.getLogger(__name__)
 
 
-async def _process_file(
+async def _process_one_file(
     config: TranspileConfig,
     validator: Validator | None,
     transpiler: TranspileEngine,
@@ -71,29 +71,23 @@ async def _process_file(
     return transpile_result.success_count, error_list
 
 
-async def _process_directory(
+async def _process_many_files(
     config: TranspileConfig,
     validator: Validator | None,
     transpiler: TranspileEngine,
-    root: Path,
-    base_root: Path,
+    output_folder: Path,
     files: list[Path],
 ) -> tuple[int, list[TranspileError]]:
 
-    output_folder = config.output_folder
-    output_folder_base = root / ("transpiled" if output_folder is None else base_root)
-    make_dir(output_folder_base)
-
     counter = 0
     all_errors: list[TranspileError] = []
 
     for file in files:
         logger.info(f"Processing file :{file}")
         if not is_sql_file(file):
             continue
-
-        output_file_name = output_folder_base / file.name
-        success_count, error_list = await _process_file(config, validator, transpiler, file, output_file_name)
+        output_file_name = output_folder / file.name
+        success_count, error_list = await _process_one_file(config, validator, transpiler, file, output_file_name)
         counter = counter + success_count
         all_errors.extend(error_list)
 
@@ -104,15 +98,17 @@ async def _process_input_dir(config: TranspileConfig, validator: Validator | Non
     error_list = []
     file_list = []
     counter = 0
-    input_source = str(config.input_source)
-    input_path = Path(input_source)
-    for root, _, files in dir_walk(input_path):
-        base_root = Path(str(root).replace(input_source, ""))
-        folder = str(input_path.resolve().joinpath(base_root))
-        msg = f"Processing for sqls under this folder: {folder}"
-        logger.info(msg)
+    input_path = config.input_path
+    output_folder = config.output_path
+    if output_folder is None:
+        output_folder = input_path.parent / "transpiled"
+    make_dir(output_folder)
+    for source_dir, _, files in dir_walk(input_path):
+        relative_path = cast(Path, source_dir).relative_to(input_path)
+        transpiled_dir = output_folder / relative_path
+        logger.info(f"Transpiling sql files from folder: {source_dir!s} into {transpiled_dir!s}")
         file_list.extend(files)
-        no_of_sqls, errors = await _process_directory(config, validator, transpiler, root, base_root, files)
+        no_of_sqls, errors = await _process_many_files(config, validator, transpiler, transpiled_dir, files)
         counter = counter + no_of_sqls
         error_list.extend(errors)
     return TranspileStatus(file_list, counter, error_list)
@@ -126,23 +122,21 @@ async def _process_input_file(
         logger.warning(msg)
         # silently ignore non-sql files
         return TranspileStatus([], 0, [])
-    msg = f"Processing sql from this file: {config.input_source}"
+    msg = f"Transpiling sql file: {config.input_path!s}"
     logger.info(msg)
-    if config.output_path is None:
+    output_path = config.output_path
+    if output_path is None:
         output_path = config.input_path.parent / "transpiled"
-    else:
-        output_path = config.output_path
-
     make_dir(output_path)
     output_file = output_path / config.input_path.name
-    no_of_sqls, error_list = await _process_file(config, validator, transpiler, config.input_path, output_file)
+    no_of_sqls, error_list = await _process_one_file(config, validator, transpiler, config.input_path, output_file)
     return TranspileStatus([config.input_path], no_of_sqls, error_list)
 
 
 @timeit
 async def transpile(
     workspace_client: WorkspaceClient, engine: TranspileEngine, config: TranspileConfig
-) -> tuple[list[dict[str, Any]], list[TranspileError]]:
+) -> tuple[dict[str, Any], list[TranspileError]]:
     await engine.initialize(config)
     status, errors = await _do_transpile(workspace_client, engine, config)
     await engine.shutdown()
@@ -151,7 +145,7 @@ async def transpile(
 
 async def _do_transpile(
     workspace_client: WorkspaceClient, engine: TranspileEngine, config: TranspileConfig
-) -> tuple[list[dict[str, Any]], list[TranspileError]]:
+) -> tuple[dict[str, Any], list[TranspileError]]:
     """
     [Experimental] Transpiles the SQL queries from one dialect to another.
 
@@ -163,8 +157,6 @@ async def _do_transpile(
         logger.error("Input SQL path is not provided.")
         raise ValueError("Input SQL path is not provided.")
 
-    status = []
-
     validator = None
     if not config.skip_validation:
         sql_backend = db_sql.get_sql_backend(workspace_client)
@@ -194,17 +186,16 @@ async def _do_transpile(
         with cast(Path, error_log_path).open("a", encoding="utf-8") as e:
             e.writelines(f"{err!s}\n" for err in result.error_list)
 
-    status.append(
-        {
-            "total_files_processed": len(result.file_list),
-            "total_queries_processed": result.no_of_transpiled_queries,
-            "no_of_sql_failed_while_analysing": result.analysis_error_count,
-            "no_of_sql_failed_while_parsing": result.parsing_error_count,
-            "no_of_sql_failed_while_generating": result.generation_error_count,
-            "no_of_sql_failed_while_validating": result.validation_error_count,
-            "error_log_file": str(error_log_path),
-        }
-    )
+    status = {
+        "total_files_processed": len(result.file_list),
+        "total_queries_processed": result.no_of_transpiled_queries,
+        "analysis_error_count": result.analysis_error_count,
+        "parsing_error_count": result.parsing_error_count,
+        "validation_error_count": result.validation_error_count,
+        "generation_error_count": result.generation_error_count,
+        "error_log_file": str(error_log_path),
+    }
+
     return status, result.error_list
 
 

src/databricks/labs/remorph/transpiler/sqlglot/sqlglot_engine.py

@@ -80,7 +80,7 @@ async def transpile(
         read_dialect = get_dialect(source_dialect)
         error: TranspileError | None = self._check_supported(read_dialect, source_code, file_path)
         if error:
-            return TranspileResult(str(file_path), 1, [error])
+            return TranspileResult(source_code, 1, [error])
         write_dialect = get_dialect(target_dialect)
         try:
             transpiled_expressions = transpile(

tests/unit/conftest.py

@@ -1,5 +1,6 @@
 import io
 import re
+import shutil
 from pathlib import Path
 from collections.abc import Sequence
 from unittest.mock import create_autospec
@@ -15,6 +16,7 @@
 from databricks.sdk.errors import NotFound
 
 from databricks.labs.remorph.config import TranspileConfig
+from databricks.labs.remorph.helpers.file_utils import make_dir
 from databricks.labs.remorph.transpiler.sqlglot.dialect_utils import SQLGLOT_DIALECTS
 from databricks.labs.remorph.transpiler.sqlglot.generator.databricks import Databricks
 from databricks.labs.remorph.transpiler.sqlglot.parsers.snowflake import Snowflake
@@ -270,3 +272,143 @@ def download(path: str) -> io.StringIO | io.BytesIO:
     config.cluster_id = None
     workspace_client.config = config
     return workspace_client
+
+
+def safe_remove_dir(dir_path: Path):
+    if dir_path.exists():
+        shutil.rmtree(dir_path)
+
+
+def safe_remove_file(file_path: Path):
+    if file_path.exists():
+        file_path.unlink()
+
+
+def write_data_to_file(path: Path, content: str):
+    with path.open("w") as writable:
+        # added encoding to avoid UnicodeEncodeError while writing to file for token error test
+        writable.write(content.encode("utf-8", "ignore").decode("utf-8"))
+
+
+@pytest.fixture
+def input_source(tmp_path: Path):
+    input_dir = tmp_path / "remorph_source"
+    query_1_sql_file = input_dir / "query1.sql"
+    query_2_sql_file = input_dir / "query2.sql"
+    query_3_sql_file = input_dir / "query3.sql"
+    query_4_sql_file = input_dir / "query4.sql"
+    query_5_sql_file = input_dir / "query5.sql"
+    stream_1_sql_file = input_dir / "stream1.sql"
+    call_center_ddl_file = input_dir / "call_center.ddl"
+    file_text = input_dir / "file.txt"
+    safe_remove_dir(input_dir)
+    make_dir(input_dir)
+
+    query_1_sql = """select  i_manufact, sum(ss_ext_sales_price) ext_price from date_dim, store_sales where
+    d_date_sk = ss_sold_date_sk and substr(ca_zip,1,5) <> substr(s_zip,1,5) group by i_manufact order by i_manufact
+    limit 100 ;"""
+
+    call_center_ddl = """create table call_center
+        (
+            cc_call_center_sk         int                           ,
+            cc_call_center_id         varchar(16)
+        )
+
+         CLUSTER BY(cc_call_center_sk)
+         """
+
+    query_2_sql = """select wswscs.d_week_seq d_week_seq1,sun_sales sun_sales1,mon_sales mon_sales1 from wswscs,
+    date_dim where date_dim.d_week_seq = wswscs.d_week_seq and d_year = 2001"""
+
+    query_3_sql = """with wscs as
+     (select sold_date_sk
+            ,sales_price
+      from (select ws_sold_date_sk sold_date_sk
+                  ,ws_ext_sales_price sales_price
+            from web_sales
+            union all
+            select cs_sold_date_sk sold_date_sk
+                  ,cs_ext_sales_price sales_price
+            from catalog_sales)),
+     wswscs as
+     (select d_week_seq,
+            sum(case when (d_day_name='Sunday') then sales_price else null end) sun_sales,
+            sum(case when (d_day_name='Monday') then sales_price else null end) mon_sales,
+            sum(case when (d_day_name='Tuesday') then sales_price else  null end) tue_sales,
+            sum(case when (d_day_name='Wednesday') then sales_price else null end) wed_sales,
+            sum(case when (d_day_name='Thursday') then sales_price else null end) thu_sales,
+            sum(case when (d_day_name='Friday') then sales_price else null end) fri_sales,
+            sum(case when (d_day_name='Saturday') then sales_price else null end) sat_sales
+     from wscs
+         ,date_dim
+     where d_date_sk = sold_date_sk
+     group by d_week_seq)
+     select d_week_seq1
+           ,round(sun_sales1/sun_sales2,2)
+           ,round(mon_sales1/mon_sales2,2)
+           ,round(tue_sales1/tue_sales2,2)
+           ,round(wed_sales1/wed_sales2,2)
+           ,round(thu_sales1/thu_sales2,2)
+           ,round(fri_sales1/fri_sales2,2)
+           ,round(sat_sales1/sat_sales2,2)
+     from
+     (select wswscs.d_week_seq d_week_seq1
+            ,sun_sales sun_sales1
+            ,mon_sales mon_sales1
+            ,tue_sales tue_sales1
+            ,wed_sales wed_sales1
+            ,thu_sales thu_sales1
+            ,fri_sales fri_sales1
+            ,sat_sales sat_sales1
+      from wswscs,date_dim
+      where date_dim.d_week_seq = wswscs.d_week_seq and
+            d_year = 2001) y,
+     (select wswscs.d_week_seq d_week_seq2
+            ,sun_sales sun_sales2
+            ,mon_sales mon_sales2
+            ,tue_sales tue_sales2
+            ,wed_sales wed_sales2
+            ,thu_sales thu_sales2
+            ,fri_sales fri_sales2
+            ,sat_sales sat_sales2
+      from wswscs
+          ,date_dim
+      where date_dim.d_week_seq = wswscs.d_week_seq2 and
+            d_year = 2001+1) z
+     where d_week_seq1=d_week_seq2-53
+     order by d_week_seq1;
+     """
+
+    stream_1_sql = """CREATE STREAM unsupported_stream AS SELECT * FROM some_table;"""
+
+    query_4_sql = """create table(
+    col1 int
+    col2 string
+    );"""
+
+    query_5_sql = """1SELECT * from ~v\ud83d' table;"""
+
+    write_data_to_file(query_1_sql_file, query_1_sql)
+    write_data_to_file(call_center_ddl_file, call_center_ddl)
+    write_data_to_file(query_2_sql_file, query_2_sql)
+    write_data_to_file(query_3_sql_file, query_3_sql)
+    write_data_to_file(query_4_sql_file, query_4_sql)
+    write_data_to_file(query_5_sql_file, query_5_sql)
+    write_data_to_file(stream_1_sql_file, stream_1_sql)
+    write_data_to_file(file_text, "This is a test file")
+    yield input_dir
+    safe_remove_dir(input_dir)
+
+
+@pytest.fixture
+def output_folder(tmp_path: Path):
+    output_dir = tmp_path / "remorph_transpiled"
+    yield output_dir
+    safe_remove_dir(output_dir)
+
+
+@pytest.fixture
+def error_file(tmp_path: Path):
+    file_path = tmp_path / "transpile_errors.lst"
+    yield file_path
+    safe_remove_file(file_path)