Merge pull request #30 from Snowflake-Labs/specify-target-schema

Added target schemas to read database object from

CHANGELOG.md

@@ -1,6 +1,11 @@
 # DLSync Changelog
 
 This project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
+## [2.3.0] - 2025-03-31
+### Added
+- Added support to specify target schemas for creating script from database
+- Added unit test for parsing object types
+
 ## [2.2.0] - 2025-03-24
 ### Added
 - Added support streamlit object type

README.md

@@ -2,6 +2,8 @@
 
 <img src="https://github.com/user-attachments/assets/24da3d86-f58e-4b55-8d9e-b3194117a566" height="300" title="logo" alt="logo">
 
+[![unit-test](https://github.com/Snowflake-Labs/dlsync/actions/workflows/test.yml/badge.svg)](https://github.com/Snowflake-Labs/dlsync/actions/workflows/test.yml)
+[![release](https://img.shields.io/github/release/Snowflake-Labs/dlsync.svg?style=flat)](https://github.com/Snowflake-Labs/dlsync/releases/latest)
 ---
 
 DLSync is a database change management tool designed to streamline the development and deployment of snowflake changes. 
@@ -339,11 +341,15 @@ To verify the changes use the following command:
 dlsync verify --script-root path/to/db_scripts --profile qa
 ```
 #### Create script
-This module is used to create script files for each database object. This can be used to create script files for the existing database objects. This might be helpful when you are migrating from other tools to DLSync. To achieve it first identifies the schemas inside the current database. Then for each schema retrieves the ddl of each object. Then based on the parameter profile provided it will replace the static values with the parameter keys. Then it will create the script file for each object. 
+This module is used to create script files from database. This can be used to create script files for the existing database objects. This might be helpful when you are migrating from other tools to DLSync. To achieve it first identifies the schemas inside the current database. Then for each schema retrieves the ddl of each object. Then based on the parameter profile provided it will replace the static values with the parameter keys. Then it will create the script file for each object. 
 If you have configuration tables where you want the data also to be included in the script file, you can provide the list of table names in the config file. 
 ```
 dlsync create_script --script-root path/to/db_scripts --profile uat
 ```
+You can also provide the list of schemas you want to create the script files for. If you don't provide any schemas, it will create the script files for all schemas. 
+```
+dlsync create_script --script-root path/to/db_scripts --profile uat --target-schemas schema1,schema2
+```
 
 ## Tables used by this tool
 DLSync stores script meta data, deployment history and logs in the database.

gradle.properties

@@ -1 +1 @@
-releaseVersion=2.2.0
+releaseVersion=2.3.0

src/main/java/com/snowflake/dlsync/ChangeManager.java

@@ -160,7 +160,7 @@ public class ChangeManager {
         return true;
     }
 
-    public void createAllScriptsFromDB() throws SQLException, IOException {
+    public void createAllScriptsFromDB(String targetSchemas) throws SQLException, IOException {
         log.info("Started create scripts.");
         startSync(ChangeType.CREATE_SCRIPT);
         HashSet<String> configTableWithParameter = new HashSet<>();
@@ -168,7 +168,13 @@ public class ChangeManager {
             configTableWithParameter.addAll(config.getConfigTables());
         }
         Set<String> configTables = parameterInjector.injectParameters(configTableWithParameter);
-        List<String> schemaNames = scriptRepo.getAllSchemasInDatabase(scriptRepo.getDatabaseName());
+        List<String> schemaNames = new ArrayList<>();
+        if(targetSchemas != null) {
+            schemaNames = Arrays.asList(targetSchemas.split(","));
+        }
+        else {
+            schemaNames = scriptRepo.getAllSchemasInDatabase(scriptRepo.getDatabaseName());
+        }
         int count = 0;
         for(String schema: schemaNames) {
             List<Script> scripts = scriptRepo.getAllScriptsInSchema(schema);

src/main/java/com/snowflake/dlsync/Main.java

@@ -14,7 +14,7 @@ import java.util.List;
 public class Main {
 
     public static void main(String[] args) throws SQLException {
-        log.info("DlSync change Manager started.");
+        log.info("DlSync change Manager started with args: {}", Arrays.toString(args));
         ChangeManager changeManager = null;
         ChangeType changeType = null;
         boolean onlyHashes = false;
@@ -25,6 +25,7 @@ public class Main {
             onlyHashes = commandLine.hasOption("only-hashes");
             String scriptRoot = commandLine.getOptionValue("script-root");
             String profile = commandLine.getOptionValue("profile");
+            String targetSchemas = commandLine.getOptionValue("target-schemas");
             changeManager = ChangeMangerFactory.createChangeManger(scriptRoot, profile);
             switch (changeType) {
                 case DEPLOY:
@@ -44,7 +45,7 @@ public class Main {
                     }
                     break;
                 case CREATE_SCRIPT:
-                    changeManager.createAllScriptsFromDB();
+                    changeManager.createAllScriptsFromDB(targetSchemas);
                     log.info("DLsync created all scripts from DB.");
                     break;
                 case CREATE_LINEAGE:
@@ -97,6 +98,8 @@ public class Main {
             options.addOption(scriptRoot);
             Option profile = new Option("p", "profile", true, "Profile to use");
             options.addOption(profile);
+            Option targetSchemas = new Option("t", "target-schemas", true, "Comma separated list of schemas to use");
+            options.addOption(targetSchemas);
             CommandLine commandLine = new DefaultParser().parse(options, argsWithoutCommand);
             return commandLine;
         } catch (ParseException e) {

src/test/java/com/snowflake/dlsync/MainTest.java

@@ -5,8 +5,7 @@ import com.snowflake.dlsync.models.ChangeType;
 import org.apache.commons.cli.*;
 import org.junit.jupiter.api.Test;
 
-import static org.junit.jupiter.api.Assertions.assertThrows;
-import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.junit.jupiter.api.Assertions.*;
 
 public class MainTest {
 
@@ -82,5 +81,36 @@ public class MainTest {
         assertTrue(changeType == ChangeType.DEPLOY);
     }
 
+    @Test
+    void testTargetSchema() throws ParseException {
+        String[] args = {"create_script", "--target-schemas", "schema1,schema2"};
+        CommandLine commandLine = Main.buildCommandOptions(args);
+        assertTrue(commandLine.hasOption("--target-schemas"));
+        assertTrue(commandLine.getOptionValue("target-schemas").equals("schema1,schema2"));
+
+        String[] args1 = {"create_script", "-t", "schema1,schema2"};
+        commandLine = Main.buildCommandOptions(args);
+        assertTrue(commandLine.hasOption("--target-schemas"));
+        assertTrue(commandLine.getOptionValue("target-schemas").equals("schema1,schema2"));
+
+        String[] args2 = {"create_script", "--script-root", "test/scripts", "--profile", "prod", "--target-schemas", "schema1,schema2"};
+        commandLine = Main.buildCommandOptions(args2);
+        assertTrue(commandLine.hasOption("--target-schemas"));
+        assertTrue(commandLine.getOptionValue("target-schemas").equals("schema1,schema2"));
+
+        String[] args3 = {"create_script", "--target-schemas"};
+        assertThrows(MissingArgumentException.class,  () -> Main.buildCommandOptions(args3));
+
+        String[] args4 = {"create_script", "--script-root", "test/scripts", "--profile", "prod"};
+        commandLine = Main.buildCommandOptions(args4);
+        assertTrue(!commandLine.hasOption("--target-schemas"));
+        assertTrue(commandLine.getOptionValue("--target-schemas") == null);
+
+        String[] args5 = {"create_script", "-s", "captured_scripts", "-p", "qa", "-t", "schema1,schema2"};
+        commandLine = Main.buildCommandOptions(args5);
+        assertTrue(commandLine.hasOption("--target-schemas"));
+        assertEquals(commandLine.getOptionValue("--target-schemas"), "schema1,schema2");
+    }
+
 
 }

src/test/java/com/snowflake/dlsync/parser/SqlTokenizerTest.java

@@ -2,6 +2,7 @@ package com.snowflake.dlsync.parser;
 
 import com.snowflake.dlsync.ScriptFactory;
 import com.snowflake.dlsync.models.Migration;
+import com.snowflake.dlsync.models.MigrationScript;
 import com.snowflake.dlsync.models.Script;
 import com.snowflake.dlsync.models.ScriptObjectType;
 import org.junit.jupiter.api.Test;
@@ -267,6 +268,227 @@ class SqlTokenizerTest {
         assertEquals(expected, actual, "Remove string literal assertion failed!");
     }
 
+    @Test
+    void parseScriptTypeView() {
+        String filePath = "db_scripts/db1/schema1/VIEWS/VIEW1.SQL";
+        String name = "VIEW1.SQL";
+        String scriptType = "VIEWS";
+        String content = "CREATE OR REPLACE VIEW db1.schema1.VIEW1 AS SELECT * FROM table1;";
+
+        Set<Script> scripts = SqlTokenizer.parseScript(filePath, name, scriptType, content);
+
+        assertNotNull(scripts, "Scripts should not be null");
+        assertEquals(1, scripts.size(), "There should be exactly one script parsed");
+
+        Script script = scripts.iterator().next();
+        assertEquals("VIEW1", script.getObjectName(), "Object name should be VIEW1");
+        assertEquals("db1".toUpperCase(), script.getDatabaseName(), "Database name should be db1");
+        assertEquals("schema1".toUpperCase(), script.getSchemaName(), "Schema name should be schema1");
+        assertEquals(ScriptObjectType.VIEWS, script.getObjectType(), "Object type should be VIEWS");
+        assertEquals(content, script.getContent(), "Script content should match the input content");
+    }
+
+    @Test
+    void parseScriptTypeFunction() {
+        String filePath = "db_scripts/db1/schema1/FUNCTIONS/FUNCTION1.SQL";
+        String name = "FUNCTION1.SQL";
+        String scriptType = "FUNCTIONS";
+        String content = "CREATE OR REPLACE FUNCTION db1.schema1.FUNCTION1() RETURNS STRING LANGUAGE JAVASCRIPT AS 'return \"Hello\";';";
+
+        Set<Script> scripts = SqlTokenizer.parseScript(filePath, name, scriptType, content);
+
+        assertNotNull(scripts, "Scripts should not be null");
+        assertEquals(1, scripts.size(), "There should be exactly one script parsed");
+
+        Script script = scripts.iterator().next();
+        assertEquals("FUNCTION1", script.getObjectName(), "Object name should be FUNCTION1");
+        assertEquals("db1".toUpperCase(), script.getDatabaseName(), "Database name should be db1");
+        assertEquals("schema1".toUpperCase(), script.getSchemaName(), "Schema name should be schema1");
+        assertEquals(ScriptObjectType.FUNCTIONS, script.getObjectType(), "Object type should be FUNCTIONS");
+        assertEquals(content, script.getContent(), "Script content should match the input content");
+    }
+
+    @Test
+    void parseScriptTypeProcedure() {
+        String filePath = "db_scripts/db1/schema1/PROCEDURES/PROCEDURE1.SQL";
+        String name = "PROCEDURE1.SQL";
+        String scriptType = "PROCEDURES";
+        String content = "CREATE OR REPLACE PROCEDURE db1.schema1.PROCEDURE1() RETURNS STRING LANGUAGE JAVASCRIPT AS 'return \"Hello\";';";
+
+        Set<Script> scripts = SqlTokenizer.parseScript(filePath, name, scriptType, content);
+
+        assertNotNull(scripts, "Scripts should not be null");
+        assertEquals(1, scripts.size(), "There should be exactly one script parsed");
+
+        Script script = scripts.iterator().next();
+        assertEquals("PROCEDURE1", script.getObjectName(), "Object name should be PROCEDURE1");
+        assertEquals("db1".toUpperCase(), script.getDatabaseName(), "Database name should be db1");
+        assertEquals("schema1".toUpperCase(), script.getSchemaName(), "Schema name should be schema1");
+        assertEquals(ScriptObjectType.PROCEDURES, script.getObjectType(), "Object type should be PROCEDURES");
+        assertEquals(content, script.getContent(), "Script content should match the input content");
+    }
+
+    @Test
+    void parseScriptTypeFileFormat() {
+        String filePath = "db_scripts/db1/schema1/FILE_FORMATS/FILE_FORMAT1.SQL";
+        String name = "FILE_FORMAT1.SQL";
+        String scriptType = "FILE_FORMATS";
+        String content = "CREATE OR REPLACE FILE FORMAT db1.schema1.FILE_FORMAT1 TYPE = 'CSV';";
+
+        Set<Script> scripts = SqlTokenizer.parseScript(filePath, name, scriptType, content);
+
+        assertNotNull(scripts, "Scripts should not be null");
+        assertEquals(1, scripts.size(), "There should be exactly one script parsed");
+
+        Script script = scripts.iterator().next();
+        assertEquals("FILE_FORMAT1", script.getObjectName(), "Object name should be FILE_FORMAT1");
+        assertEquals("db1".toUpperCase(), script.getDatabaseName(), "Database name should be db1");
+        assertEquals("schema1".toUpperCase(), script.getSchemaName(), "Schema name should be schema1");
+        assertEquals(ScriptObjectType.FILE_FORMATS, script.getObjectType(), "Object type should be FILE_FORMATS");
+        assertEquals(content, script.getContent(), "Script content should match the input content");
+    }
+
+    @Test
+    void parseScriptTypeTable() {
+        String filePath = "db_scripts/db1/schema1/TABLES/TABLE1.SQL";
+        String name = "TABLE1.SQL";
+        String scriptType = "TABLES";
+        String content = "---version: 0, author: dlsync \n" +
+                "CREATE OR REPLACE TABLE db1.schema1.TABLE1 (id INT, name STRING);\n" +
+                "---rollback: drop table db1.schema1.TABLE1;\n" +
+                "---verify: select * from db1.schema1.TABLE1;";
+
+        String expected_rollback = "drop table db1.schema1.TABLE1;";
+        String expected_verify = "select * from db1.schema1.TABLE1;";
+        Set<Script> scripts = SqlTokenizer.parseScript(filePath, name, scriptType, content);
+
+        assertNotNull(scripts, "Scripts should not be null");
+        assertEquals(1, scripts.size(), "There should be exactly one script parsed");
+
+        MigrationScript script = (MigrationScript) scripts.iterator().next();
+        assertEquals(0, script.getVersion(), "Version should be 0");
+        assertEquals(expected_rollback, script.getRollback(), "Rollback should match the input content");
+        assertEquals(expected_verify, script.getVerify(), "Verify should match the input content");
+        assertEquals("TABLE1", script.getObjectName(), "Object name should be TABLE1");
+        assertEquals("db1".toUpperCase(), script.getDatabaseName(), "Database name should be db1");
+        assertEquals("schema1".toUpperCase(), script.getSchemaName(), "Schema name should be schema1");
+        assertEquals(ScriptObjectType.TABLES, script.getObjectType(), "Object type should be TABLES");
+        assertEquals(content, script.getContent(), "Script content should match the input content");
+    }
+
+    @Test
+    void parseScriptTypeStream() {
+        String filePath = "db_scripts/db1/schema1/STREAMS/STREAM1.SQL";
+        String name = "STREAM1.SQL";
+        String scriptType = "STREAMS";
+
+        String content = "---version: 0, author: dlsync \n" +
+                "CREATE OR REPLACE STREAM db1.schema1.STREAM1 ON TABLE db1.schema1.TABLE1;\n" +
+                "---rollback: drop stream db1.schema1.STREAM1;\n" +
+                "---verify: select * from db1.schema1.STREAM1;";
+
+        String expected_rollback = "drop stream db1.schema1.STREAM1;";
+        String expected_verify = "select * from db1.schema1.STREAM1;";
+        Set<Script> scripts = SqlTokenizer.parseScript(filePath, name, scriptType, content);
+
+        assertNotNull(scripts, "Scripts should not be null");
+        assertEquals(1, scripts.size(), "There should be exactly one script parsed");
+
+        MigrationScript script = (MigrationScript) scripts.iterator().next();
+        assertEquals(0, script.getVersion(), "Version should be 0");
+        assertEquals(expected_rollback, script.getRollback(), "Rollback should match the input content");
+        assertEquals(expected_verify, script.getVerify(), "Verify should match the input content");
+        assertEquals("STREAM1", script.getObjectName(), "Object name should be STREAM1");
+        assertEquals("db1".toUpperCase(), script.getDatabaseName(), "Database name should be db1");
+        assertEquals("schema1".toUpperCase(), script.getSchemaName(), "Schema name should be schema1");
+        assertEquals(ScriptObjectType.STREAMS, script.getObjectType(), "Object type should be STREAMS");
+        assertEquals(content, script.getContent(), "Script content should match the input content");
+    }
+
+    @Test
+    void parseScriptTypeSequence() {
+        String filePath = "db_scripts/db1/schema1/SEQUENCES/SEQUENCE1.SQL";
+        String name = "SEQUENCE1.SQL";
+        String scriptType = "SEQUENCES";
+        String content = "---version: 0, author: dlsync \n" +
+                "CREATE OR REPLACE SEQUENCE db1.schema1.SEQUENCE1 START WITH 1 INCREMENT BY 1;\n" +
+                "---rollback: drop sequence db1.schema1.SEQUENCE1;\n" +
+                "---verify: select * from db1.schema1.SEQUENCE1;";
+
+        String expected_rollback = "drop sequence db1.schema1.SEQUENCE1;";
+        String expected_verify = "select * from db1.schema1.SEQUENCE1;";
+        Set<Script> scripts = SqlTokenizer.parseScript(filePath, name, scriptType, content);
+
+        assertNotNull(scripts, "Scripts should not be null");
+        assertEquals(1, scripts.size(), "There should be exactly one script parsed");
+
+        MigrationScript script = (MigrationScript) scripts.iterator().next();
+        assertEquals(0, script.getVersion(), "Version should be 0");
+        assertEquals(expected_rollback, script.getRollback(), "Rollback should match the input content");
+        assertEquals(expected_verify, script.getVerify(), "Verify should match the input content");
+        assertEquals("SEQUENCE1", script.getObjectName(), "Object name should be SEQUENCE1");
+        assertEquals("db1".toUpperCase(), script.getDatabaseName(), "Database name should be db1");
+        assertEquals("schema1".toUpperCase(), script.getSchemaName(), "Schema name should be schema1");
+        assertEquals(ScriptObjectType.SEQUENCES, script.getObjectType(), "Object type should be SEQUENCES");
+        assertEquals(content, script.getContent(), "Script content should match the input content");
+    }
+
+    @Test
+    void parseScriptTypeStage() {
+        String filePath = "db_scripts/db1/schema1/STAGES/STAGE1.SQL";
+        String name = "STAGE1.SQL";
+        String scriptType = "STAGES";
+        String content = "---version: 0, author: dlsync \n" +
+                "CREATE OR REPLACE STAGE db1.schema1.STAGE1 URL='s3://mybucket/mypath/';\n" +
+                "---rollback: drop stage db1.schema1.STAGE1;\n" +
+                "---verify: list @db1.schema1.STAGE1;";
+
+        String expected_rollback = "drop stage db1.schema1.STAGE1;";
+        String expected_verify = "list @db1.schema1.STAGE1;";
+        Set<Script> scripts = SqlTokenizer.parseScript(filePath, name, scriptType, content);
+
+        assertNotNull(scripts, "Scripts should not be null");
+        assertEquals(1, scripts.size(), "There should be exactly one script parsed");
+
+        MigrationScript script = (MigrationScript) scripts.iterator().next();
+        assertEquals(0, script.getVersion(), "Version should be 0");
+        assertEquals(expected_rollback, script.getRollback(), "Rollback should match the input content");
+        assertEquals(expected_verify, script.getVerify(), "Verify should match the input content");
+        assertEquals("STAGE1", script.getObjectName(), "Object name should be STAGE1");
+        assertEquals("db1".toUpperCase(), script.getDatabaseName(), "Database name should be db1");
+        assertEquals("schema1".toUpperCase(), script.getSchemaName(), "Schema name should be schema1");
+        assertEquals(ScriptObjectType.STAGES, script.getObjectType(), "Object type should be STAGES");
+        assertEquals(content, script.getContent(), "Script content should match the input content");
+    }
+
+    @Test
+    void parseScriptTypeTask() {
+        String filePath = "db_scripts/db1/schema1/TASKS/TASK1.SQL";
+        String name = "TASK1.SQL";
+        String scriptType = "TASKS";
+        String content = "---version: 0, author: dlsync \n" +
+                "CREATE OR REPLACE TASK db1.schema1.TASK1 WAREHOUSE = 'my_warehouse' SCHEDULE = 'USING CRON 0 0 * * * UTC' AS CALL my_procedure();\n" +
+                "---rollback: drop task db1.schema1.TASK1;\n" +
+                "---verify: select * from db1.schema1.TASK1;";
+
+        String expected_rollback = "drop task db1.schema1.TASK1;";
+        String expected_verify = "select * from db1.schema1.TASK1;";
+        Set<Script> scripts = SqlTokenizer.parseScript(filePath, name, scriptType, content);
+
+        assertNotNull(scripts, "Scripts should not be null");
+        assertEquals(1, scripts.size(), "There should be exactly one script parsed");
+
+        MigrationScript script = (MigrationScript) scripts.iterator().next();
+        assertEquals(0, script.getVersion(), "Version should be 0");
+        assertEquals(expected_rollback, script.getRollback(), "Rollback should match the input content");
+        assertEquals(expected_verify, script.getVerify(), "Verify should match the input content");
+        assertEquals("TASK1", script.getObjectName(), "Object name should be TASK1");
+        assertEquals("db1".toUpperCase(), script.getDatabaseName(), "Database name should be db1");
+        assertEquals("schema1".toUpperCase(), script.getSchemaName(), "Schema name should be schema1");
+        assertEquals(ScriptObjectType.TASKS, script.getObjectType(), "Object type should be TASKS");
+        assertEquals(content, script.getContent(), "Script content should match the input content");
+    }
+
     @Test
     void parseScriptTypeStreamlit() {
         String filePath = "db_scripts/db1/schema1/STREAMLITS/STREAMLIT1.SQL";
@@ -290,4 +512,6 @@ class SqlTokenizerTest {
         assertEquals(content, script.getContent(), "Script content should match the input content");
 
     }
+
+
 }