Implement use database (#3300)

extension/duckdb_scanner/test/test_files/duckdb_scanner.test

@@ -57,6 +57,23 @@ Binder exception: No database named tinysnb has been attached.
 -STATEMENT LOAD FROM tinysnb.person RETURN count(*);
 ---- 1
 8
+-LOG UseDatabaseStatement
+-STATEMENT USE other;
+---- ok
+-STATEMENT LOAD FROM person RETURN *;
+---- 4
+1
+2
+3
+5
+-STATEMENT USE tinysnb;
+---- ok
+-STATEMENT LOAD FROM person RETURN count(*);
+---- 1
+8
+-STATEMENT LOAD FROM other.person RETURN count(*);
+---- 1
+4
 
 -CASE InvalidDuckDBDatabase
 -STATEMENT LOAD FROM tinysnb1.person RETURN *;
@@ -65,3 +82,8 @@ Binder exception: No database named tinysnb1 has been attached.
 -STATEMENT LOAD FROM tinysnb1_person RETURN *;
 ---- error
 Binder exception: Variable tinysnb1_person is not in scope.
+
+-CASE InvalidUseDatabaseName
+-STATEMENT USE NONEXIST;
+---- error
+Runtime exception: No database named NONEXIST.

scripts/antlr4/Cypher.g4.copy

@@ -38,7 +38,8 @@ oC_Statement
         | kU_ExportDatabase
         | kU_ImportDatabase
         | kU_AttachDatabase
-        | kU_DetachDatabase;
+        | kU_DetachDatabase
+        | kU_UseDatabase;
 
 kU_CopyFrom
     : COPY SP oC_SchemaName ( ( SP? kU_ColumnNames SP? ) | SP ) FROM SP kU_ScanSource ( SP? kU_ParsingOptions )? ;
@@ -76,6 +77,12 @@ DBTYPE:
 kU_DetachDatabase
     : DETACH SP oC_SchemaName;
 
+kU_UseDatabase
+    : USE SP oC_SchemaName;
+
+USE:
+    ( 'U' | 'u') ( 'S' | 's') ( 'E' | 'e');
+
 kU_StandaloneCall
     : CALL SP oC_SymbolicName SP? '=' SP? oC_Literal ;
 
@@ -779,6 +786,7 @@ kU_NonReservedKeywords
         | IMPORT
         | EXPORT
         | DATABASE
+        | USE
         ;
 
 UnescapedSymbolicName

src/antlr4/Cypher.g4

@@ -38,7 +38,8 @@ oC_Statement
         | kU_ExportDatabase
         | kU_ImportDatabase
         | kU_AttachDatabase
-        | kU_DetachDatabase;
+        | kU_DetachDatabase
+        | kU_UseDatabase;
 
 kU_CopyFrom
     : COPY SP oC_SchemaName ( ( SP? kU_ColumnNames SP? ) | SP ) FROM SP kU_ScanSource ( SP? kU_ParsingOptions )? ;
@@ -76,6 +77,12 @@ DBTYPE:
 kU_DetachDatabase
     : DETACH SP oC_SchemaName;
 
+kU_UseDatabase
+    : USE SP oC_SchemaName;
+
+USE:
+    ( 'U' | 'u') ( 'S' | 's') ( 'E' | 'e');
+
 kU_StandaloneCall
     : CALL SP oC_SymbolicName SP? '=' SP? oC_Literal ;
 
@@ -779,6 +786,7 @@ kU_NonReservedKeywords
         | IMPORT
         | EXPORT
         | DATABASE
+        | USE
         ;
 
 UnescapedSymbolicName

src/binder/bind/CMakeLists.txt

@@ -21,7 +21,8 @@ add_library(
         bind_updating_clause.cpp
         bind_extension.cpp
         bind_export_database.cpp
-        bind_import_database.cpp)
+        bind_import_database.cpp
+        bind_use_database.cpp)
 
 set(ALL_OBJECT_FILES
         ${ALL_OBJECT_FILES} $<TARGET_OBJECTS:kuzu_binder_bind>

src/binder/bind/bind_attach_database.cpp

@@ -1,14 +1,12 @@
 #include "binder/binder.h"
 #include "binder/bound_attach_database.h"
-#include "common/cast.h"
 #include "parser/attach_database.h"
 
 namespace kuzu {
 namespace binder {
 
 std::unique_ptr<BoundStatement> Binder::bindAttachDatabase(const parser::Statement& statement) {
-    auto& attachDatabase =
-        common::ku_dynamic_cast<const parser::Statement&, const parser::AttachDatabase&>(statement);
+    auto& attachDatabase = statement.constCast<parser::AttachDatabase>();
     return std::make_unique<BoundAttachDatabase>(attachDatabase.getAttachInfo());
 }
 

src/binder/bind/bind_detach_database.cpp

@@ -1,14 +1,12 @@
 #include "binder/binder.h"
 #include "binder/bound_detach_database.h"
-#include "common/cast.h"
 #include "parser/detach_database.h"
 
 namespace kuzu {
 namespace binder {
 
 std::unique_ptr<BoundStatement> Binder::bindDetachDatabase(const parser::Statement& statement) {
-    auto& detachDatabase =
-        common::ku_dynamic_cast<const parser::Statement&, const parser::DetachDatabase&>(statement);
+    auto& detachDatabase = statement.constCast<parser::DetachDatabase>();
     return std::make_unique<BoundDetachDatabase>(detachDatabase.getDBName());
 }
 

src/binder/bind/bind_use_database.cpp

@@ -0,0 +1,14 @@
+#include "binder/binder.h"
+#include "binder/bound_use_database.h"
+#include "parser/use_database.h"
+
+namespace kuzu {
+namespace binder {
+
+std::unique_ptr<BoundStatement> Binder::bindUseDatabase(const parser::Statement& statement) {
+    auto useDatabase = statement.constCast<parser::UseDatabase>();
+    return std::make_unique<BoundUseDatabase>(useDatabase.getDBName());
+}
+
+} // namespace binder
+} // namespace kuzu

src/binder/bind/read/bind_load_from.cpp

@@ -16,38 +16,49 @@ using namespace kuzu::catalog;
 namespace kuzu {
 namespace binder {
 
+static TableFunction getObjectScanFunc(const ObjectScanSource* objectSource,
+    main::ClientContext* clientContext) {
+    // Bind external database table
+    auto dbName = objectSource->objectNames[0];
+    auto attachedDB = clientContext->getDatabaseManager()->getAttachedDatabase(dbName);
+    if (attachedDB == nullptr) {
+        throw BinderException{stringFormat("No database named {} has been attached.", dbName)};
+    }
+    auto tableName = objectSource->objectNames[1];
+    auto attachedCatalog = attachedDB->getCatalogContent();
+    auto tableID = attachedCatalog->getTableID(tableName);
+    auto entry = attachedCatalog->getTableCatalogEntry(tableID);
+    auto tableEntry = ku_dynamic_cast<CatalogEntry*, TableCatalogEntry*>(entry);
+    return tableEntry->getScanFunction();
+}
+
 std::unique_ptr<BoundReadingClause> Binder::bindLoadFrom(const ReadingClause& readingClause) {
     auto& loadFrom = readingClause.constCast<LoadFrom>();
     TableFunction scanFunction;
     std::unique_ptr<TableFuncBindData> bindData;
     auto source = loadFrom.getSource();
     switch (source->type) {
     case ScanSourceType::OBJECT: {
-        auto objectSource = source->constPtrCast<ObjectScanSource>();
+        auto objectSource = reinterpret_cast<ObjectScanSource*>(source);
         KU_ASSERT(!objectSource->objectNames.empty());
         if (objectSource->objectNames.size() == 1) {
             // Bind external object as table
             auto objectName = objectSource->objectNames[0];
             auto replacementData = clientContext->tryReplace(objectName);
-            if (replacementData == nullptr) {
+            if (replacementData != nullptr) {
+                scanFunction = replacementData->func;
+                bindData = scanFunction.bindFunc(clientContext, &replacementData->bindInput);
+            } else if (clientContext->getDatabaseManager()->hasDefaultDatabase()) {
+                objectSource->objectNames.insert(objectSource->objectNames.begin(),
+                    clientContext->getDatabaseManager()->getDefaultDatabase());
+                scanFunction = getObjectScanFunc(objectSource, clientContext);
+                auto bindInput = function::TableFuncBindInput();
+                bindData = scanFunction.bindFunc(clientContext, &bindInput);
+            } else {
                 throw BinderException(ExceptionMessage::variableNotInScope(objectName));
             }
-            scanFunction = replacementData->func;
-            bindData = scanFunction.bindFunc(clientContext, &replacementData->bindInput);
         } else if (objectSource->objectNames.size() == 2) {
-            // Bind external database table
-            auto dbName = objectSource->objectNames[0];
-            auto attachedDB = clientContext->getDatabaseManager()->getAttachedDatabase(dbName);
-            if (attachedDB == nullptr) {
-                throw BinderException{
-                    stringFormat("No database named {} has been attached.", dbName)};
-            }
-            auto tableName = objectSource->objectNames[1];
-            auto attachedCatalog = attachedDB->getCatalogContent();
-            auto tableID = attachedCatalog->getTableID(tableName);
-            auto entry = attachedCatalog->getTableCatalogEntry(tableID);
-            auto tableEntry = ku_dynamic_cast<CatalogEntry*, TableCatalogEntry*>(entry);
-            scanFunction = tableEntry->getScanFunction();
+            scanFunction = getObjectScanFunc(objectSource, clientContext);
             auto bindInput = function::TableFuncBindInput();
             bindData = scanFunction.bindFunc(clientContext, &bindInput);
         } else {

src/binder/binder.cpp

@@ -74,6 +74,9 @@ std::unique_ptr<BoundStatement> Binder::bind(const Statement& statement) {
     case StatementType::DETACH_DATABASE: {
         boundStatement = bindDetachDatabase(statement);
     } break;
+    case StatementType::USE_DATABASE: {
+        boundStatement = bindUseDatabase(statement);
+    } break;
     default: {
         KU_UNREACHABLE;
     }

src/binder/bound_statement_visitor.cpp

@@ -61,6 +61,9 @@ void BoundStatementVisitor::visit(const BoundStatement& statement) {
     case StatementType::DETACH_DATABASE: {
         visitDetachDatabase(statement);
     } break;
+    case StatementType::USE_DATABASE: {
+        visitUseDatabase(statement);
+    } break;
     default:
         KU_UNREACHABLE;
     }

src/include/binder/binder.h

@@ -119,6 +119,7 @@ class Binder {
 
     std::unique_ptr<BoundStatement> bindAttachDatabase(const parser::Statement& statement);
     std::unique_ptr<BoundStatement> bindDetachDatabase(const parser::Statement& statement);
+    std::unique_ptr<BoundStatement> bindUseDatabase(const parser::Statement& statement);
 
     /*** bind scan source ***/
     std::unique_ptr<BoundBaseScanSource> bindScanSource(parser::BaseScanSource* scanSource,

src/include/binder/bound_database_statement.h

@@ -0,0 +1,21 @@
+#pragma once
+
+#include "binder/bound_statement.h"
+
+namespace kuzu {
+namespace binder {
+
+class BoundDatabaseStatement : public BoundStatement {
+public:
+    explicit BoundDatabaseStatement(common::StatementType statementType, std::string dbName)
+        : BoundStatement{statementType, BoundStatementResult::createEmptyResult()},
+          dbName{std::move(dbName)} {}
+
+    std::string getDBName() const { return dbName; }
+
+private:
+    std::string dbName;
+};
+
+} // namespace binder
+} // namespace kuzu

src/include/binder/bound_detach_database.h

@@ -1,21 +1,14 @@
 #pragma once
 
-#include "binder/bound_statement.h"
+#include "binder/bound_database_statement.h"
 
 namespace kuzu {
 namespace binder {
 
-class BoundDetachDatabase final : public BoundStatement {
+class BoundDetachDatabase final : public BoundDatabaseStatement {
 public:
     explicit BoundDetachDatabase(std::string dbName)
-        : BoundStatement{common::StatementType::DETACH_DATABASE,
-              BoundStatementResult::createEmptyResult()},
-          dbName{std::move(dbName)} {}
-
-    std::string getDBName() const { return dbName; }
-
-private:
-    std::string dbName;
+        : BoundDatabaseStatement{common::StatementType::DETACH_DATABASE, std::move(dbName)} {}
 };
 
 } // namespace binder

src/include/binder/bound_statement_visitor.h

@@ -50,6 +50,7 @@ class BoundStatementVisitor {
     virtual void visitProjectionBodyPredicate(const std::shared_ptr<Expression>& /* predicate*/) {}
     virtual void visitAttachDatabase(const BoundStatement&) {}
     virtual void visitDetachDatabase(const BoundStatement&) {}
+    virtual void visitUseDatabase(const BoundStatement&) {}
 };
 
 } // namespace binder

src/include/binder/bound_use_database.h

@@ -0,0 +1,15 @@
+#pragma once
+
+#include "binder/bound_database_statement.h"
+
+namespace kuzu {
+namespace binder {
+
+class BoundUseDatabase final : public BoundDatabaseStatement {
+public:
+    explicit BoundUseDatabase(std::string dbName)
+        : BoundDatabaseStatement{common::StatementType::USE_DATABASE, std::move(dbName)} {}
+};
+
+} // namespace binder
+} // namespace kuzu

src/include/common/enums/statement_type.h

@@ -22,6 +22,7 @@ enum class StatementType : uint8_t {
     IMPORT_DATABASE = 33,
     ATTACH_DATABASE = 34,
     DETACH_DATABASE = 35,
+    USE_DATABASE = 36,
 };
 
 struct StatementTypeUtils {

src/include/main/database_manager.h

@@ -7,12 +7,18 @@ namespace main {
 
 class DatabaseManager {
 public:
+    DatabaseManager();
+
     void registerAttachedDatabase(std::unique_ptr<AttachedDatabase> attachedDatabase);
     AttachedDatabase* getAttachedDatabase(const std::string& name);
     void detachDatabase(const std::string& databaseName);
+    std::string getDefaultDatabase() const { return defaultDatabase; }
+    bool hasDefaultDatabase() const { return defaultDatabase != ""; }
+    void setDefaultDatabase(const std::string& databaseName);
 
 private:
     std::vector<std::unique_ptr<AttachedDatabase>> attachedDatabases;
+    std::string defaultDatabase;
 };
 
 } // namespace main

src/include/parser/database_statement.h

@@ -0,0 +1,22 @@
+#pragma once
+
+#include <string>
+
+#include "parser/statement.h"
+
+namespace kuzu {
+namespace parser {
+
+class DatabaseStatement : public Statement {
+public:
+    explicit DatabaseStatement(common::StatementType type, std::string dbName)
+        : Statement{type}, dbName{std::move(dbName)} {}
+
+    std::string getDBName() const { return dbName; }
+
+private:
+    std::string dbName;
+};
+
+} // namespace parser
+} // namespace kuzu

src/include/parser/detach_database.h

@@ -2,20 +2,15 @@
 
 #include <string>
 
-#include "parser/statement.h"
+#include "parser/database_statement.h"
 
 namespace kuzu {
 namespace parser {
 
-class DetachDatabase final : public Statement {
+class DetachDatabase final : public DatabaseStatement {
 public:
     explicit DetachDatabase(std::string dbName)
-        : Statement{common::StatementType::DETACH_DATABASE}, dbName{std::move(dbName)} {}
-
-    std::string getDBName() const { return dbName; }
-
-private:
-    std::string dbName;
+        : DatabaseStatement{common::StatementType::DETACH_DATABASE, std::move(dbName)} {}
 };
 
 } // namespace parser

src/include/parser/parsed_statement_visitor.h

@@ -52,6 +52,7 @@ class StatementVisitor {
     virtual void visitImportDatabase(const Statement& /*statement*/) {}
     virtual void visitAttachDatabase(const Statement& /*statement*/) {}
     virtual void visitDetachDatabase(const Statement& /*statement*/) {}
+    virtual void visitUseDatabase(const Statement& /*statement*/) {}
     // LCOV_EXCL_STOP
 };