🐛 Source CockroachDB: handled case when user can see only tables with…

… select grants (#9004) * Handled case when user can see only tables with select grants * Fixed failed test for cockroachDB * format * added integration test * bump versions * bump versions
airbytehq · Dec 24, 2021 · 6f4c3d1 · 6f4c3d1
1 parent cd30cf4
commit 6f4c3d1
Show file tree

Hide file tree

Showing 8 changed files with 77 additions and 12 deletions.
diff --git a/...ain/resources/config/STANDARD_SOURCE_DEFINITION/9fa5862c-da7c-11eb-8d19-0242ac130003.json b/...ain/resources/config/STANDARD_SOURCE_DEFINITION/9fa5862c-da7c-11eb-8d19-0242ac130003.json
@@ -2,7 +2,7 @@
   "sourceDefinitionId": "9fa5862c-da7c-11eb-8d19-0242ac130003",
   "name": "Cockroachdb",
   "dockerRepository": "airbyte/source-cockroachdb",
-  "dockerImageTag": "0.1.4",
+  "dockerImageTag": "0.1.5",
   "documentationUrl": "https://docs.airbyte.io/integrations/sources/cockroachdb",
   "icon": "cockroachdb.svg"
 }
diff --git a/airbyte-config/init/src/main/resources/seed/source_definitions.yaml b/airbyte-config/init/src/main/resources/seed/source_definitions.yaml
@@ -133,7 +133,7 @@
 - name: Cockroachdb
   sourceDefinitionId: 9fa5862c-da7c-11eb-8d19-0242ac130003
   dockerRepository: airbyte/source-cockroachdb
-  dockerImageTag: 0.1.4
+  dockerImageTag: 0.1.5
   documentationUrl: https://docs.airbyte.io/integrations/sources/cockroachdb
   icon: cockroachdb.svg
   sourceType: database

diff --git a/airbyte-integrations/connectors/source-cockroachdb-strict-encrypt/Dockerfile b/airbyte-integrations/connectors/source-cockroachdb-strict-encrypt/Dockerfile
@@ -16,5 +16,5 @@ ENV APPLICATION source-cockroachdb-strict-encrypt
 
 COPY --from=build /airbyte /airbyte
 
-LABEL io.airbyte.version=0.1.1
+LABEL io.airbyte.version=0.1.2
 LABEL io.airbyte.name=airbyte/source-cockroachdb-strict-encrypt
diff --git a/airbyte-integrations/connectors/source-cockroachdb/Dockerfile b/airbyte-integrations/connectors/source-cockroachdb/Dockerfile
@@ -16,5 +16,5 @@ ENV APPLICATION source-cockroachdb
 
 COPY --from=build /airbyte /airbyte
 
-LABEL io.airbyte.version=0.1.4
+LABEL io.airbyte.version=0.1.5
 LABEL io.airbyte.name=airbyte/source-cockroachdb
diff --git a/...ckroachdb/src/main/java/io/airbyte/integrations/source/cockroachdb/CockroachDbSource.java b/...ckroachdb/src/main/java/io/airbyte/integrations/source/cockroachdb/CockroachDbSource.java
@@ -6,20 +6,27 @@
 
 import com.fasterxml.jackson.databind.JsonNode;
 import com.google.common.collect.ImmutableMap;
+import io.airbyte.commons.functional.CheckedFunction;
 import io.airbyte.commons.json.Jsons;
 import io.airbyte.commons.util.AutoCloseableIterator;
+import io.airbyte.db.jdbc.JdbcDatabase;
 import io.airbyte.db.jdbc.PostgresJdbcStreamingQueryConfiguration;
 import io.airbyte.integrations.base.IntegrationRunner;
 import io.airbyte.integrations.base.Source;
 import io.airbyte.integrations.base.ssh.SshWrappedSource;
 import io.airbyte.integrations.source.jdbc.AbstractJdbcSource;
+import io.airbyte.integrations.source.jdbc.dto.JdbcPrivilegeDto;
 import io.airbyte.protocol.models.AirbyteConnectionStatus;
 import io.airbyte.protocol.models.AirbyteMessage;
 import io.airbyte.protocol.models.ConfiguredAirbyteCatalog;
+import java.sql.Connection;
 import java.sql.JDBCType;
