Skip to content

Commit

Permalink
Merge pull request #1712 from bharadwaj-aditya:fk_in_ddl
Browse files Browse the repository at this point in the history 
PiperOrigin-RevId: 649017205
  • Loading branch information
@cloud-teleport
cloud-teleport committed Jul 3, 2024
2 parents f9b8f5c + 5377dc5 commit 0bcf179
Show file tree
Hide file tree
Showing 5 changed files with 129 additions and 19 deletions.
24 changes: 23 additions & 1 deletion v2/spanner-common/src/main/java/com/google/cloud/teleport/v2/spanner/ddl/Ddl.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -29,11 +29,13 @@
import java.io.Serializable;
import java.util.ArrayList;
import java.util.Collection;
import java.util.HashSet;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.NavigableSet;
import java.util.Set;
import java.util.stream.Collectors;
import javax.annotation.Nullable;

/** Allows to build and introspect Cloud Spanner DDL with Java code. */
Expand Down Expand Up @@ -81,6 +83,25 @@ public Table apply(@Nullable String input) {
});
}

/**
* Return list of all the tables that this table refers to or depends on. This method only
* provides the immediate references, not the whole tree.
*
* @param tableName
* @return
*/
public List<String> tablesReferenced(String tableName) {
Set<String> tablesReferenced = new HashSet<>();
Table table = tables.get(tableName.toLowerCase());
if (table.interleaveInParent() != null) {
tablesReferenced.add(table.interleaveInParent());
}
Set<String> fkReferencedTables =
table.foreignKeys().stream().map(f -> f.referencedTable()).collect(Collectors.toSet());
tablesReferenced.addAll(fkReferencedTables);
return new ArrayList<>(tablesReferenced);
}

private NavigableSet<String> childTableNames(String table) {
return parents.get(table.toLowerCase());
}
Expand Down Expand Up @@ -166,7 +187,8 @@ public List<String> createIndexStatements() {
public List<String> addForeignKeyStatements() {
List<String> result = new ArrayList<>();
for (Table table : allTables()) {
result.addAll(table.foreignKeys());
result.addAll(
table.foreignKeys().stream().map(f -> f.prettyPrint()).collect(Collectors.toList()));
}
return result;
}
Expand Down
8 changes: 4 additions & 4 deletions v2/spanner-common/src/main/java/com/google/cloud/teleport/v2/spanner/ddl/ForeignKey.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,15 +27,15 @@
public abstract class ForeignKey implements Serializable {
private static final long serialVersionUID = 286089905L;

abstract String name();
public abstract String name();

abstract String table();

abstract String referencedTable();
public abstract String referencedTable();

abstract ImmutableList<String> columns();
public abstract ImmutableList<String> columns();

abstract ImmutableList<String> referencedColumns();
public abstract ImmutableList<String> referencedColumns();

abstract Dialect dialect();

Expand Down
4 changes: 2 additions & 2 deletions ...mmon/src/main/java/com/google/cloud/teleport/v2/spanner/ddl/InformationSchemaScanner.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -89,10 +89,10 @@ public Ddl scan() {
continue; // Skipping as table does not exist
}

ImmutableList.Builder<String> tableForeignKeys = ImmutableList.builder();
ImmutableList.Builder<ForeignKey> tableForeignKeys = ImmutableList.builder();
for (Map.Entry<String, ForeignKey.Builder> entry : tableEntry.getValue().entrySet()) {
ForeignKey.Builder foreignKeyBuilder = entry.getValue();
tableForeignKeys.add(foreignKeyBuilder.build().prettyPrint());
tableForeignKeys.add(foreignKeyBuilder.build());
}
builder.createTable(tableName).foreignKeys(tableForeignKeys.build()).endTable();
}
Expand Down
12 changes: 8 additions & 4 deletions v2/spanner-common/src/main/java/com/google/cloud/teleport/v2/spanner/ddl/Table.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -45,7 +45,7 @@ public abstract class Table implements Serializable {

public abstract ImmutableList<String> indexes();

public abstract ImmutableList<String> foreignKeys();
public abstract ImmutableList<ForeignKey> foreignKeys();

public abstract ImmutableList<String> checkConstraints();

Expand Down Expand Up @@ -132,7 +132,9 @@ private void prettyPrintPg(
}
if (includeForeignKeys) {
appendable.append("\n");
appendable.append(String.join("\n", foreignKeys()));
appendable.append(
String.join(
"\n", foreignKeys().stream().map(f -> f.prettyPrint()).collect(Collectors.toList())));
}
}

Expand Down Expand Up @@ -176,7 +178,9 @@ private void prettyPrintGsql(
}
if (includeForeignKeys) {
appendable.append("\n");
appendable.append(String.join("\n", foreignKeys()));
appendable.append(
String.join(
"\n", foreignKeys().stream().map(f -> f.prettyPrint()).collect(Collectors.toList())));
}
}

