Examples with TrackedSpannerTable com.google.cloud.teleport.v2.templates.spannerchangestreamstobigquery.model.TrackedSpannerTable used on opensource projects

Example 1 with TrackedSpannerTable

use of com.google.cloud.teleport.v2.templates.spannerchangestreamstobigquery.model.TrackedSpannerTable in project DataflowTemplates by GoogleCloudPlatform.

the class SchemaUtilsTest method testChangeStreamTrackAll.

@Test
public void testChangeStreamTrackAll() {
    mockInformationSchemaChangeStreamsQuery(true);
    mockInformationSchemaTablesQuery();
    mockInformationSchemaColumnsQuery();
    mockInformationSchemaKeyColumnUsageQuery();
    String sql = "SELECT TABLE_NAME, COLUMN_NAME FROM INFORMATION_SCHEMA.CHANGE_STREAM_COLUMNS " + "WHERE CHANGE_STREAM_NAME = @changeStreamName";
    when(mockReadContext.executeQuery(Statement.newBuilder(sql).bind("changeStreamName").to(changeStreamName).build())).thenReturn(ResultSets.forRows(Type.struct(Type.StructField.of("TABLE_NAME", Type.string()), Type.StructField.of("COLUMN_NAME", Type.string())), Collections.emptyList()));
    Map<String, TrackedSpannerTable> actualSpannerTableByName = new SpannerUtils(mockDatabaseClient, changeStreamName).getSpannerTableByName();
    List<TrackedSpannerColumn> singersPkColumns = ImmutableList.of(TrackedSpannerColumn.create("SingerId", Type.int64(), 1));
    List<TrackedSpannerColumn> singersNonPkColumns = ImmutableList.of(TrackedSpannerColumn.create("FirstName", Type.string(), 2), TrackedSpannerColumn.create("LastName", Type.string(), 3));
    Map<String, TrackedSpannerTable> expectedSpannerTableByName = new HashMap<>();
    expectedSpannerTableByName.put("Singers", new TrackedSpannerTable("Singers", singersPkColumns, singersNonPkColumns));
    assertThat(actualSpannerTableByName).isEqualTo(expectedSpannerTableByName);
}

Example 2 with TrackedSpannerTable

use of com.google.cloud.teleport.v2.templates.spannerchangestreamstobigquery.model.TrackedSpannerTable in project DataflowTemplates by GoogleCloudPlatform.

the class SchemaUtilsTest method testChangeStreamTrackOneColumn.

@Test
public void testChangeStreamTrackOneColumn() {
    mockInformationSchemaChangeStreamsQuery(false);
    mockInformationSchemaChangeStreamTablesQuery();
    mockInformationSchemaColumnsQuery();
    mockInformationSchemaKeyColumnUsageQuery();
    String sql = "SELECT TABLE_NAME, COLUMN_NAME FROM INFORMATION_SCHEMA.CHANGE_STREAM_COLUMNS " + "WHERE CHANGE_STREAM_NAME = @changeStreamName";
    // spotless:off
    when(mockReadContext.executeQuery(Statement.newBuilder(sql).bind("changeStreamName").to(changeStreamName).build())).thenReturn(ResultSets.forRows(Type.struct(Type.StructField.of("TABLE_NAME", Type.string()), Type.StructField.of("COLUMN_NAME", Type.string())), ImmutableList.of(Struct.newBuilder().set("TABLE_NAME").to(Value.string("Singers")).set("COLUMN_NAME").to(Value.string("SingerId")).build(), Struct.newBuilder().set("TABLE_NAME").to(Value.string("Singers")).set("COLUMN_NAME").to(Value.string("FirstName")).build())));
    // spotless:on
    Map<String, TrackedSpannerTable> actualSpannerTableByName = new SpannerUtils(mockDatabaseClient, changeStreamName).getSpannerTableByName();
    List<TrackedSpannerColumn> singersPkColumns = Collections.singletonList(TrackedSpannerColumn.create("SingerId", Type.int64(), 1));
    List<TrackedSpannerColumn> singersNonPkColumns = Collections.singletonList(TrackedSpannerColumn.create("FirstName", Type.string(), 2));
    Map<String, TrackedSpannerTable> expectedSpannerTableByName = new HashMap<>();
    expectedSpannerTableByName.put("Singers", new TrackedSpannerTable("Singers", singersPkColumns, singersNonPkColumns));
    assertThat(actualSpannerTableByName).isEqualTo(expectedSpannerTableByName);
}

Example 3 with TrackedSpannerTable

use of com.google.cloud.teleport.v2.templates.spannerchangestreamstobigquery.model.TrackedSpannerTable in project DataflowTemplates by GoogleCloudPlatform.

