Improve Cassandra visualization

packages/core/dialect/cql/index.ts

@@ -153,12 +153,20 @@ export default class CQLDialect extends GenericDialect<CassandraLib.Client> impl
   }
 
   public async getTables(): Promise<DatabaseInterface.Table[]> {
-    const [queryResults] = await this.query(this.queries.fetchTables);
+    const [queryResults, numberOfColumnsResults] = await this.query(this.queries.fetchTables);
+    const numberOfColumnsMap: {string: {string: number}} = numberOfColumnsResults.results.reduce((prev, curr) => prev.concat(curr), []).reduce((acc, obj: any) => {
+      if (typeof acc[obj.keyspace_name] === 'undefined') {
+        acc[obj.keyspace_name] = {};
+      }
+      acc[obj.keyspace_name][obj.table_name] = parseInt(JSON.parse(obj.count), 10);
+      return acc;
+    }, {});
     return queryResults.results.reduce((prev, curr) => prev.concat(curr), []).map((obj: any) => {
       const table: DatabaseInterface.Table = {
         name: obj.table_name,
         isView: false,
         tableSchema: obj.keyspace_name,
+        numberOfColumns: numberOfColumnsMap[obj.keyspace_name] ? numberOfColumnsMap[obj.keyspace_name][obj.table_name] : undefined,
         tree: [obj.keyspace_name, 'tables', obj.table_name].join(TREE_SEP)
       };
       return table;
@@ -174,6 +182,7 @@ export default class CQLDialect extends GenericDialect<CassandraLib.Client> impl
         type: obj.type,
         isNullable: obj.kind === 'regular',
         isPk: obj.kind !== 'regular',
+        isPartitionKey: obj.kind === 'partition_key',
         tableSchema: obj.keyspace_name,
         tree: [obj.keyspace_name, 'tables', obj.table_name, obj.column_name].join(TREE_SEP)
       };

packages/core/dialect/cql/queries.ts

@@ -7,6 +7,10 @@ export default {
     AND table_name = ':table'`,
   fetchColumns: `SELECT * FROM system_schema.columns`,
   fetchRecords: `SELECT * FROM :keyspace.:table LIMIT :limit`,
-  fetchTables: `SELECT * FROM system_schema.tables`,
-  fetchFunctions: `SELECT * FROM system_schema.functions`
+  fetchTables: `
+    SELECT * FROM system_schema.tables;
+    SELECT keyspace_name, table_name, COUNT(*)
+    FROM system_schema.columns
+    GROUP BY keyspace_name, table_name`,
+  fetchFunctions: `SELECT * FROM system_schema.functions`,
 } as DialectQueries;

packages/core/plugin-api.d.ts

@@ -140,6 +140,7 @@ export namespace DatabaseInterface {
     tableCatalog?: string;
     defaultValue?: string;
     isNullable: boolean;
+    isPartitionKey?: boolean;
     isPk?: boolean;
     isFk?: boolean;
     columnKey?: string;

packages/plugins/connection-manager/explorer/SidebarColumn.ts

@@ -12,7 +12,7 @@ export default class SidebarColumn extends SidebarAbstractItem<null> {
   }
   public get description() {
     let typeSize = '';
-    if (this.column.size !== null) {
+    if (typeof this.column.size !== 'undefined' && this.column.size !== null) {
       typeSize = `(${this.column.size})`;
     }
     return `${(this.column.type || '').toUpperCase()}${typeSize}`;
@@ -35,6 +35,10 @@ export default class SidebarColumn extends SidebarAbstractItem<null> {
           dark: context.asAbsolutePath('icons/fk-dark.svg'),
           light: context.asAbsolutePath('icons/fk-light.svg'),
         },
+        partitionKey: {
+          dark: context.asAbsolutePath('icons/partition-key-dark.svg'),
+          light: context.asAbsolutePath('icons/partition-key-light.svg'),
+        },
       };
     }
     this.updateIconPath();
@@ -46,7 +50,10 @@ export default class SidebarColumn extends SidebarAbstractItem<null> {
   }
   public updateIconPath() {
     this.iconPath = SidebarColumn.icons.default;
-    if (this.column.isPk) {
+    if (this.column.isPartitionKey) {
+      this.iconPath = SidebarColumn.icons.partitionKey;
+    }
+    else if (this.column.isPk) {
       this.iconPath = SidebarColumn.icons.primaryKey;
     }
     else if (this.column.isFk) {

packages/plugins/connection-manager/explorer/index.ts

@@ -17,6 +17,7 @@ import logger from '@sqltools/core/log/vscode';
 const DialectHierarchyChildNames = {
   [DatabaseDialect.PostgreSQL]: ['Database', 'Schema'],
   [DatabaseDialect['AWS Redshift']]: ['Database', 'Schema'],
+  [DatabaseDialect.Cassandra]: ['Keyspace'],
 }
 
 

yarn.lock

@@ -1061,6 +1061,11 @@
     "@types/minimatch" "*"
     "@types/node" "*"
 
+"@types/ibm_db@^2.0.5":
+  version "2.0.5"
+  resolved "https://registry.yarnpkg.com/@types/ibm_db/-/ibm_db-2.0.5.tgz#6804bb7f4b3ed5018ed2db54755ac37ae482f9cb"
+  integrity sha512-Tb4CRJ4xz+iwixGMjV0WJ5612WZBot/zwYMti/kRmnSLM6K0Wk57yuw+hyjAq9NJp5rY9tHQVzhDzkxh9QuQqQ==
+
 "@types/istanbul-lib-coverage@*", "@types/istanbul-lib-coverage@^2.0.0":
   version "2.0.1"
   resolved "https://registry.yarnpkg.com/@types/istanbul-lib-coverage/-/istanbul-lib-coverage-2.0.1.tgz#42995b446db9a48a11a07ec083499a860e9138ff"