Expand Down Expand Up @@ -222,7 +226,7 @@ Builder ddlBuilder(Ddl.Builder ddlBuilder) {

public abstract Builder indexes(ImmutableList<String> indexes);

public abstract Builder foreignKeys(ImmutableList<String> foreignKeys);
public abstract Builder foreignKeys(ImmutableList<ForeignKey> foreignKeys);

public abstract Builder checkConstraints(ImmutableList<String> checkConstraints);

Expand Down
100 changes: 92 additions & 8 deletions v2/spanner-common/src/test/java/com/google/cloud/teleport/v2/spanner/ddl/DdlTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -58,7 +58,13 @@ public void testEmptyDbPG() {

@Test
public void testDdlGSQL() {
ForeignKey.Builder usersForeignKeyBuilder =
ForeignKey.builder().name("fk").table("Users").referencedTable("AllowedNames");
usersForeignKeyBuilder.columnsBuilder().add("first_name");
usersForeignKeyBuilder.referencedColumnsBuilder().add("first_name");

Ddl.Builder builder = Ddl.builder();

builder
.createTable("Users")
.column("id")
Expand All @@ -83,10 +89,7 @@ public void testDdlGSQL() {
.asc("id")
.end()
.indexes(ImmutableList.of("CREATE INDEX `UsersByFirstName` ON `Users` (`first_name`)"))
.foreignKeys(
ImmutableList.of(
"ALTER TABLE `Users` ADD CONSTRAINT `fk` FOREIGN KEY (`first_name`)"
+ " REFERENCES `AllowedNames` (`first_name`)"))
.foreignKeys(ImmutableList.of(usersForeignKeyBuilder.build()))
.checkConstraints(ImmutableList.of("CONSTRAINT `ck` CHECK (`first_name` != `last_name`)"))
.endTable();
Ddl ddl = builder.build();
Expand Down Expand Up @@ -124,10 +127,19 @@ public void testDdlGSQL() {
"ALTER TABLE `Users` ADD CONSTRAINT `fk` FOREIGN KEY (`first_name`) REFERENCES"
+ " `AllowedNames` (`first_name`)"));
assertNotNull(ddl.hashCode());
assertTrue(ddl.tablesReferenced("Users").contains("AllowedNames"));
assertTrue(ddl.tablesReferenced("Users").size() == 1);
}

@Test
public void testDdlPG() {
ForeignKey.Builder usersForeignKeyBuilder =
ForeignKey.builder(Dialect.POSTGRESQL)
.name("fk")
.table("Users")
.referencedTable("AllowedNames");
usersForeignKeyBuilder.columnsBuilder().add("first_name");
usersForeignKeyBuilder.referencedColumnsBuilder().add("first_name");
Ddl.Builder builder = Ddl.builder(Dialect.POSTGRESQL);
builder
.createTable("Users")
Expand All @@ -154,10 +166,7 @@ public void testDdlPG() {
.end()
.indexes(
ImmutableList.of("CREATE INDEX \"UsersByFirstName\" ON \"Users\" (\"first_name\")"))
.foreignKeys(
ImmutableList.of(
"ALTER TABLE \"Users\" ADD CONSTRAINT \"fk\" FOREIGN KEY (\"first_name\")"
+ " REFERENCES \"AllowedNames\" (\"first_name\")"))
.foreignKeys(ImmutableList.of(usersForeignKeyBuilder.build()))
.checkConstraints(
ImmutableList.of("CONSTRAINT \"ck\" CHECK (\"first_name\" != \"last_name\")"))
.endTable();
Expand All @@ -178,6 +187,8 @@ public void testDdlPG() {
+ " ALTER TABLE \"Users\" ADD CONSTRAINT \"fk\" FOREIGN KEY (\"first_name\")"
+ " REFERENCES \"AllowedNames\" (\"first_name\")"));
assertNotNull(ddl.hashCode());
assertTrue(ddl.tablesReferenced("Users").contains("AllowedNames"));
assertTrue(ddl.tablesReferenced("Users").size() == 1);
}

@Test
Expand Down Expand Up @@ -245,6 +256,10 @@ public void testInterleavingGSQL() {
assertTrue(perLevelView.containsKey(1));
assertEquals("account", perLevelView.get(1).iterator().next());
assertNotNull(ddl.hashCode());

List<String> tablesReferenced = ddl.tablesReferenced("Account");
assertTrue(tablesReferenced.contains("Users"));
assertTrue(tablesReferenced.size() == 1);
}

@Test
Expand Down Expand Up @@ -307,6 +322,75 @@ public void testInterleavingPG() {
assertNotNull(ddl.hashCode());
}

@Test
public void testReferencedTables() {
ForeignKey.Builder accountsForeignKeyBuilder =
ForeignKey.builder(Dialect.POSTGRESQL)
.name("fk")
.table("Account")
.referencedTable("BalanceNames");
accountsForeignKeyBuilder.columnsBuilder().add("name");
accountsForeignKeyBuilder.referencedColumnsBuilder().add("first_name");
Ddl ddl =
Ddl.builder()
.createTable("Users")
.column("id")
.int64()
.notNull()
.endColumn()
.column("first_name")
.string()
.size(10)
.endColumn()
.column("last_name")
.type(Type.string())
.max()
.endColumn()
.primaryKey()
.asc("id")
.end()
.endTable()
.createTable("Account")
.column("id")
.int64()
.notNull()
.endColumn()
.column("name")
.int64()
.notNull()
.endColumn()
.column("balanceId")
.int64()
.notNull()
.endColumn()
.column("balance")
.float64()
.notNull()
.endColumn()
.primaryKey()
.asc("id")
.end()
.foreignKeys(ImmutableList.of(accountsForeignKeyBuilder.build()))
.interleaveInParent("Users")
.onDeleteCascade()
.endTable()
.createTable("BalanceNames")
.column("id")
.int64()
.notNull()
.endColumn()
.column("first_name")
.string()
.size(10)
.endColumn()
.endTable()
.build();

List<String> accountTablesReferenced = ddl.tablesReferenced("Account");
assertTrue(accountTablesReferenced.containsAll(List.of("Users", "BalanceNames")));
assertTrue(accountTablesReferenced.size() == 2);
}

@Test
public void testIndexGSQL() {
Index.Builder builder =
Expand Down

0 comments on commit 0bcf179

Please sign in to comment.