the class SchemaUtilsTest method testChangeStreamTrackOneTable.

@Test
public void testChangeStreamTrackOneTable() {
    mockInformationSchemaChangeStreamsQuery(false);
    mockInformationSchemaChangeStreamTablesQuery();
    mockInformationSchemaColumnsQuery();
    mockInformationSchemaKeyColumnUsageQuery();
    String sql = "SELECT TABLE_NAME, COLUMN_NAME FROM INFORMATION_SCHEMA.CHANGE_STREAM_COLUMNS " + "WHERE CHANGE_STREAM_NAME = @changeStreamName";
    when(mockReadContext.executeQuery(Statement.newBuilder(sql).bind("changeStreamName").to(changeStreamName).build())).thenReturn(ResultSets.forRows(Type.struct(Type.StructField.of("TABLE_NAME", Type.string()), Type.StructField.of("COLUMN_NAME", Type.string())), Collections.emptyList()));
    Map<String, TrackedSpannerTable> actualSpannerTableByName = new SpannerUtils(mockDatabaseClient, changeStreamName).getSpannerTableByName();
    List<TrackedSpannerColumn> singersPkColumns = ImmutableList.of(TrackedSpannerColumn.create("SingerId", Type.int64(), 1));
    List<TrackedSpannerColumn> singersNonPkColumns = ImmutableList.of(TrackedSpannerColumn.create("FirstName", Type.string(), 2), TrackedSpannerColumn.create("LastName", Type.string(), 3));
    Map<String, TrackedSpannerTable> expectedSpannerTableByName = new HashMap<>();
    expectedSpannerTableByName.put("Singers", new TrackedSpannerTable("Singers", singersPkColumns, singersNonPkColumns));
    assertThat(actualSpannerTableByName).isEqualTo(expectedSpannerTableByName);
}

Example 4 with TrackedSpannerTable

use of com.google.cloud.teleport.v2.templates.spannerchangestreamstobigquery.model.TrackedSpannerTable in project DataflowTemplates by GoogleCloudPlatform.

the class BigQueryDynamicDestinations method getSchema.

@Override
public TableSchema getSchema(KV<TableId, TableRow> destination) {
    TableRow tableRow = destination.getValue();
    String spannerTableName = (String) tableRow.get(BigQueryUtils.BQ_CHANGELOG_FIELD_NAME_TABLE_NAME);
    TrackedSpannerTable spannerTable = spannerTableByName.get(spannerTableName);
    List<TableFieldSchema> fields = getFields(spannerTable);
    List<TableFieldSchema> filteredFields = new ArrayList<>();
    for (TableFieldSchema field : fields) {
        if (!ignoreFields.contains(field.getName())) {
            filteredFields.add(field);
        }
    }
    return new TableSchema().setFields(filteredFields);
}

Example 5 with TrackedSpannerTable

use of com.google.cloud.teleport.v2.templates.spannerchangestreamstobigquery.model.TrackedSpannerTable in project DataflowTemplates by GoogleCloudPlatform.

the class SpannerUtils method getSpannerTableByName.

/**
 * @return a map where the key is the table name tracked by the change stream and the value is the
 *     {@link SpannerTable} object of the table name.
 */
private Map<String, TrackedSpannerTable> getSpannerTableByName(Set<String> spannerTableNames, Map<String, Set<String>> spannerColumnNamesExplicitlyTrackedByChangeStreamByTableName) {
    Map<String, Set<String>> keyColumnNameByTableName = getKeyColumnNameByTableName(spannerTableNames);
    Map<String, List<TrackedSpannerColumn>> spannerColumnsByTableName = getSpannerColumnsByTableName(spannerTableNames, keyColumnNameByTableName, spannerColumnNamesExplicitlyTrackedByChangeStreamByTableName);
    Map<String, TrackedSpannerTable> result = new HashMap<>();
    for (String tableName : spannerColumnsByTableName.keySet()) {
        List<TrackedSpannerColumn> pkColumns = new ArrayList<>();
        List<TrackedSpannerColumn> nonPkColumns = new ArrayList<>();
        Set<String> keyColumnNames = keyColumnNameByTableName.get(tableName);
        for (TrackedSpannerColumn spannerColumn : spannerColumnsByTableName.get(tableName)) {
            if (keyColumnNames.contains(spannerColumn.getName())) {
                pkColumns.add(spannerColumn);
            } else {
                nonPkColumns.add(spannerColumn);
            }
        }
        result.put(tableName, new TrackedSpannerTable(tableName, pkColumns, nonPkColumns));
    }
    return result;
}