Add schema support for bal persist sql modules

ballerina/annotations.bal

@@ -24,6 +24,16 @@ public type NameConfig record {|
 # The Annotation used to specify the mapping of an entity/field to a database table/column.
 public annotation NameConfig Name on type, record field;
 
+# Groups the entity to a specific schema in the database.
+#
+# + value - name of the schema in the database
+public type SchemaConfig record {|
+    string value;
+|};
+
+# The Annotation used to specify the schema of an entity in the database.
+public annotation SchemaConfig Schema on type;
+
 # Marks the entity field as an index field.
 #
 # + name - specify a single index name or an array of index names

ballerina/build.gradle

@@ -206,8 +206,7 @@ task stopMySQLTestDockerContainer() {
 task createMSSQLTestDockerImage(type: Exec) {
    if (!Os.isFamily(Os.FAMILY_WINDOWS)) {
        def standardOutput = new ByteArrayOutputStream()
-       commandLine 'sh', '-c', "docker build -f $project.projectDir/tests/resources/mssql/Dockerfile -t ballerina-persist-mssql" +
-               " -q $project.projectDir/tests/resources/mssql/"
+       commandLine 'sh', '-c', "docker build -f $project.projectDir/tests/resources/mssql/Dockerfile -t ballerina-persist-mssql -q $project.projectDir/tests/resources/mssql/"
        doLast {
            checkExecResult(executionResult, 'Error', standardOutput)
            sleep(10 * 1000)

ballerina/metadata_types.bal

@@ -18,12 +18,14 @@
 #
 # + entityName - Name of the entity
 # + tableName - Table name of the entity
+# + schemaName - Schema name of the entity
 # + fieldMetadata - Metadata of all the fields of the entity
 # + keyFields - Names of the identity fields
 # + joinMetadata - Metadata of the fields that are used for `JOIN` operations
 public type SQLMetadata record {|
     string entityName;
     string tableName;
+    string schemaName?;
     map<FieldMetadata> fieldMetadata;
     string[] keyFields;
     map<JoinMetadata> joinMetadata?;
@@ -67,7 +69,8 @@ public type RelationMetadata record {|
 # Only used by the generated persist clients and `persist:SQLClient`.
 #
 # + entity - The name of the entity that is being joined  
-# + fieldName - The name of the field in the `entity` that is being joined  
+# + fieldName - The name of the field in the `entity` that is being joined
+# + refSchema - The name of the SQL schema to be joined
 # + refTable - The name of the SQL table to be joined  
 # + refColumns - The names of the referenced columns of the referenced table
 # + joinColumns - The names of the join columns
@@ -78,6 +81,7 @@ public type RelationMetadata record {|
 public type JoinMetadata record {|
     typedesc<record {}> entity;
     string fieldName;
+    string refSchema?;
     string refTable;
     string[] refColumns;
     string[] joinColumns;

ballerina/sql_client.bal

@@ -25,6 +25,7 @@ public isolated client class SQLClient {
 
     private final string & readonly entityName;
     private final string & readonly tableName;
+    private final string? & readonly schemaName;
     private final map<FieldMetadata> & readonly fieldMetadata;
     private final string[] & readonly keyFields;
     private final map<JoinMetadata> & readonly joinMetadata;
@@ -40,6 +41,7 @@ public isolated client class SQLClient {
         self.tableName = metadata.tableName;
         self.fieldMetadata = metadata.fieldMetadata;
         self.keyFields = metadata.keyFields;
+        self.schemaName = metadata.schemaName;
         self.dbClient = dbClient;
         if metadata.joinMetadata is map<JoinMetadata> {
             self.joinMetadata = <map<JoinMetadata> & readonly>metadata.joinMetadata;
@@ -457,12 +459,12 @@ public isolated client class SQLClient {
 
     private isolated function getInsertQueries(record {}[] insertRecords) returns sql:ParameterizedQuery[] {
         return from record {} insertRecord in insertRecords
-            select sql:queryConcat(`INSERT INTO `, stringToParameterizedQuery(self.escape(self.tableName)), ` (`, self.getInsertColumnNames(), ` ) `, `VALUES `, self.getInsertQueryParams(insertRecord));
+            select sql:queryConcat(`INSERT INTO `, stringToParameterizedQuery(self.getTableName()), ` (`, self.getInsertColumnNames(), ` ) `, `VALUES `, self.getInsertQueryParams(insertRecord));
     }
 
     private isolated function getSelectQuery(string[] selectableFields) returns sql:ParameterizedQuery {
         return sql:queryConcat(
-            `SELECT `, self.getSelectColumnNames(selectableFields), ` FROM `, stringToParameterizedQuery(self.escape(self.tableName)), ` AS `, stringToParameterizedQuery(self.escape(self.entityName))
+            `SELECT `, self.getSelectColumnNames(selectableFields), ` FROM `, stringToParameterizedQuery(self.getTableName()), ` AS `, stringToParameterizedQuery(self.escape(self.entityName))
         );
     }
 
@@ -472,16 +474,25 @@ public isolated client class SQLClient {
 
     private isolated function getUpdateQuery(record {} updateRecord) returns sql:ParameterizedQuery|persist:Error {
         if self.dataSourceSpecifics == MSSQL_SPECIFICS {
-            return sql:queryConcat(`UPDATE `, stringToParameterizedQuery(self.escape(self.entityName)), ` SET `, check self.getSetClauses(updateRecord), ` FROM `, stringToParameterizedQuery(self.escape(self.tableName)), ` `, stringToParameterizedQuery(self.escape(self.entityName)));
+            return sql:queryConcat(`UPDATE `, stringToParameterizedQuery(self.escape(self.entityName)), ` SET `, check self.getSetClauses(updateRecord), ` FROM `, stringToParameterizedQuery(self.getTableName()), ` `, stringToParameterizedQuery(self.escape(self.entityName)));
         }
-        return sql:queryConcat(`UPDATE `, stringToParameterizedQuery(self.escape(self.tableName)), ` AS `, stringToParameterizedQuery(self.escape(self.entityName)), ` SET `, check self.getSetClauses(updateRecord));
+        return sql:queryConcat(`UPDATE `, stringToParameterizedQuery(self.getTableName()), ` AS `, stringToParameterizedQuery(self.escape(self.entityName)), ` SET `, check self.getSetClauses(updateRecord));
     }
 
     private isolated function getDeleteQuery() returns sql:ParameterizedQuery {
         if self.dataSourceSpecifics == MSSQL_SPECIFICS {
-            return sql:queryConcat(`DELETE `, stringToParameterizedQuery(self.escape(self.entityName)), ` FROM `, stringToParameterizedQuery(self.escape(self.tableName)), ` AS `, stringToParameterizedQuery(self.escape(self.entityName)));
+            return sql:queryConcat(`DELETE `, stringToParameterizedQuery(self.escape(self.entityName)), ` FROM `, stringToParameterizedQuery(self.getTableName()), ` AS `, stringToParameterizedQuery(self.escape(self.entityName)));
         }
-        return sql:queryConcat(`DELETE FROM `, stringToParameterizedQuery(self.escape(self.tableName)), ` AS `, stringToParameterizedQuery(self.escape(self.entityName)));
+        return sql:queryConcat(`DELETE FROM `, stringToParameterizedQuery(self.getTableName()), ` AS `, stringToParameterizedQuery(self.escape(self.entityName)));
+    }
+
+    // Constructs table name with schema name if schema name is available
+    private isolated function getTableName() returns string {
+        string? schemaName = self.schemaName;
+        if schemaName is () {
+            return self.escape(self.tableName);
+        }
+        return string `${schemaName}.${self.escape(self.tableName)}`;
     }
 
     private isolated function getJoinFields(string[] include) returns string[] {
@@ -512,7 +523,9 @@ public isolated client class SQLClient {
 
     private isolated function getJoinQuery(string joinKey) returns sql:ParameterizedQuery|persist:Error {
         JoinMetadata joinMetadata = self.joinMetadata.get(joinKey);
-        return sql:queryConcat(` LEFT JOIN `, stringToParameterizedQuery(self.escape(joinMetadata.refTable) + " " + self.escape(joinKey)),
+        string refSchema = joinMetadata.refSchema ?: "";
+        refSchema = joinMetadata.refSchema != () ? string `${refSchema}.` : refSchema;
+        return sql:queryConcat(` LEFT JOIN `, stringToParameterizedQuery(string `${refSchema}${self.escape(joinMetadata.refTable)} ${self.escape(joinKey)}`),
                                 ` ON `, check self.getJoinFilters(joinKey, joinMetadata.refColumns, <string[]>joinMetadata.joinColumns));
     }
 

ballerina/tests/Config.toml

@@ -27,3 +27,29 @@ database="test"
 url="jdbc:h2:./build/test"
 user="sa"
 password=""
+
+[mssqlWithSchema]
+user="sa"
+password="Test123#"
+host="localhost"
+port=1433
+database="testschema"
+defaultSchema="persist"
+
+[postgresqlWithSchema]
+user="postgres"
+password="postgres"
+host="localhost"
+port=5432
+database="test"
+defaultSchema="persist"
+
+[h2WithSchema]
+url="jdbc:h2:./build/test"
+user="sa"
+password=""
+defaultSchema="persist"
+
+[[ballerina.log.modules]]
+name = "ballerina/persist"
+level = "DEBUG"

ballerina/tests/h2_api_subscription_client.bal

@@ -30,7 +30,7 @@ public isolated client class H2ApimClient {
 
     private final map<SQLClient> persistClients;
 
-    private final record {|SQLMetadata...;|} & readonly metadata = {
+    private final record {|SQLMetadata...;|} metadata = {
         [SUBSCRIPTION]: {
             entityName: "Subscription",
             tableName: "Subscription",
@@ -72,8 +72,8 @@ public isolated client class H2ApimClient {
         }
         self.dbClient = dbClient;
         self.persistClients = {
-            [SUBSCRIPTION]: check new (dbClient, self.metadata.get(SUBSCRIPTION), H2_SPECIFICS),
-            [API_METADATA]: check new (dbClient, self.metadata.get(API_METADATA), H2_SPECIFICS)
+            [SUBSCRIPTION]: check new (dbClient, self.metadata.get(SUBSCRIPTION).cloneReadOnly(), H2_SPECIFICS),
+            [API_METADATA]: check new (dbClient, self.metadata.get(API_METADATA).cloneReadOnly(), H2_SPECIFICS)
         };
     }
 

ballerina/tests/h2_hospital_persist_client.bal

@@ -34,7 +34,7 @@ public isolated client class H2HospitalClient {
 
     private final map<SQLClient> persistClients;
 
-    private final record {|SQLMetadata...;|} & readonly metadata = {
+    private final record {|SQLMetadata...;|} metadata = {
         [APPOINTMENT]: {
             entityName: "Appointment",
             tableName: "appointment",
@@ -110,10 +110,30 @@ public isolated client class H2HospitalClient {
             return <persist:Error>error(dbClient.message());
         }
         self.dbClient = dbClient;
+        // Update the metadata with the schema name
+        if h2.defaultSchema != () {
+            lock {
+                foreach string key in self.metadata.keys() {
+                    SQLMetadata metadata = self.metadata.get(key);
+                    if metadata.schemaName == () {
+                        metadata.schemaName = h2.defaultSchema;
+                    }
+                    map<JoinMetadata>? joinMetadataMap = metadata.joinMetadata;
+                    if joinMetadataMap != () {
+                        foreach string joinKey in joinMetadataMap.keys() {
+                            JoinMetadata joinMetadata = joinMetadataMap.get(joinKey);
+                            if joinMetadata.refSchema == () {
+                                joinMetadata.refSchema = h2.defaultSchema;
+                            }
+                        }
+                    }
+                }
+            }
+        }
         self.persistClients = {
-            [APPOINTMENT]: check new (dbClient, self.metadata.get(APPOINTMENT), H2_SPECIFICS),
-            [PATIENT]: check new (dbClient, self.metadata.get(PATIENT), H2_SPECIFICS),
-            [DOCTOR]: check new (dbClient, self.metadata.get(DOCTOR), H2_SPECIFICS)
+            [APPOINTMENT]: check new (dbClient, self.metadata.get(APPOINTMENT).cloneReadOnly(), H2_SPECIFICS),
+            [PATIENT]: check new (dbClient, self.metadata.get(PATIENT).cloneReadOnly(), H2_SPECIFICS),
+            [DOCTOR]: check new (dbClient, self.metadata.get(DOCTOR).cloneReadOnly(), H2_SPECIFICS)
         };
     }