+import java.sql.PreparedStatement;
+import java.sql.SQLException;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Set;
+import java.util.stream.Collectors;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -90,6 +97,32 @@ public AutoCloseableIterator<AirbyteMessage> read(final JsonNode config,
     return super.read(config, catalog, state);
   }
 
+  @Override
+  public Set<JdbcPrivilegeDto> getPrivilegesTableForCurrentUser(final JdbcDatabase database, final String schema) throws SQLException {
+    return database
+        .query(getPrivileges(database), sourceOperations::rowToJson)
+        .map(this::getPrivilegeDto)
+        .collect(Collectors.toSet());
+  }
+
+  private CheckedFunction<Connection, PreparedStatement, SQLException> getPrivileges(JdbcDatabase database) {
+    return connection -> {
+      final PreparedStatement ps = connection.prepareStatement(
+          "SELECT DISTINCT table_catalog, table_schema, table_name, privilege_type\n"
+              + "FROM   information_schema.table_privileges\n"
+              + "WHERE  grantee = ? AND privilege_type in ('SELECT', 'ALL')");
+      ps.setString(1, database.getDatabaseConfig().get("username").asText());
+      return ps;
+    };
+  }
+
+  private JdbcPrivilegeDto getPrivilegeDto(JsonNode jsonNode) {
+    return JdbcPrivilegeDto.builder()
+        .schemaName(jsonNode.get("table_schema").asText())
+        .tableName(jsonNode.get("table_name").asText())
+        .build();
+  }
+
   public static void main(final String[] args) throws Exception {
     final Source source = new CockroachDbSource();
     LOGGER.info("starting source: {}", CockroachDbSource.class);

diff --git a/.../java/io/airbyte/integrations/source/cockroachdb/CockroachDbJdbcSourceAcceptanceTest.java b/.../java/io/airbyte/integrations/source/cockroachdb/CockroachDbJdbcSourceAcceptanceTest.java
@@ -12,7 +12,6 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Lists;
-import io.airbyte.commons.io.IOs;
 import io.airbyte.commons.json.Jsons;
 import io.airbyte.commons.string.Strings;
 import io.airbyte.commons.util.MoreIterators;
@@ -36,7 +35,6 @@
 import io.airbyte.protocol.models.Field;
 import io.airbyte.protocol.models.JsonSchemaPrimitive;
 import io.airbyte.protocol.models.SyncMode;
-import io.airbyte.test.utils.CockroachDBContainerHelper;
 import java.sql.JDBCType;
 import java.util.ArrayList;
 import java.util.Comparator;
@@ -48,7 +46,6 @@
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.testcontainers.containers.CockroachContainer;
-import org.testcontainers.utility.MountableFile;
 
 class CockroachDbJdbcSourceAcceptanceTest extends JdbcSourceAcceptanceTest {
 

diff --git a/...achdb/src/test/java/io/airbyte/integrations/source/cockroachdb/CockroachDbSourceTest.java b/...achdb/src/test/java/io/airbyte/integrations/source/cockroachdb/CockroachDbSourceTest.java
@@ -11,7 +11,6 @@
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Lists;
 import com.google.common.collect.Sets;
-import io.airbyte.commons.io.IOs;
 import io.airbyte.commons.json.Jsons;
 import io.airbyte.commons.string.Strings;
 import io.airbyte.commons.util.MoreIterators;
@@ -27,7 +26,6 @@
 import io.airbyte.protocol.models.Field;
 import io.airbyte.protocol.models.JsonSchemaPrimitive;
 import io.airbyte.protocol.models.SyncMode;
-import io.airbyte.test.utils.CockroachDBContainerHelper;
 import java.math.BigDecimal;
 import java.util.HashMap;
 import java.util.List;
@@ -41,7 +39,6 @@
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.testcontainers.containers.CockroachContainer;
-import org.testcontainers.utility.MountableFile;
 
 class CockroachDbSourceTest {
 
@@ -104,7 +101,6 @@ static void init() {
   void setup() throws Exception {
     dbName = Strings.addRandomSuffix("db", "_", 10).toLowerCase();
 
-
     final JsonNode config = getConfig(PSQL_DB, dbName);
     final Database database = getDatabaseFromConfig(config);
     database.query(ctx -> {
@@ -142,11 +138,15 @@ private static Database getDatabaseFromConfig(final JsonNode config) {
   }
 
   private JsonNode getConfig(final CockroachContainer psqlDb, final String dbName) {
+    return getConfig(psqlDb, dbName, psqlDb.getUsername());
+  }
+
+  private JsonNode getConfig(final CockroachContainer psqlDb, final String dbName, final String username) {
     return Jsons.jsonNode(ImmutableMap.builder()
         .put("host", psqlDb.getHost())
         .put("port", psqlDb.getFirstMappedPort() - 1)
         .put("database", dbName)
-        .put("username", psqlDb.getUsername())
+        .put("username", username)
         .put("password", psqlDb.getPassword())
         .put("ssl", false)
         .build());
@@ -206,6 +206,39 @@ void testDiscoverWithPk() throws Exception {
     });
   }
 
+  @Test
+  void testDiscoverWithPermissions() throws Exception {
+    JsonNode config = getConfig(PSQL_DB, dbName);
+    final Database database = getDatabaseFromConfig(config);
+    database.query(ctx -> {
+      ctx.fetch(
+          "CREATE USER cock;");
+      ctx.fetch(
+          "CREATE TABLE id_and_name_perm1(id NUMERIC(20, 10), name VARCHAR(200), power double precision, PRIMARY KEY (id));");
+      ctx.fetch(
+          "CREATE TABLE id_and_name_perm2(id NUMERIC(20, 10), name VARCHAR(200), power double precision, PRIMARY KEY (id));");
+      ctx.fetch(
+          "CREATE TABLE id_and_name_perm3(id NUMERIC(20, 10), name VARCHAR(200), power double precision, PRIMARY KEY (id));");
+      ctx.fetch("grant all on database " + dbName + " to cock;");
+      ctx.fetch("grant all on table " + dbName + ".public.id_and_name_perm1 to cock;");
+      ctx.fetch("grant select on table " + dbName + ".public.id_and_name_perm2 to cock;");
+      return null;
+    });
+
+    List<String> expected = List.of("id_and_name_perm1", "id_and_name_perm2");
+
+    AirbyteCatalog airbyteCatalog = new CockroachDbSource().discover(getConfig(PSQL_DB, dbName, "cock"));
+    final List<String> actualNamesWithPermission =
+        airbyteCatalog
+            .getStreams()
+            .stream()
+            .map(AirbyteStream::getName)
+            .toList();
+
+    assertEquals(expected.size(), actualNamesWithPermission.size());
+    assertEquals(expected, actualNamesWithPermission);
+  }
+
   @Test
   void testReadSuccess() throws Exception {
     final ConfiguredAirbyteCatalog configuredCatalog =

diff --git a/docs/integrations/sources/cockroachdb.md b/docs/integrations/sources/cockroachdb.md
@@ -95,6 +95,7 @@ Your database user should now be ready for use with Airbyte.
 
 | Version | Date | Pull Request | Subject |
 | :--- | :--- | :--- | :--- |
+| 0.1.5 | 2021-12-24 | [9004](https://github.com/airbytehq/airbyte/pull/9004) | User can see only permmited tables during discovery |
 | 0.1.4 | 2021-12-24 | [8958](https://github.com/airbytehq/airbyte/pull/8958) | Add support for JdbcType.ARRAY |
 | 0.1.3 | 2021-10-10 | [7819](https://github.com/airbytehq/airbyte/pull/7819) | Fixed Datatype errors during Cockroach DB parsing |
 | 0.1.2 | 2021-08-13 | [4699](https://github.com/airbytehq/airbyte/pull/4699) | Added json config validator |
@@ -103,6 +104,7 @@ Your database user should now be ready for use with Airbyte.
 
 | Version | Date | Pull Request | Subject |
 |:--------| :--- | :--- | :--- |
+| 0.1.2   | 2021-12-24 | [9004](https://github.com/airbytehq/airbyte/pull/9004) | User can see only permmited tables during discovery |
 | 0.1.1   | 2021-12-24 | [8958](https://github.com/airbytehq/airbyte/pull/8958) | Add support for JdbcType.ARRAY |
 | 0.1.0   | 2021-11-23 | [7457](https://github.com/airbytehq/airbyte/pull/7457) | CockroachDb source: Add only encrypted version for the connector |